Fix Broken Reflection for 1.4 FutureEngine (sqlalchemy-redshift#277)

* fix 269

* Update CHANGES.rst

Author: Brooke-white

CHANGES.rst

@@ -4,6 +4,8 @@
 - Override new upstream postgres method that fails against redshift (`Pull #266 <https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/pull/266>`_)
 - Fix table reflection broken for non-superusers
   (`Pull #276 <https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/pull/276>`_)
+- Fix Broken Reflection for 1.4 FutureEngine
+  (`Pull #277 <https://github.com/sqlalchemy-redshift/sqlalchemy-redshift/pull/277>`_)
 
 
 0.8.13 (2023-03-28)

sqlalchemy_redshift/dialect.py

@@ -753,7 +753,7 @@ def get_check_constraints(self, connection, table_name, schema=None, **kw):
         )
         table_oid = 'NULL' if not table_oid else table_oid
 
-        result = connection.execute("""
+        result = connection.execute(sa.text("""
                         SELECT
                             cons.conname as name,
                             pg_get_constraintdef(cons.oid) as src
@@ -762,7 +762,7 @@ def get_check_constraints(self, connection, table_name, schema=None, **kw):
                         WHERE
                             cons.conrelid = {} AND
                             cons.contype = 'c'
-                        """.format(table_oid))
+                        """.format(table_oid)))
         ret = []
         for name, src in result:
             # samples:
@@ -1122,15 +1122,14 @@ def _get_schema_column_info(self, connection, **kw):
         )
 
         all_columns = defaultdict(list)
-        with connection.connect() as cc:
-            result = cc.execute(sa.text(REFLECTION_SQL.format(
-                schema_clause=schema_clause,
-                table_clause=table_clause
-            )))
-
-            for col in result:
-                key = RelationKey(col.table_name, col.schema, connection)
-                all_columns[key].append(col)
+        result = connection.execute(sa.text(REFLECTION_SQL.format(
+            schema_clause=schema_clause,
+            table_clause=table_clause
+        )))
+
+        for col in result:
+            key = RelationKey(col.table_name, col.schema, connection)
+            all_columns[key].append(col)
 
         return dict(all_columns)
 

tests/conftest.py

@@ -176,14 +176,23 @@ def migrate(self, engine):
     @contextlib.contextmanager
     def _database(self):
         from sqlalchemy_redshift.dialect import \
-            RedshiftDialect_redshift_connector
+            RedshiftDialect_psycopg2cffi
 
         db_name = database_name()
-        with self.engine.connect() as conn:
-            conn.execute('COMMIT')  # Can't create databases in a transaction
-            if isinstance(conn.dialect, RedshiftDialect_redshift_connector):
-                conn.execution_options(isolation_level="AUTOCOMMIT")
-            conn.execute('CREATE DATABASE {db_name}'.format(db_name=db_name))
+        opts = (
+            {"isolation_level": "AUTOCOMMIT"}
+            if not isinstance(
+                self.engine.dialect, RedshiftDialect_psycopg2cffi
+            )
+            else {}
+        )
+
+        with self.engine.connect().execution_options(**opts) as conn:
+            if isinstance(self.engine.dialect, RedshiftDialect_psycopg2cffi):
+                conn.execute(sa.text("COMMIT"))
+            conn.execute(
+                sa.text('CREATE DATABASE {db_name}'.format(db_name=db_name))
+            )
 
         dburl = copy.deepcopy(self.engine.url)
         try:
@@ -197,13 +206,14 @@ def _database(self):
                 connect_args=self.engine_definition.connect_args,
             )
         finally:
-            with self.engine.connect() as conn:
-                conn.execute('COMMIT')  # Can't drop databases in a transaction
+            with self.engine.connect().execution_options(**opts) as conn:
                 if isinstance(
-                        conn.dialect, RedshiftDialect_redshift_connector
+                        self.engine.dialect, RedshiftDialect_psycopg2cffi
                 ):
