diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py index d13b0072..ba491fc4 100644 --- a/pybigquery/sqlalchemy_bigquery.py +++ b/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 diff --git a/test/test_sqlalchemy_bigquery.py b/test/test_sqlalchemy_bigquery.py index 7633a421..acb2f82c 100644 --- a/test/test_sqlalchemy_bigquery.py +++ b/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