Repository: googleapis/python-spanner-django

Issue:
django: Cursor.description returns None on raw queries
I'm not familiar enough with [the code](https://github.com/orijtech/spanner-orm/blob/f0d203df800a52243ef2410b141fa2fd0e313525/spanner/dbapi/cursor.py#L38-L47) to say if this is a bug or if it's a limitation of Spanner. Example failure:
```
======================================================================
ERROR: test_FK_raw_query (raw_query.tests.RawQueryTests)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/tim/code/django/tests/raw_query/tests.py", line 124, in test_FK_raw_query
    self.assertSuccessfulRawQuery(Book, query, books)
  File "/home/tim/code/django/tests/raw_query/tests.py", line 53, in assertSuccessfulRawQuery
    results = list(model.objects.raw(query, params=params, translations=translations))
  File "/home/tim/code/django/django/db/models/query.py", line 1395, in __iter__
    self._fetch_all()
  File "/home/tim/code/django/django/db/models/query.py", line 1382, in _fetch_all
    self._result_cache = list(self.iterator())
  File "/home/tim/code/django/django/db/models/query.py", line 1408, in iterator
    model_init_names, model_init_pos, annotation_fields = self.resolve_model_init_order()
  File "/home/tim/code/django/django/db/models/query.py", line 1351, in resolve_model_init_order
    model_init_fields = [f for f in self.model._meta.fields if converter(f.column) in self.columns]
  File "/home/tim/code/django/django/db/models/query.py", line 1351, in <listcomp>
    model_init_fields = [f for f in self.model._meta.fields if converter(f.column) in self.columns]
  File "/home/tim/code/django/django/utils/functional.py", line 80, in __get__
    res = instance.__dict__[self.name] = self.func(instance)
  File "/home/tim/code/django/django/db/models/query.py", line 1457, in columns
    columns = self.query.get_columns()
  File "/home/tim/code/django/django/db/models/sql/query.py", line 97, in get_columns
    for column_meta in self.cursor.description]
TypeError: 'NoneType' object is not iterable
```


