Fetch table and column descriptions (#82)

* Fetch table and column descriptions

* Reset coltype for unrecognized columns

* add tests

Co-authored-by: Tim Swast <swast@google.com>

Author: hsheth2

dev_requirements.txt

@@ -1,7 +1,7 @@
 sqlalchemy>=1.1.9
 google-cloud-bigquery>=1.6.0
-future==0.16.0
+future==0.18.2
 
-pytest==3.2.2
-pytest-flake8==1.0.6
-pytz==2017.2
+pytest==6.2.2
+pytest-flake8==1.0.7
+pytz==2021.1

pybigquery/sqlalchemy_bigquery.py

@@ -495,16 +495,24 @@ def get_columns(self, connection, table_name, schema=None, **kw):
                 coltype = _type_map[col.field_type]
             except KeyError:
                 util.warn("Did not recognize type '%s' of column '%s'" % (col.field_type, col.name))
+                coltype = types.NullType
 
             result.append({
                 'name': col.name,
                 'type': types.ARRAY(coltype) if col.mode == 'REPEATED' else coltype,
                 'nullable': col.mode == 'NULLABLE' or col.mode == 'REPEATED',
+                'comment': col.description,
                 'default': None,
             })
 
         return result
 
+    def get_table_comment(self, connection, table_name, schema=None, **kw):
+        table = self._get_table(connection, table_name, schema)
+        return {
+            'text': table.description,
+        }
+
     def get_foreign_keys(self, connection, table_name, schema=None, **kw):
         # BigQuery has no support for foreign keys.
         return []

scripts/load_test_data.sh

@@ -9,7 +9,7 @@ bq rm -f -t test_pybigquery.sample_dml
 bq rm -f -t test_pybigquery.sample_view
 bq rm -f -t test_pybigquery_location.sample_one_row
 
-bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery.sample
+bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string --description 'A sample table containing most data types.' test_pybigquery.sample
 bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery_alt.sample_alt
 bq load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery.sample $(dirname $0)/sample.json
 

scripts/schema.json

@@ -52,6 +52,7 @@
   "mode": "NULLABLE",
   "name": "record",
   "type": "RECORD",
+  "description": "In Standard SQL this data type is a STRUCT<name STRING, age INT64>.",
   "fields": [
     {
       "mode": "NULLABLE",

test/test_sqlalchemy_bigquery.py

@@ -86,7 +86,13 @@
     {'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': 'record', 'type': types.JSON(), 'nullable': True, 'default': None},
+    {
+        'name': 'record',
+        'type': types.JSON(),
+        'nullable': True,
+        'default': None,
+        'comment': 'In Standard SQL this data type is a STRUCT<name STRING, age INT64>.',
+    },
     {'name': 'record.name', 'type': types.String(), 'nullable': True, 'default': None},
     {'name': 'record.age', 'type': types.Integer(), 'nullable': True, 'default': None},
     {'name': 'nested_record', 'type': types.JSON(), 'nullable': True, 'default': None},
@@ -225,8 +231,9 @@ def test_dataset_location(engine_with_location):
 
 def test_reflect_select(table, table_using_test_dataset):
     for table in [table, table_using_test_dataset]:
-        assert len(table.c) == 18
+        assert table.comment == "A sample table containing most data types."
 
+        assert len(table.c) == 18
         assert isinstance(table.c.integer, Column)
         assert isinstance(table.c.integer.type, types.Integer)
         assert isinstance(table.c.timestamp.type, types.TIMESTAMP)
@@ -526,9 +533,10 @@ def test_get_columns(inspector, inspector_using_test_dataset):
     for columns in columns_queries:
         for i, col in enumerate(columns):
             sample_col = SAMPLE_COLUMNS[i]
+            assert col['comment'] == sample_col.get('comment')
+            assert col['default'] == sample_col['default']
             assert col['name'] == sample_col['name']
             assert col['nullable'] == sample_col['nullable']
-            assert col['default'] == sample_col['default']
             assert col['type'].__class__.__name__ == sample_col['type'].__class__.__name__
 
     columns_without_schema = inspector_using_test_dataset.get_columns('sample')
@@ -537,9 +545,10 @@ def test_get_columns(inspector, inspector_using_test_dataset):
     for columns in columns_queries:
         for i, col in enumerate(columns):
             sample_col = SAMPLE_COLUMNS[i]
+            assert col['comment'] == sample_col.get('comment')
+            assert col['default'] == sample_col['default']
             assert col['name'] == sample_col['name']
             assert col['nullable'] == sample_col['nullable']
-            assert col['default'] == sample_col['default']
             assert col['type'].__class__.__name__ == sample_col['type'].__class__.__name__