Merge pull request #32 from jmcarp/type-compiler

Add type compiler.

Author: mxmzdlv

packages/sqlalchemy-bigquery/pybigquery/sqlalchemy_bigquery.py

@@ -12,7 +12,7 @@
 from google.api_core.exceptions import NotFound
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy import types, util
-from sqlalchemy.sql.compiler import SQLCompiler, IdentifierPreparer
+from sqlalchemy.sql.compiler import SQLCompiler, GenericTypeCompiler, DDLCompiler, IdentifierPreparer
 from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.sql.schema import Column
@@ -179,11 +179,45 @@ def visit_label(self, *args, **kwargs):
         return result
 
 
+class BigQueryTypeCompiler(GenericTypeCompiler):
+
+    def visit_integer(self, type_, **kw):
+        return 'INT64'
+
+    def visit_float(self, type_, **kw):
+        return 'FLOAT64'
+
+    def visit_text(self, type_, **kw):
+        return 'STRING'
+
+    def visit_string(self, type_, **kw):
+        return 'STRING'
+
+    def visit_BINARY(self, type_, **kw):
+        return 'BYTES'
+
+    def visit_DECIMAL(self, type_, **kw):
+        return 'NUMERIC'
+
+
+class BigQueryDDLCompiler(DDLCompiler):
+
+    # BigQuery has no support for foreign keys.
+    def visit_foreign_key_constraint(self, constraint):
+        return None
+
+    # BigQuery has no support for primary keys.
+    def visit_primary_key_constraint(self, constraint):
+        return None
+
+
 class BigQueryDialect(DefaultDialect):
     name = 'bigquery'
     driver = 'bigquery'
     preparer = BigQueryIdentifierPreparer
     statement_compiler = BigQueryCompiler
+    type_compiler = BigQueryTypeCompiler
+    ddl_compiler = BigQueryDDLCompiler
     execution_ctx_cls = BigQueryExecutionContext
     supports_alter = False
     supports_pk_autoincrement = False

packages/sqlalchemy-bigquery/test/test_sqlalchemy_bigquery.py

@@ -5,6 +5,7 @@
 from pybigquery.api import ApiClient
 from sqlalchemy.engine import create_engine
 from sqlalchemy.schema import Table, MetaData, Column
+from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import types, func, case, inspect
 from sqlalchemy.sql import expression, select, literal_column
 from sqlalchemy.exc import NoSuchTableError
@@ -382,6 +383,36 @@ def test_dml(engine, session, table_dml):
     assert len(result) == 0
 
 
+def test_create_table(engine):
+    meta = MetaData()
+    table = Table(
+        'test_pybigquery.test_table_create', meta,
+        Column('integer_c', sqlalchemy.Integer),
+        Column('float_c', sqlalchemy.Float),
+        Column('decimal_c', sqlalchemy.DECIMAL),
+        Column('string_c', sqlalchemy.String),
+        Column('text_c', sqlalchemy.Text),
+        Column('boolean_c', sqlalchemy.Boolean),
+        Column('timestamp_c', sqlalchemy.TIMESTAMP),
+        Column('datetime_c', sqlalchemy.DATETIME),
+        Column('date_c', sqlalchemy.DATE),
+        Column('time_c', sqlalchemy.TIME),
+        Column('binary_c', sqlalchemy.BINARY)
+    )
+    meta.create_all(engine)
+    meta.drop_all(engine)
+
+    # Test creating tables with declarative_base
+    Base = declarative_base()
+
+    class TableTest(Base):
+        __tablename__ = 'test_pybigquery.test_table_create2'
+        integer_c = Column(sqlalchemy.Integer, primary_key=True)
+        float_c = Column(sqlalchemy.Float)
+
+    Base.metadata.create_all(engine)
+    Base.metadata.drop_all(engine)
+
 def test_schemas_names(inspector, inspector_using_test_dataset):
     datasets = inspector.get_schema_names()
     assert 'test_pybigquery' in datasets