-                    conn.execution_options(isolation_level="AUTOCOMMIT")
-                conn.execute('DROP DATABASE {db_name}'.format(db_name=db_name))
+                    conn.execute(sa.text("COMMIT"))
+                conn.execute(
+                    sa.text('DROP DATABASE {db_name}'.format(db_name=db_name))
+                )
 
     @contextlib.contextmanager
     def migrated_database(self):

tests/test_reflection.py

@@ -1,8 +1,8 @@
 import pytest
 from sqlalchemy import MetaData, Table, inspect
-from sqlalchemy.dialects.postgresql.psycopg2cffi import PGDialect_psycopg2cffi
 from sqlalchemy.schema import CreateTable
 from sqlalchemy.exc import NoSuchTableError
+import sqlalchemy as sa
 
 from rs_sqla_test_utils import models, utils
 
@@ -195,7 +195,7 @@ def test_no_table_reflection(redshift_session):
 def test_no_search_path_leak(redshift_session):
     metadata = MetaData(bind=redshift_session.bind)
     Table('basic', metadata, autoload=True)
-    result = redshift_session.execute("SHOW search_path")
+    result = redshift_session.execute(sa.text("SHOW search_path"))
     search_path = result.scalar()
     assert 'other_schema' not in search_path
 
@@ -207,12 +207,12 @@ def test_external_table_reflection(redshift_engine, iam_role_arn):
                     iam_role '{iam_role_arn}'
                     create external database if not exists;
                     """
-
-    conn = redshift_engine.connect()
-    conn.execute(schema_ddl)
-    insp = inspect(redshift_engine)
-    all_schemas = insp.get_schema_names()
-    assert 'bananas' in all_schemas
+    with redshift_engine.connect() as conn:
+        conn.execute(sa.text(schema_ddl))
+        conn.execute(sa.text('COMMIT'))
+        insp = inspect(redshift_engine)
+        all_schemas = insp.get_schema_names()
+        assert 'bananas' in all_schemas
 
     table_ddl = """create external table bananas.sales(
         salesid integer,
@@ -226,13 +226,19 @@ def test_external_table_reflection(redshift_engine, iam_role_arn):
         location 's3://awssampledbuswest2/tickit/spectrum/sales/'
         table properties ('numRows'='172000');
     """
-    with redshift_engine.connect() as conn:
-        # Redshift can't run CREATE EXTERNAL TABLE inside a transaction
-        # e.g. (BEGIN … END)
-        if not isinstance(conn.dialect, PGDialect_psycopg2cffi):
-            conn.execution_options(isolation_level="AUTOCOMMIT")
+    from sqlalchemy_redshift.dialect import \
+        RedshiftDialect_psycopg2cffi
+    opts = (
+        {"isolation_level": "AUTOCOMMIT"}
+        if not isinstance(
+            redshift_engine.dialect, RedshiftDialect_psycopg2cffi
+        ) else {}
+    )
 
-        conn.execute(table_ddl)
+    with redshift_engine.connect().execution_options(**opts) as conn:
+        conn.execute(sa.text(table_ddl))
+        if isinstance(redshift_engine.dialect, RedshiftDialect_psycopg2cffi):
+            conn.execute(sa.text("COMMIT"))
 
         insp = inspect(redshift_engine)
         table_columns_definition = insp.get_columns(
@@ -245,8 +251,12 @@ def test_external_table_reflection(redshift_engine, iam_role_arn):
         assert 'pricepaid' in table_columns
 
         # Drop external table because we are using `AUTOCOMMIT`
-        conn.execute("DROP TABLE IF EXISTS bananas.sales")
+        conn.execute(sa.text("DROP TABLE IF EXISTS bananas.sales"))
 
         # Also drop the external db:
         # https://docs.aws.amazon.com/redshift/latest/dg/r_DROP_DATABASE.html
-        conn.execute("DROP SCHEMA IF EXISTS bananas DROP EXTERNAL DATABASE")
+        conn.execute(
+            sa.text("DROP SCHEMA IF EXISTS bananas DROP EXTERNAL DATABASE")
+        )
+        if isinstance(redshift_engine.dialect, RedshiftDialect_psycopg2cffi):
+            conn.execute(sa.text("COMMIT"))

tests/test_reflection_views.py

@@ -1,6 +1,6 @@
 from sqlalchemy import MetaData, Table, inspect
 from sqlalchemy.schema import CreateTable
-
+import sqlalchemy as sa
 from rs_sqla_test_utils.utils import clean, compile_query
 
 
@@ -13,42 +13,48 @@ def test_view_reflection(redshift_engine):
     table_ddl = "CREATE TABLE my_table (col1 INTEGER, col2 INTEGER)"
     view_query = "SELECT my_table.col1, my_table.col2 FROM my_table"
     view_ddl = "CREATE VIEW my_view AS %s" % view_query
-    conn = redshift_engine.connect()
-    try:
-        conn.execute(table_ddl)
-        conn.execute(view_ddl)
-        insp = inspect(redshift_engine)
-        view_definition = insp.get_view_definition('my_view')
-        assert clean(
-            compile_query(view_definition, redshift_engine.dialect)
-        ) == clean(view_query)
-        view = Table('my_view', MetaData(),
-                     autoload=True, autoload_with=redshift_engine)
-        assert(len(view.columns) == 2)
-    finally:
-        conn.execute('DROP TABLE IF EXISTS my_table CASCADE')
-        conn.execute('DROP VIEW IF EXISTS my_view CASCADE')
+
+    with redshift_engine.connect() as conn:
+        try:
+            conn.execute(sa.text(table_ddl))
+            conn.execute(sa.text(view_ddl))
+            conn.execute(sa.text("COMMIT"))
+            insp = inspect(redshift_engine)
+            view_definition = insp.get_view_definition('my_view')
+            assert clean(
+                compile_query(view_definition, redshift_engine.dialect)
+            ) == clean(view_query)
+            view = Table('my_view', MetaData(),
+                         autoload=True, autoload_with=redshift_engine)
+            assert(len(view.columns) == 2)
+        finally:
+            conn.execute(sa.text('DROP TABLE IF EXISTS my_table CASCADE'))
+            conn.execute(sa.text('DROP VIEW IF EXISTS my_view CASCADE'))
+            conn.execute(sa.text("COMMIT"))
 
 
 def test_late_binding_view_reflection(redshift_engine):
     table_ddl = "CREATE TABLE my_table (col1 INTEGER, col2 INTEGER)"
     view_query = "SELECT my_table.col1, my_table.col2 FROM public.my_table"
     view_ddl = ("CREATE VIEW my_late_view AS "
                 "%s WITH NO SCHEMA BINDING" % view_query)
-    conn = redshift_engine.connect()
-    try:
-        conn.execute(table_ddl)
-        conn.execute(view_ddl)
-        insp = inspect(redshift_engine)
-        view_definition = insp.get_view_definition('my_late_view')
-
-        # Redshift returns the entire DDL for late binding views.
-        assert clean(
-            compile_query(view_definition, redshift_engine.dialect)
-        ) == clean(view_ddl)
-        view = Table('my_late_view', MetaData(),
-                     autoload=True, autoload_with=redshift_engine)
-        assert(len(view.columns) == 2)
-    finally:
-        conn.execute('DROP TABLE IF EXISTS my_table CASCADE')
-        conn.execute('DROP VIEW IF EXISTS my_late_view CASCADE')
+
+    with redshift_engine.connect() as conn:
+        try:
+            conn.execute(sa.text(table_ddl))
+            conn.execute(sa.text(view_ddl))
+            conn.execute(sa.text("COMMIT"))
+            insp = inspect(redshift_engine)
+            view_definition = insp.get_view_definition('my_late_view')
+
+            # Redshift returns the entire DDL for late binding views.
+            assert clean(
+                compile_query(view_definition, redshift_engine.dialect)
+            ) == clean(view_ddl)
+            view = Table('my_late_view', MetaData(),
+                         autoload=True, autoload_with=redshift_engine)
+            assert(len(view.columns) == 2)
+        finally:
+            conn.execute(sa.text('DROP TABLE IF EXISTS my_table CASCADE'))
+            conn.execute(sa.text('DROP VIEW IF EXISTS my_late_view CASCADE'))
+            conn.execute(sa.text('COMMIT'))