[Issue 171] Improve sqlalchemy dialect backward compatibility with 1.3.24

examples/sqlalchemy.py

@@ -42,9 +42,15 @@
 """
 
 import os
-from sqlalchemy.orm import declarative_base, Session
+import sqlalchemy
+from sqlalchemy.orm import Session
 from sqlalchemy import Column, String, Integer, BOOLEAN, create_engine, select
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+
 host = os.getenv("DATABRICKS_SERVER_HOSTNAME")
 http_path = os.getenv("DATABRICKS_HTTP_PATH")
 access_token = os.getenv("DATABRICKS_TOKEN")
@@ -59,10 +65,20 @@
     "_user_agent_entry": "PySQL Example Script",
 }
 
-engine = create_engine(
-    f"databricks://token:{access_token}@{host}?http_path={http_path}&catalog={catalog}&schema={schema}",
-    connect_args=extra_connect_args,
-)
+if sqlalchemy.__version__.startswith("1.3"):
+    # SQLAlchemy 1.3.x fails to parse the http_path, catalog, and schema from our connection string
+    # Pass these in as connect_args instead
+
+    conn_string = f"databricks://token:{access_token}@{host}"
+    connect_args = dict(catalog=catalog, schema=schema, http_path=http_path)
+    all_connect_args = {**extra_connect_args, **connect_args}
+    engine = create_engine(conn_string, connect_args=all_connect_args)
+else:
+    engine = create_engine(
+        f"databricks://token:{access_token}@{host}?http_path={http_path}&catalog={catalog}&schema={schema}",
+        connect_args=extra_connect_args,
+    )
+
 session = Session(bind=engine)
 base = declarative_base(bind=engine)
 
@@ -73,7 +89,7 @@ class SampleObject(base):
 
     name = Column(String(255), primary_key=True)
     episodes = Column(Integer)
-    some_bool = Column(BOOLEAN)
+    some_bool = Column(BOOLEAN(create_constraint=False))
 
 
 base.metadata.create_all()
@@ -86,9 +102,14 @@ class SampleObject(base):
 
 session.commit()
 
-stmt = select(SampleObject).where(SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"]))
+# SQLAlchemy 1.3 has slightly different methods
+if sqlalchemy.__version__.startswith("1.3"):
+    stmt = select([SampleObject]).where(SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"]))
+    output = [i for i in session.execute(stmt)]
+else:
+    stmt = select(SampleObject).where(SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"]))
+    output = [i for i in session.scalars(stmt)]
 
-output = [i for i in session.scalars(stmt)]
 assert len(output) == 2
 
 base.metadata.drop_all()

src/databricks/sqlalchemy/dialect/__init__.py

@@ -4,9 +4,10 @@
 import decimal, re, datetime
 from dateutil.parser import parse
 
+import sqlalchemy
 from sqlalchemy import types, processors, event
 from sqlalchemy.engine import default, Engine
-from sqlalchemy.exc import DatabaseError
+from sqlalchemy.exc import DatabaseError, SQLAlchemyError
 from sqlalchemy.engine import reflection
 
 from databricks import sql
@@ -153,9 +154,7 @@ def get_columns(self, connection, table_name, schema=None, **kwargs):
             "date": DatabricksDate,
         }
 
-        with self.get_driver_connection(
-            connection
-        )._dbapi_connection.dbapi_connection.cursor() as cur:
+        with self.get_connection_cursor(connection) as cur:
             resp = cur.columns(
                 catalog_name=self.catalog,
                 schema_name=schema or self.schema,
@@ -244,9 +243,7 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
 
     def get_table_names(self, connection, schema=None, **kwargs):
         TABLE_NAME = 1
-        with self.get_driver_connection(
-            connection
-        )._dbapi_connection.dbapi_connection.cursor() as cur:
+        with self.get_connection_cursor(connection) as cur:
             sql_str = "SHOW TABLES FROM {}".format(
                 ".".join([self.catalog, schema or self.schema])
             )
@@ -257,9 +254,7 @@ def get_table_names(self, connection, schema=None, **kwargs):
 
     def get_view_names(self, connection, schema=None, **kwargs):
         VIEW_NAME = 1
-        with self.get_driver_connection(
-            connection
-        )._dbapi_connection.dbapi_connection.cursor() as cur:
+        with self.get_connection_cursor(connection) as cur:
             sql_str = "SHOW VIEWS FROM {}".format(
                 ".".join([self.catalog, schema or self.schema])
             )
@@ -291,6 +286,21 @@ def has_table(self, connection, table_name, schema=None, **kwargs) -> bool:
                 return False
             else:
                 raise e
+
+
+    def get_connection_cursor(self, connection):
+        """Added for backwards compatibility with 1.3.x
+        """
+        if hasattr(connection, "_dbapi_connection"):
+            return connection._dbapi_connection.dbapi_connection.cursor()
+        elif hasattr(connection, "raw_connection"):
+            return connection.raw_connection().cursor()
+        elif hasattr(connection, "connection"):
+            return connection.connection.cursor()
+
+        raise SQLAlchemyError("Databricks dialect can't obtain a cursor context manager from the dbapi")
+
+
 
     @reflection.cache
     def get_schema_names(self, connection, **kw):
@@ -314,3 +324,13 @@ def receive_do_connect(dialect, conn_rec, cargs, cparams):
         new_user_agent = "sqlalchemy"
 
     cparams["_user_agent_entry"] = new_user_agent
+
+    if sqlalchemy.__version__.startswith("1.3"):
+        # SQLAlchemy 1.3.x fails to parse the http_path, catalog, and schema from our connection string
+        # These should be passed in as connect_args when building the Engine
+
+        if "schema" in cparams:
+            dialect.schema = cparams["schema"]
+
+        if "catalog" in cparams:
+            dialect.catalog = cparams["catalog"]

tests/e2e/sqlalchemy/test_basic.py

@@ -2,12 +2,34 @@
 import pytest
 from unittest import skipIf
 from sqlalchemy import create_engine, select, insert, Column, MetaData, Table
-from sqlalchemy.orm import declarative_base, Session
+from sqlalchemy.orm import Session
 from sqlalchemy.types import SMALLINT, Integer, BOOLEAN, String, DECIMAL, Date
 
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+
 
 USER_AGENT_TOKEN = "PySQL e2e Tests"
 
+def sqlalchemy_1_3():
+    import sqlalchemy
+    return sqlalchemy.__version__.startswith("1.3")
+
+def version_agnostic_select(object_to_select, *args, **kwargs):
+    """
+    SQLAlchemy==1.3.x requires arguments to select() to be a Python list
+
+    https://docs.sqlalchemy.org/en/20/changelog/migration_14.html#orm-query-is-internally-unified-with-select-update-delete-2-0-style-execution-available
+    """
+
+    if sqlalchemy_1_3():
+        return select([object_to_select], *args, **kwargs)
+    else:
+        return select(object_to_select, *args, **kwargs)
+
+
 
 @pytest.fixture
 def db_engine():
@@ -19,6 +41,20 @@ def db_engine():
     SCHEMA = os.environ.get("schema")
 
     connect_args = {"_user_agent_entry": USER_AGENT_TOKEN}
+
+    connect_args = {
+        **connect_args,
+        "http_path": HTTP_PATH,
+        "server_hostname": HOST,
+        "catalog": CATALOG,
+        "schema": SCHEMA
+    }
+
+    engine = create_engine(
+        f"databricks://token:{ACCESS_TOKEN}@{HOST}",
+        connect_args=connect_args,
+    )
+    return engine
 
     engine = create_engine(
         f"databricks://token:{ACCESS_TOKEN}@{HOST}?http_path={HTTP_PATH}&catalog={CATALOG}&schema={SCHEMA}",
@@ -62,6 +98,7 @@ def test_connect_args(db_engine):
     assert expected in user_agent
 
 
+@pytest.mark.skipif(sqlalchemy_1_3(), reason="Pandas requires SQLAlchemy >= 1.4")
 def test_pandas_upload(db_engine, metadata_obj):
 
     import pandas as pd
@@ -95,7 +132,7 @@ def test_create_table_not_null(db_engine, metadata_obj):
         metadata_obj,
         Column("name", String(255)),
         Column("episodes", Integer),
-        Column("some_bool", BOOLEAN, nullable=False),
+        Column("some_bool", BOOLEAN(create_constraint=False), nullable=False),
     )
 
     metadata_obj.create_all()
@@ -135,7 +172,7 @@ def test_bulk_insert_with_core(db_engine, metadata_obj, session):
     metadata_obj.create_all()
     db_engine.execute(insert(SampleTable).values(rows))
 
-    rows = db_engine.execute(select(SampleTable)).fetchall()
+    rows = db_engine.execute(version_agnostic_select(SampleTable)).fetchall()
 
     assert len(rows) == num_to_insert
 
@@ -148,7 +185,7 @@ def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
         metadata_obj,
         Column("name", String(255)),
         Column("episodes", Integer),
-        Column("some_bool", BOOLEAN),
+        Column("some_bool", BOOLEAN(create_constraint=False)),
         Column("dollars", DECIMAL(10, 2)),
     )
 
@@ -161,7 +198,7 @@ def test_create_insert_drop_table_core(base, db_engine, metadata_obj: MetaData):
     with db_engine.connect() as conn:
         conn.execute(insert_stmt)
 
-    select_stmt = select(SampleTable)
+    select_stmt = version_agnostic_select(SampleTable)
     resp = db_engine.execute(select_stmt)
 
     result = resp.fetchall()
@@ -187,7 +224,7 @@ class SampleObject(base):
 
         name = Column(String(255), primary_key=True)
         episodes = Column(Integer)
-        some_bool = Column(BOOLEAN)
+        some_bool = Column(BOOLEAN(create_constraint=False))
 
     base.metadata.create_all()
 
@@ -197,11 +234,15 @@ class SampleObject(base):
     session.add(sample_object_2)
     session.commit()
 
-    stmt = select(SampleObject).where(
+    stmt = version_agnostic_select(SampleObject).where(
         SampleObject.name.in_(["Bim Adewunmi", "Miki Meek"])
     )
 
-    output = [i for i in session.scalars(stmt)]
+    if sqlalchemy_1_3():
+        output = [i for i in session.execute(stmt)]
+    else:
+        output = [i for i in session.scalars(stmt)]
+
     assert len(output) == 2
 
     base.metadata.drop_all()
@@ -215,7 +256,7 @@ def test_dialect_type_mappings(base, db_engine, metadata_obj: MetaData):
         metadata_obj,
         Column("string_example", String(255)),
         Column("integer_example", Integer),
-        Column("boolean_example", BOOLEAN),
+        Column("boolean_example", BOOLEAN(create_constraint=False)),
         Column("decimal_example", DECIMAL(10, 2)),
         Column("date_example", Date),
     )
@@ -239,7 +280,7 @@ def test_dialect_type_mappings(base, db_engine, metadata_obj: MetaData):
     with db_engine.connect() as conn:
         conn.execute(insert_stmt)
 
-    select_stmt = select(SampleTable)
+    select_stmt = version_agnostic_select(SampleTable)
     resp = db_engine.execute(select_stmt)
 
     result = resp.fetchall()