Add tests for RECORD and REPEATED types.

tswast · tswast · commit e5a07b93a24c · 2018-08-10T15:42:07.000-07:00
diff --git a/README.rst b/README.rst
@@ -70,8 +70,3 @@ Set up an environment and run tests::
     source .env/bin/activate
     pip install -r dev_requirements.txt
     pytest
-
-TODO
-====
-
-- Support for Record column type
diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py
@@ -115,6 +115,14 @@ def create_connect_args(self, url):
             client = bigquery.Client(url.host)
         return ([client], {})
 
+    def _json_deserializer(self, row):
+        """JSON deserializer for RECORD types.
+
+        The DB-API layer already deserializes JSON to a dictionary, so this
+        just returns the input.
+        """
+        return row
+
     def _split_table_name(self, full_table_name):
         # Split full_table_name to get project, dataset and table name
         dataset = None
@@ -164,7 +172,7 @@ def get_columns(self, connection, table_name, schema=None, **kw):
 
             result.append({
                 'name': col.name,
-                'type': types.ARRAY if col.mode == 'REPEATED' else coltype,
+                'type': types.ARRAY(coltype) if col.mode == 'REPEATED' else coltype,
                 'nullable': col.mode == 'NULLABLE' or col.mode == 'REPEATED',
                 'default': None,
             })
diff --git a/scripts/sample_one_row.json b/scripts/sample_one_row.json
@@ -1 +1 @@
-{"integer":"588","timestamp":"1381404436","string":"W 52 St & 11 Ave","float":"40.76727216","boolean":"false","date":"2013-10-10","datetime":"2013-10-10T11:27:16","time":"11:27:16","bytes":"7w=="}
+{"integer":"588","timestamp":"1381404436","string":"W 52 St & 11 Ave","float":"40.76727216","boolean":"false","date":"2013-10-10","datetime":"2013-10-10T11:27:16","time":"11:27:16","bytes":"7w==","record":{"name": "John Doe","age":"100"},"array":[1,2,3]}
diff --git a/scripts/schema.json b/scripts/schema.json
@@ -42,4 +42,26 @@
   "mode": "NULLABLE",
   "name": "bytes",
   "type": "BYTES"
+},
+{
+  "mode": "NULLABLE",
+  "name": "record",
+  "type": "RECORD",
+  "fields": [
+    {
+      "mode": "NULLABLE",
+      "name": "name",
+      "type": "STRING"
+    },
+    {
+      "mode": "NULLABLE",
+      "name": "age",
+      "type": "INTEGER"
+    }
+  ]
+},
+{
+  "mode": "REPEATED",
+  "name": "array",
+  "type": "INTEGER"
 }]
diff --git a/test/test_sqlalchemy_bigquery.py b/test/test_sqlalchemy_bigquery.py
@@ -22,7 +22,12 @@
     datetime.date(2013, 10, 10),
     datetime.datetime(2013, 10, 10, 11, 27, 16),
     datetime.time(11, 27, 16),
-    b'\xef'
+    b'\xef',
+    {
+        'name': 'John Doe',
+        'age': 100,
+    },
+    [1, 2, 3],
 ]
 
 
@@ -47,7 +52,9 @@
     {'name': 'date', 'type': types.DATE(), 'nullable': True, 'default': None},
     {'name': 'datetime', 'type': types.DATETIME(), 'nullable': True, 'default': None},
     {'name': 'time', 'type': types.TIME(), 'nullable': True, 'default': None},
-    {'name': 'bytes', 'type': types.BINARY(), 'nullable': True, 'default': None}
+    {'name': 'bytes', 'type': types.BINARY(), 'nullable': True, 'default': None},
+    {'name': 'record', 'type': types.JSON(), 'nullable': True, 'default': None},
+    {'name': 'array', 'type': types.ARRAY(types.Integer()), 'nullable': True, 'default': None},
 ]
 
 
@@ -101,7 +108,7 @@ def query(table):
 
 
 def test_reflect_select(engine, table):
-    assert len(table.c) == 9
+    assert len(table.c) == 11
     assert isinstance(table.c.integer, Column)
     assert isinstance(table.c.integer.type, types.Integer)
     assert isinstance(table.c.timestamp.type, types.TIMESTAMP)
@@ -112,6 +119,8 @@ def test_reflect_select(engine, table):
     assert isinstance(table.c.datetime.type, types.DATETIME)
     assert isinstance(table.c.time.type, types.TIME)
     assert isinstance(table.c.bytes.type, types.BINARY)
+    assert isinstance(table.c.record.type, types.JSON)
+    assert isinstance(table.c.array.type, types.ARRAY)
 
     rows = table.select().execute().fetchall()
     assert len(rows) == 1000

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-{"integer":"588","timestamp":"1381404436","string":"W 52 St & 11 Ave","float":"40.76727216","boolean":"false","date":"2013-10-10","datetime":"2013-10-10T11:27:16","time":"11:27:16","bytes":"7w=="}`
	`1`	`+{"integer":"588","timestamp":"1381404436","string":"W 52 St & 11 Ave","float":"40.76727216","boolean":"false","date":"2013-10-10","datetime":"2013-10-10T11:27:16","time":"11:27:16","bytes":"7w==","record":{"name": "John Doe","age":"100"},"array":[1,2,3]}`