Discussion:
This issue comes up because after a DQL statement from .execute_sql, Cloud Spanner Python returns a [google.cloud.spanner_v1.streamed.StreamedResultSet](https://googleapis.dev/python/spanner/latest/streamed-api.html#module-google.cloud.spanner_v1.streamed) of which to get the metadata, we have to read at least one element for the metadata to be populated and here is an example
```python
from spanner.dbapi import connect

instance_url = 'cloudspanner:/projects/orijtech-161805/instances/django-tests/databases/travisdb1?instanceConfig=projects/orijtech-161805/instanceConfigs/us-west2'

conn = connect(instance_url)
with conn.cursor() as cur:
        for sql, params in inserts:
            cur.execute(sql, params)

        print('Without fetching single item')
        cur.execute('SELECT * from raw_query_book')
        print('description ', cur.description)

        print('\nWith fetching single item')
        res = cur.execute('SELECT * from raw_query_book')
        cur.fetchone()
        print('description ', cur.description)
   
 conn.close()
```
which prints out
```shell
Without fetching single item
description  None

With fetching single item
description  (Column(name='id', type_code=2), Column(name='title', type_code=6), Column(name='author_id', type_code=2), Column(name='paperback', type_code=1), Column(name='opening_line', type_code=6))
```

What we can do is perhaps always try to read a single item after any DQL, then place it back as the head of the queue or have a single variable called `peeked` that retains the first read item or some sort of merging iterator that reads from the head and the rest of the prior iterator.

Resolving patch:
diff --git a/spanner/dbapi/cursor.py b/spanner/dbapi/cursor.py
--- a/spanner/dbapi/cursor.py
+++ b/spanner/dbapi/cursor.py
@@ -84,6 +84,8 @@ def execute(self, sql, args=None):
         # ValueError("Specify 'param_types' when passing 'params'.").
         # See https://github.com/orijtech/spanner-orm/issues/35
         param_types = {} if args else None
+        self.__res = None
+
         # Classify whether this is a read-only SQL statement.
         try:
             classification = classify_stmt(sql)
@@ -196,8 +198,9 @@ def __handle_DQL(self, sql, params, param_types=None):
                 # many items returns a RuntimeError if .fetchone() is
                 # invoked and vice versa.
                 self.__res = res
-                self.__itr = iter(self.__res)
-
+                # Read the first element so that StreamedResult can
+                # return the metadata after a DQL statement. See issue #155.
+                self.__itr = PeekIterator(self.__res)
                 # Unfortunately, Spanner doesn't seem to send back
                 # information about the number of rows available.
                 self.__row_count = _UNSET_COUNT
@@ -293,6 +296,22 @@ def __init__(self, name, type_code, display_size=None, internal_size=None,
     def __repr__(self):
         return self.__str__()
 
+    def __getitem__(self, index):
+        if index == 0:
+            return self.name
+        elif index == 1:
+            return self.type_code
+        elif index == 2:
+            return self.display_size
+        elif index == 3:
+            return self.internal_size
+        elif index == 4:
+            return self.precision
+        elif index == 5:
+            return self.scale
+        elif index == 6:
+            return self.null_ok
+
     def __str__(self):
         rstr = ', '.join([field for field in [
             "name='%s'" % self.name,
@@ -305,3 +324,40 @@ def __str__(self):
         ] if field])
 
         return 'Column(%s)' % rstr
+
+
+class PeekIterator(object):
+    """
+    PeekIterator peeks at the first element out of an iterator
+    for the sake of operations like auto-population of fields on reading
+    the first element.
+    """
+    def __init__(self, source):
+        itr_src = iter(source)
+
+        self.__iters = []
+        self.__index = 0
+
+        try:
+            head = next(itr_src)
+            # Restitch and prepare to read from multiple iterators.
+            self.__iters = [iter(itr) for itr in [[head], itr_src]]
+        except StopIteration:
+            pass
+
+    def __next__(self):
+        if self.__index >= len(self.__iters):
+            raise StopIteration
+
+        iterator = self.__iters[self.__index]
+        try:
+            head = next(iterator)
+        except StopIteration:
+            # That iterator has been exhausted, try with the next one.
+            self.__index += 1
+            return self.__next__()
+        else:
+            return head
+
+    def __iter__(self):
+        return self
diff --git a/spanner/django/features.py b/spanner/django/features.py
--- a/spanner/django/features.py
+++ b/spanner/django/features.py
@@ -49,6 +49,8 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         'queries.tests.SubqueryTests.test_related_sliced_subquery',
         'queries.tests.Ticket14056Tests.test_ticket_14056',
         'queries.tests.RelatedLookupTypeTests.test_values_queryset_lookup',
+        'raw_query.tests.RawQueryTests.test_annotations',
+        'raw_query.tests.RawQueryTests.test_get_item',
         'syndication_tests.tests.SyndicationFeedTest.test_rss2_feed',
         'syndication_tests.tests.SyndicationFeedTest.test_latest_post_date',
         'syndication_tests.tests.SyndicationFeedTest.test_rss091_feed',
@@ -59,10 +61,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         # datetimes retrieved from the database with the wrong hour when
         # USE_TZ = True: https://github.com/orijtech/spanner-orm/issues/193
         'datetimes.tests.DateTimesTests.test_21432',
-        # Cursor.description returns None on raw queries:
-        # https://github.com/orijtech/spanner-orm/issues/155
-        'timezones.tests.LegacyDatabaseTests.test_raw_sql',
-        'timezones.tests.NewDatabaseTests.test_raw_sql',
         # Unable to infer type for parameter:
         # https://github.com/orijtech/spanner-orm/issues/185
         'timezones.tests.LegacyDatabaseTests.test_cursor_execute_returns_naive_datetime',


Test patch:
diff --git a/tests/spanner/dbapi/test_types.py b/tests/spanner/dbapi/test_types.py
--- a/tests/spanner/dbapi/test_types.py
+++ b/tests/spanner/dbapi/test_types.py
@@ -15,6 +15,7 @@
 import datetime
 from unittest import TestCase
 
+from spanner.dbapi.cursor import PeekIterator
 from spanner.dbapi.types import (
     Date, DateFromTicks, Time, TimeFromTicks, Timestamp, TimestampFromTicks,
 )
@@ -76,3 +77,18 @@ def test_TimestampFromTicks(self):
         )
         matches = got in want
         self.assertTrue(matches, '`%s` not present in any of\n`%s`' % (got, want))
+
+    def test_PeekIterator(self):
+        cases = [
+            ('list', [1, 2, 3, 4, 6, 7], [1, 2, 3, 4, 6, 7]),
+            ('iter_from_list', iter([1, 2, 3, 4, 6, 7]), [1, 2, 3, 4, 6, 7]),
+            ('tuple', ('a', 12, 0xff,), ['a', 12, 0xff]),
+            ('iter_from_tuple', iter(('a', 12, 0xff,)), ['a', 12, 0xff]),
+            ('no_args', (), []),
+        ]
+
+        for name, data_in, want in cases:
+            with self.subTest(name=name):
+                pitr = PeekIterator(data_in)
+                got = list(pitr)
+                self.assertEqual(got, want)