Merge pull request googleapis#16 from tswast/patch-1

mxmzdlv · web-flow · commit d214bb089ca0 · 2018-08-11T16:14:29.000+12:00
Add support for STRUCT and ARRAY types.
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
@@ -52,9 +52,8 @@ def format_label(self, *args, **kwargs):
     'DATETIME': types.DATETIME,
     'DATE': types.DATE,
     'BYTES': types.BINARY,
-    'TIME': types.TIME
-    # TODO
-    # 'RECORD'
+    'TIME': types.TIME,
+    'RECORD': types.JSON,
 }
 
 
@@ -116,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
@@ -165,8 +172,8 @@ def get_columns(self, connection, table_name, schema=None, **kw):
 
             result.append({
                 'name': col.name,
-                'type': coltype,
-                'nullable': True if col.mode == 'NULLABLE' else False,
+                '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]}`