Merge pull request googleapis#3404 from GoogleCloudPlatform/3017-spanner-read_with_limit_and_range

tseaver · web-flow · commit ae98aca682a5 · 2017-05-12T16:37:33.000-04:00
Add tests for range reads and reads with limits.
diff --git a/spanner/tests/system.py b/spanner/tests/system.py
@@ -492,24 +492,29 @@ def _row_data(max_index):
             yield [
                 index,
                 'First%09d' % (index,),
-                'Last09%d' % (index),
+                'Last%09d' % (max_index - index),
                 'test-%09d@example.com' % (index,),
             ]
 
-    def _set_up_table(self, row_count):
-        retry = RetryInstanceState(_has_all_ddl)
-        retry(self._db.reload)()
+    def _set_up_table(self, row_count, db=None):
 
-        session = self._db.session()
+        if db is None:
+            db = self._db
+            retry = RetryInstanceState(_has_all_ddl)
+            retry(db.reload)()
+
+        session = db.session()
         session.create()
         self.to_delete.append(session)
 
-        with session.transaction() as transaction:
-            transaction.delete(self.TABLE, self.ALL)
+        def _unit_of_work(transaction, test):
+            transaction.delete(test.TABLE, test.ALL)
             transaction.insert(
-                self.TABLE, self.COLUMNS, self._row_data(row_count))
+                test.TABLE, test.COLUMNS, test._row_data(row_count))
+
+        committed = session.run_in_transaction(_unit_of_work, test=self)
 
-        return session, transaction.committed
+        return session, committed
 
     def test_read_w_manual_consume(self):
         ROW_COUNT = 4000
@@ -531,6 +536,63 @@ def test_read_w_manual_consume(self):
         self.assertEqual(streamed._current_row, [])
         self.assertEqual(streamed._pending_chunk, None)
 
+    def test_read_w_index(self):
+        ROW_COUNT = 2000
+        # Indexed reads cannot return non-indexed columns
+        MY_COLUMNS = self.COLUMNS[0], self.COLUMNS[2]
+        EXTRA_DDL = [
+            'CREATE INDEX contacts_by_last_name ON contacts(last_name)',
+        ]
+        pool = BurstyPool()
+        temp_db = Config.INSTANCE.database(
+            'test_read_w_index', ddl_statements=DDL_STATEMENTS + EXTRA_DDL,
+            pool=pool)
+        operation = temp_db.create()
+        self.to_delete.append(_DatabaseDropper(temp_db))
+
+        # We want to make sure the operation completes.
+        operation.result(30)  # raises on failure / timeout.
+
+        session, committed = self._set_up_table(ROW_COUNT, db=temp_db)
+
+        snapshot = session.snapshot(read_timestamp=committed)
+        rows = list(snapshot.read(
+            self.TABLE, MY_COLUMNS, self.ALL, index='contacts_by_last_name'))
+
+        expected = list(reversed(
+            [(row[0], row[2]) for row in self._row_data(ROW_COUNT)]))
+        self._check_row_data(rows, expected)
+
+    def test_read_w_limit(self):
+        ROW_COUNT = 4000
+        LIMIT = 100
+        session, committed = self._set_up_table(ROW_COUNT)
+
+        snapshot = session.snapshot(read_timestamp=committed)
+        rows = list(snapshot.read(
+            self.TABLE, self.COLUMNS, self.ALL, limit=LIMIT))
+
+        all_data_rows = list(self._row_data(ROW_COUNT))
+        expected = all_data_rows[:LIMIT]
+        self._check_row_data(rows, expected)
+
+    def test_read_w_range(self):
+        from google.cloud.spanner.keyset import KeyRange
+        ROW_COUNT = 4000
+        START_CLOSED = 1000
+        END_OPEN = 2000
+        session, committed = self._set_up_table(ROW_COUNT)
+        key_range = KeyRange(start_closed=[START_CLOSED], end_open=[END_OPEN])
+        keyset = KeySet(ranges=(key_range,))
+
+        snapshot = session.snapshot(read_timestamp=committed)
+        rows = list(snapshot.read(
+            self.TABLE, self.COLUMNS, keyset))
+
+        all_data_rows = list(self._row_data(ROW_COUNT))
+        expected = all_data_rows[START_CLOSED:END_OPEN]
+        self._check_row_data(rows, expected)
+
     def test_execute_sql_w_manual_consume(self):
         ROW_COUNT = 4000
         session, committed = self._set_up_table(ROW_COUNT)
@@ -662,3 +724,13 @@ def test_execute_sql_w_query_param(self):
             param_types={'my_list': array_type},
             expected=[(u'dog',), (u'cat',)],
         )
+
+
+class _DatabaseDropper(object):
+    """Helper for cleaning up databases created on-the-fly."""
+
+    def __init__(self, db):
+        self._db = db
+
+    def delete(self):
+        self._db.drop()
