Fix compatibility with sqlalchemy>=2 (#117)

Author: matthewwardrop

omniduct/databases/_schemas.py

@@ -61,7 +61,7 @@ class SchemasMixin:
 
     It is currently implemented as a mixin rather than directly provided on the
     base class because it requires that the host `DatabaseClient` instance have a
-    `sqlalchemy` metadata object handle, and not all backends support this.
+    `sqlalchemy` engine object handle, and not all backends support this.
 
     If we are willing to forgo the ability to actually make queries using the
     SQLAlchemy ORM, we could instead use an SQL agnostic version.
@@ -79,9 +79,9 @@ def schemas(self):
         def get_schemas():
             if not getattr(self, "_schemas", None):
                 assert (
-                    getattr(self, "_sqlalchemy_metadata", None) is not None
-                ), f"`{self.__class__.__name__}` instances do not provide the required sqlalchemy metadata for schema exploration."
-                self._schemas = Schemas(self._sqlalchemy_metadata)
+                    getattr(self, "_sqlalchemy_engine", None) is not None
+                ), f"`{self.__class__.__name__}` instances do not provide the required sqlalchemy engine for schema exploration."
+                self._schemas = Schemas(self._sqlalchemy_engine)
             return self._schemas
 
         return Proxy(get_schemas)
@@ -138,22 +138,20 @@ class Schemas:
     An object which has as its attributes all of the schemas in a nominated database.
 
     Args:
-        metadata (sqlalchemy.MetaData): A SQL Alchemy `MetaData` instance
+        engine (sqlalchemy.Engine): A SQL Alchemy `Engine` instance
             configured for the nominated database.
     """
 
-    def __init__(self, metadata):
-        self._metadata = metadata
+    def __init__(self, engine):
+        self._engine = engine
         self._schema_names = None
         self._schema_cache = {}
 
     @property
     def all(self):
         "list<str>: The list of schema names."
         if self._schema_names is None:
-            self._schema_names = sqlalchemy.inspect(
-                self._metadata.bind
-            ).get_schema_names()
+            self._schema_names = sqlalchemy.inspect(self._engine).get_schema_names()
         return self._schema_names
 
     def __dir__(self):
@@ -162,9 +160,7 @@ def __dir__(self):
     def __getattr__(self, value):
         if value in self.all:
             if value not in self._schema_cache:
-                self._schema_cache[value] = Schema(
-                    metadata=self._metadata, schema=value
-                )
+                self._schema_cache[value] = Schema(engine=self._engine, schema=value)
             return self._schema_cache[value]
         raise AttributeError(f"No such schema {value}")
 
@@ -184,13 +180,13 @@ class Schema:
     An object which has as its attributes all of the tables in a nominated database schema.
 
     Args:
-        metadata (sqlalchemy.MetaData): A SQL Alchemy `MetaData` instance
+        engine (sqlalchemy.Engine): A SQL Alchemy `Engine` instance
             configured for the nominated database.
         schema (str): The schema within which to expose tables.
     """
 
-    def __init__(self, metadata, schema):
-        self._metadata = metadata
+    def __init__(self, engine, schema):
+        self._engine = engine
         self._schema = schema
         self._table_cache = {}
         self._table_names = None
@@ -199,7 +195,7 @@ def __init__(self, metadata, schema):
     def all(self):
         """list<str>: The table names in this database schema."""
         if self._table_names is None:
-            self._table_names = sqlalchemy.inspect(self._metadata.bind).get_table_names(
+            self._table_names = sqlalchemy.inspect(self._engine).get_table_names(
                 self._schema
             )
         return self._table_names

omniduct/databases/base.py

@@ -122,7 +122,6 @@ def __init__(
         self._templates = templates or {}
         self._template_context = template_context or {}
         self._sqlalchemy_engine = None
-        self._sqlalchemy_metadata = None
         self._default_format_opts = default_format_opts or {}
 
         self._init(**kwargs)

omniduct/databases/hiveserver2.py

@@ -117,7 +117,7 @@ def _init(
 
     @override
     def _connect(self):
-        from sqlalchemy import create_engine, MetaData
+        from sqlalchemy import create_engine
 
         if self.driver == "pyhive":
             try:
@@ -142,7 +142,6 @@ def _connect(self):
             self._sqlalchemy_engine = create_engine(
                 f"hive://{self.host}:{self.port}/{self.schema}"
             )
-            self._sqlalchemy_metadata = MetaData(self._sqlalchemy_engine)
         elif self.driver == "impyla":
             try:
                 import impala.dbapi
@@ -165,7 +164,6 @@ def _connect(self):
             self._sqlalchemy_engine = create_engine(
                 f"impala://{self.host}:{self.port}/{self.schema}"
             )
-            self._sqlalchemy_metadata = MetaData(self._sqlalchemy_engine)
 
     def __hive_cursor(self):
         if (
@@ -192,7 +190,6 @@ def _disconnect(self):
         # pylint: disable-next=attribute-defined-outside-init
         self.__hive = None
         self._sqlalchemy_engine = None
-        self._sqlalchemy_metadata = None
         # pylint: disable-next=attribute-defined-outside-init
         self._schemas = None
 

omniduct/databases/presto.py

@@ -90,14 +90,13 @@ def source(self, source):
 
     @override
     def _connect(self):
-        from sqlalchemy import create_engine, MetaData
+        from sqlalchemy import create_engine
 
         logging.getLogger("pyhive").setLevel(1000)  # Silence pyhive logging.
         logger.info("Connecting to Presto coordinator...")
         self._sqlalchemy_engine = create_engine(
             f"presto://{self.host}:{self.port}/{self.catalog}/{self.schema}"
         )
-        self._sqlalchemy_metadata = MetaData(self._sqlalchemy_engine)
 
     @override
     def _is_connected(self):
@@ -107,7 +106,6 @@ def _is_connected(self):
     def _disconnect(self):
         logger.info("Disconnecting from Presto coordinator...")
         self._sqlalchemy_engine = None
-        self._sqlalchemy_metadata = None
         self._schemas = None  # pylint: disable=attribute-defined-outside-init
 
     # Querying

omniduct/databases/sqlalchemy.py

@@ -70,6 +70,13 @@ def db_uri(self):
             database=self.database,
         )
 
+    @property
+    def _sqlalchemy_engine(self):
+        """
+        The SQLAlchemy engine object for the SchemasMixin.
+        """
+        return self.engine
+
     @override
     def _connect(self):
         import sqlalchemy
@@ -85,7 +92,6 @@ def _connect(self):
 
         # pylint: disable-next=attribute-defined-outside-init
         self.engine = sqlalchemy.create_engine(self.db_uri, **self.engine_opts)
-        self._sqlalchemy_metadata = sqlalchemy.MetaData(self.engine)
 
     @override
     def _is_connected(self):
@@ -95,7 +101,6 @@ def _is_connected(self):
     def _disconnect(self):
         # pylint: disable-next=attribute-defined-outside-init
         self.engine = None
-        self._sqlalchemy_metadata = None
         # pylint: disable-next=attribute-defined-outside-init
         self._schemas = None
 

pyproject.toml

@@ -7,7 +7,7 @@ name = "omniduct"
 dynamic = ["version"]
 description = "A toolkit providing a uniform interface for connecting to and extracting data from a wide variety of (potentially remote) data stores (including HDFS, Hive, Presto, MySQL, etc)."
 readme = "README.md"
-license = ""
+license = "MIT"
 authors = [
     { name = "Matthew Wardrop", email = "[email protected]" },
     { name = "Dan Frank", email = "[email protected]" },