From 918a7de0d773ff24addd8fb7e082909ca6feae91 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Mon, 15 May 2023 14:55:14 -0500 Subject: [PATCH 1/6] TINYINT is specific to Databricks so must be implemented as a vendor-specific type per this link: https://docs.sqlalchemy.org/en/13/core/type_basics.html#vendor-specific-types Signed-off-by: Jesse Whitehouse From 96774b1aaca821330cc21e1b6877850d48fe5101 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 16 May 2023 10:27:19 -0500 Subject: [PATCH 2/6] Define the TINYINT type Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index da508bb09..ddc6ffe01 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -65,6 +65,15 @@ def adapt(self, impltype, **kwargs): return self.impl +class TINYINT(types.Integer): + """ + A one-byte signed integer. Can represent any integer number in the range from -128 to 127. + + Implementation copied from mssql dialect + + Details about this type: https://docs.databricks.com/sql/language-manual/data-types/tinyint-type.html + """ + __visit_name__ = "TINYINT" class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" @@ -136,6 +145,7 @@ def get_columns(self, connection, table_name, schema=None, **kwargs): _type_map = { "boolean": types.Boolean, "smallint": types.SmallInteger, + "tinyint": TINYINT, "int": types.Integer, "bigint": types.BigInteger, "float": types.Float, From b35f0c9d12c675f77b7daec48a17745774d9e9c8 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 16 May 2023 10:28:55 -0500 Subject: [PATCH 3/6] First implementation that passes the integration test. There seems to be two mechanisms where I can hook into the type compiler In this one, I explicitly set a dialect-specific compilation overwrite. Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 6 ++++++ tests/e2e/sqlalchemy/test_basic.py | 2 ++ 2 files changed, 8 insertions(+) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index ddc6ffe01..df2d603d6 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -74,6 +74,12 @@ class TINYINT(types.Integer): Details about this type: https://docs.databricks.com/sql/language-manual/data-types/tinyint-type.html """ __visit_name__ = "TINYINT" + +from sqlalchemy.ext.compiler import compiles +@compiles(TINYINT, "databricks") +def compile_tinyint_databricks(type_, compiler, **kw): + return "TINYINT" + class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" diff --git a/tests/e2e/sqlalchemy/test_basic.py b/tests/e2e/sqlalchemy/test_basic.py index 4f4df91b6..e7bf983e3 100644 --- a/tests/e2e/sqlalchemy/test_basic.py +++ b/tests/e2e/sqlalchemy/test_basic.py @@ -4,6 +4,7 @@ from sqlalchemy import create_engine, select, insert, Column, MetaData, Table from sqlalchemy.orm import declarative_base, Session from sqlalchemy.types import SMALLINT, Integer, BOOLEAN, String, DECIMAL, Date +from databricks.sqlalchemy.dialect import TINYINT USER_AGENT_TOKEN = "PySQL e2e Tests" @@ -150,6 +151,7 @@ def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData): Column("episodes", Integer), Column("some_bool", BOOLEAN), Column("dollars", DECIMAL(10, 2)), + Column("tiny_int", TINYINT) ) metadata_obj.create_all() From 71321aa3b02eeb31e4b952a615e2ec1b7e8c7617 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 16 May 2023 10:30:08 -0500 Subject: [PATCH 4/6] This is the second way to make the integration test pass using the visit mechanism. This is simpler than the approach in the commit I'm reverting. This partially reverts commit b35f0c9d12c675f77b7daec48a17745774d9e9c8. Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 6 ------ src/databricks/sqlalchemy/dialect/compiler.py | 4 ++++ 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index df2d603d6..ddc6ffe01 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -74,12 +74,6 @@ class TINYINT(types.Integer): Details about this type: https://docs.databricks.com/sql/language-manual/data-types/tinyint-type.html """ __visit_name__ = "TINYINT" - -from sqlalchemy.ext.compiler import compiles -@compiles(TINYINT, "databricks") -def compile_tinyint_databricks(type_, compiler, **kw): - return "TINYINT" - class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" diff --git a/src/databricks/sqlalchemy/dialect/compiler.py b/src/databricks/sqlalchemy/dialect/compiler.py index f77807ed4..1c3e6315e 100644 --- a/src/databricks/sqlalchemy/dialect/compiler.py +++ b/src/databricks/sqlalchemy/dialect/compiler.py @@ -1,6 +1,7 @@ from sqlalchemy.sql import compiler + class DatabricksTypeCompiler(compiler.GenericTypeCompiler): """Originally forked from pyhive""" @@ -36,3 +37,6 @@ def visit_DATE(self, type_): def visit_DATETIME(self, type_): return "TIMESTAMP" + + def visit_TINYINT(self, type_): + return "TINYINT" From 7591028ecc1e3c27f292d22040787b815c1ccf25 Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 16 May 2023 10:37:59 -0500 Subject: [PATCH 5/6] Black the source code Signed-off-by: Jesse Whitehouse --- src/databricks/sqlalchemy/dialect/__init__.py | 3 +++ src/databricks/sqlalchemy/dialect/compiler.py | 3 +-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/databricks/sqlalchemy/dialect/__init__.py b/src/databricks/sqlalchemy/dialect/__init__.py index ddc6ffe01..17d539d8a 100644 --- a/src/databricks/sqlalchemy/dialect/__init__.py +++ b/src/databricks/sqlalchemy/dialect/__init__.py @@ -73,7 +73,10 @@ class TINYINT(types.Integer): Details about this type: https://docs.databricks.com/sql/language-manual/data-types/tinyint-type.html """ + __visit_name__ = "TINYINT" + + class DatabricksDialect(default.DefaultDialect): """This dialect implements only those methods required to pass our e2e tests""" diff --git a/src/databricks/sqlalchemy/dialect/compiler.py b/src/databricks/sqlalchemy/dialect/compiler.py index 1c3e6315e..01af0ee19 100644 --- a/src/databricks/sqlalchemy/dialect/compiler.py +++ b/src/databricks/sqlalchemy/dialect/compiler.py @@ -1,7 +1,6 @@ from sqlalchemy.sql import compiler - class DatabricksTypeCompiler(compiler.GenericTypeCompiler): """Originally forked from pyhive""" @@ -37,6 +36,6 @@ def visit_DATE(self, type_): def visit_DATETIME(self, type_): return "TIMESTAMP" - + def visit_TINYINT(self, type_): return "TINYINT" From 2896665163296b64ea93bdfee41fa548dbb2715a Mon Sep 17 00:00:00 2001 From: Jesse Whitehouse Date: Tue, 16 May 2023 10:44:14 -0500 Subject: [PATCH 6/6] Add tiny_int to the INSERT test. Signed-off-by: Jesse Whitehouse --- tests/e2e/sqlalchemy/test_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/e2e/sqlalchemy/test_basic.py b/tests/e2e/sqlalchemy/test_basic.py index e7bf983e3..95fc6b9ad 100644 --- a/tests/e2e/sqlalchemy/test_basic.py +++ b/tests/e2e/sqlalchemy/test_basic.py @@ -157,7 +157,7 @@ def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData): metadata_obj.create_all() insert_stmt = insert(SampleTable).values( - name="Bim Adewunmi", episodes=6, some_bool=True, dollars=decimal.Decimal(125) + name="Bim Adewunmi", episodes=6, some_bool=True, dollars=decimal.Decimal(125), tiny_int=25 ) with db_engine.connect() as conn: