A couple of changes to make this connector ready for Great Expectations

- Replace parameterized stmts - Dremio does not support them
- Only return schema which was passed in connection URL. If no schema is passed in URL, all schemas will be returned.
- Bugfix in get_columns, when Schema is None
- Alias dialect for DremioDialect_pyodbc created
- Version uplift to 1.2.0
- Expose supported types to pyodbc level, so that it works with "Great Expectations"

setup.py

@@ -22,7 +22,7 @@
 
 setup(
     name='sqlalchemy_dremio',
-    version='1.1.6',
+    version='1.2.0',
     description="A SQLAlchemy dialect for Dremio via the ODBC and Flight interface.",
     long_description=readme,
     long_description_content_type='text/markdown',

sqlalchemy_dremio/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '1.1.6'
+__version__ = '1.2.0'
 
 from .db import Connection, connect
 from sqlalchemy.dialects import registry

sqlalchemy_dremio/base.py

@@ -147,6 +147,7 @@ class DremioDialect(default.DefaultDialect):
     ddl_compiler = DremioDDLCompiler
     preparer = DremioIdentifierPreparer
     execution_ctx_cls = DremioExecutionContext
+    default_paramstyle = "qmark"
 
     @classmethod
     def dbapi(cls):
@@ -173,10 +174,10 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
         return []
 
     def get_columns(self, connection, table_name, schema, **kw):
-        q = "DESCRIBE \"{1}\"".format(schema, table_name)
+        sql = "DESCRIBE \"{0}\"".format(table_name)
         if schema != None and schema != "":
-            q = "DESCRIBE \"{0}\".\"{1}\"".format(schema, table_name)
-        cursor = connection.execute(q)
+            sql = "DESCRIBE \"{0}\".\"{1}\"".format(schema, table_name)
+        cursor = connection.execute(sql)
         result = []
         for col in cursor:
             cname = col[0]
@@ -204,9 +205,32 @@ def get_table_names(self, connection, schema, **kw):
         return table_names
 
     def get_schema_names(self, connection, schema=None, **kw):
+        database = connection.engine.url.database
+        if database != None and ';' in database:
+            database = database.split(';')[0]
+        # This is only returning the database from connection string.
+        # If you want to get all databases, don't pass a database name.
+        if database != None and database != "":
+            return [database]
+
         result = connection.execute("SHOW SCHEMAS")
         schema_names = [r[0] for r in result]
         return schema_names
 
     def get_view_names(self, connection, schema=None, **kwargs):
         return []
+
+    # Workaround since Dremio does not support parameterized stmts
+    # Old queries should not have used queries with parameters, since Dremio does not support it
+    # and these queries failed. If there is no parameter, everything should work as before.
+    def do_execute(self, cursor, statement, parameters, context):
+        replaced_stmt = statement
+        for v in parameters:
+            if isinstance(v, (int, float)):
+                replaced_stmt = replaced_stmt.replace('?', str(v), 1)
+            else:
+                replaced_stmt = replaced_stmt.replace('?', "'" + str(v) + "'", 1)
+
+        super(DremioDialect, self).do_execute_no_params(
+            cursor, replaced_stmt, context
+        )

sqlalchemy_dremio/pyodbc.py

@@ -2,8 +2,7 @@
 
 from sqlalchemy.connectors.pyodbc import PyODBCConnector
 
-from .base import DremioExecutionContext, DremioDialect
-
+from .base import DremioExecutionContext, DremioDialect, _type_map
 
 class DremioExecutionContext_pyodbc(DremioExecutionContext):
     pass
@@ -32,3 +31,8 @@ def create_connect_args(self, url):
                       ]
 
         return [[";".join(connectors)], connect_args]
+
+
+dialect = DremioDialect_pyodbc
+# some tools like "Great Expectations" use this to find the types on root level.
+locals().update(_type_map)