Skip to content

Add type compiler. #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Dec 21, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 35 additions & 1 deletion pybigquery/sqlalchemy_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
31 changes: 31 additions & 0 deletions test/test_sqlalchemy_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down