[Chore] Remove SQLA 2.0 deprecations (#76)

Author: tchapi

conftest.py

@@ -3,9 +3,8 @@
 
 import pytest
 import sqlalchemy as sa
-from sqlalchemy import create_engine
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy import create_engine, text
+from sqlalchemy.orm import close_all_sessions, declarative_base, sessionmaker
 
 from postgresql_audit import VersioningManager
 
@@ -45,10 +44,12 @@ def engine(dns):
 
 @pytest.fixture
 def connection(engine):
-    conn = engine.connect()
-    conn.execute('CREATE EXTENSION IF NOT EXISTS btree_gist')
-    yield conn
-    conn.close()
+    with engine.connect() as conn:
+        with conn.begin():
+            conn.execute(text('CREATE EXTENSION IF NOT EXISTS btree_gist'))
+        
+        yield conn
+        conn.close()
 
 
 @pytest.fixture
@@ -57,7 +58,7 @@ def session(connection):
     session = Session()
     yield session
     session.expunge_all()
-    session.close_all()
+    close_all_sessions()
 
 
 @pytest.fixture

postgresql_audit/base.py

@@ -5,7 +5,7 @@
 from weakref import WeakSet
 
 import sqlalchemy as sa
-from sqlalchemy import orm
+from sqlalchemy import orm, text
 from sqlalchemy.dialects.postgresql import (
     array,
     ExcludeConstraint,
@@ -205,13 +205,17 @@ def get_transaction_values(self):
     @contextmanager
     def disable(self, session):
         session.execute(
-            "SET LOCAL postgresql_audit.enable_versioning = 'false'"
+            text(
+                "SET LOCAL postgresql_audit.enable_versioning = 'false'"
+            )
         )
         try:
             yield
         finally:
             session.execute(
-                "SET LOCAL postgresql_audit.enable_versioning = 'true'"
+                text(
+                    "SET LOCAL postgresql_audit.enable_versioning = 'true'"
+                )
             )
 
     def render_tmpl(self, tmpl_name):
@@ -299,7 +303,7 @@ def audit_table(self, table, exclude_columns=None):
             func = sa.func.audit_table
         else:
             func = getattr(getattr(sa.func, self.schema_name), 'audit_table')
-        query = sa.select([func(*args)])
+        query = sa.select(func(*args))
         if query not in cached_statements:
             cached_statements[query] = StatementExecutor(query)
         listener = (table, 'after_create', cached_statements[query])

postgresql_audit/expressions.py

@@ -15,7 +15,7 @@ class jsonb_change_key_name(expression.FunctionElement):
 
 
         data = {'key1': 1, 'key3': 4}
-        query = sa.select([jsonb_merge(data, 'key1', 'key2')])
+        query = sa.select(jsonb_merge(data, 'key1', 'key2'))
         session.execute(query).scalar()  # {'key2': 1, 'key3': 4}
     """
     type = JSONB()

tests/test_custom_schema.py

@@ -37,12 +37,16 @@ def table_creator(
     schema_name
 ):
     sa.orm.configure_mappers()
-    connection.execute('DROP SCHEMA IF EXISTS {} CASCADE'.format(schema_name))
-    tx = connection.begin()
-    versioning_manager.transaction_cls.__table__.create(connection)
-    versioning_manager.activity_cls.__table__.create(connection)
-    base.metadata.create_all(connection)
-    tx.commit()
+    with connection.begin():
+        connection.execute(
+            sa.text(
+                'DROP SCHEMA IF EXISTS {} CASCADE'.format(schema_name)
+            )
+        )
+        versioning_manager.transaction_cls.__table__.create(connection)
+        versioning_manager.activity_cls.__table__.create(connection)
+        base.metadata.create_all(connection)
+
     yield
     session.expunge_all()
     base.metadata.drop_all(connection)

tests/test_json_functions.py

@@ -31,14 +31,16 @@ class TestJSONBChangeKeyName(object):
     )
     def test_raw_sql(self, session, data, old_key, new_key, expected):
         result = session.execute(
-            '''SELECT jsonb_change_key_name(
-                '{data}'::jsonb,
-                '{old_key}',
-                '{new_key}'
-            )'''.format(
-                data=data,
-                old_key=old_key,
-                new_key=new_key
+            sa.text(
+                '''SELECT jsonb_change_key_name(
+                    '{data}'::jsonb,
+                    '{old_key}',
+                    '{new_key}'
+                )'''.format(
+                    data=data,
+                    old_key=old_key,
+                    new_key=new_key
+                )
             )
         ).scalar()
         assert result == expected
@@ -75,7 +77,7 @@ def test_sqlalchemy_function_expr(
         expected
     ):
         result = session.execute(
-            sa.select([jsonb_change_key_name(data, old_key, new_key)])
+            sa.select(jsonb_change_key_name(data, old_key, new_key))
         ).scalar()
         assert result == expected
 
@@ -104,9 +106,11 @@ class TestJSONBConcatOperator(object):
     )
     def test_raw_sql(self, session, data, merge_data, expected):
         result = session.execute(
-            '''SELECT '{data}'::jsonb || '{merge_data}'::jsonb'''.format(
-                data=data,
-                merge_data=merge_data
+            sa.text(
+                '''SELECT '{data}'::jsonb || '{merge_data}'::jsonb'''.format(
+                    data=data,
+                    merge_data=merge_data
+                )
             )
         ).scalar()
         assert result == expected

tests/test_operators.py

@@ -1,6 +1,8 @@
 import json
 import pytest
 
+from sqlalchemy import text
+
 
 @pytest.mark.parametrize(
     ('old', 'new', 'result'),
@@ -23,7 +25,9 @@
 )
 def test_jsonb_subtract(table_creator, session, old, new, result):
     assert session.execute(
-        'SELECT (:new)::jsonb - (:old)::jsonb',
+        text(
+            'SELECT (:new)::jsonb - (:old)::jsonb'
+        ),
         dict(old=json.dumps(old), new=json.dumps(new))
     ).scalar() == result
 
@@ -47,6 +51,8 @@ def test_jsonb_subtract(table_creator, session, old, new, result):
 )
 def test_jsonb_subtract_text_array(table_creator, session, old, new, result):
     assert session.execute(
-        'SELECT (:new)::jsonb - (:old)::text[]',
+        text(
+            'SELECT (:new)::jsonb - (:old)::text[]'
+        ),
         dict(old=old, new=json.dumps(new))
     ).scalar() == result

tests/test_sql_files.py

@@ -1,6 +1,8 @@
 # -*- coding: utf-8 -*-
 import pytest
 
+from sqlalchemy import text
+
 from .utils import last_activity
 
 
@@ -9,8 +11,10 @@ class TestActivityCreationWithColumnExclusion(object):
     @pytest.fixture
     def audit_trigger_creator(self, session, user_class):
         session.execute(
-            '''SELECT audit_table('{0}', '{{"age"}}')'''.format(
-                user_class.__tablename__
+            text(
+                '''SELECT audit_table('{0}', '{{"age"}}')'''.format(
+                    user_class.__tablename__
+                )
             )
         )
 

tests/utils.py

@@ -1,11 +1,16 @@
+from sqlalchemy import text
+
+
 def last_activity(connection, schema=None):
     if schema is not None:
         schema_prefix = '{}.'.format(schema)
     else:
         schema_prefix = ''
     return dict(
         connection.execute(
-            'SELECT * FROM {}activity ORDER BY issued_at '
-            'DESC LIMIT 1'.format(schema_prefix)
+            text(
+                'SELECT * FROM {}activity ORDER BY issued_at '
+                'DESC LIMIT 1'.format(schema_prefix)
+            )
         ).fetchone()
     )

tox.ini

@@ -1,4 +1,4 @@
-# Tox (http://tox.testrun.org/) is a tool for running tests
+# Tox (https://tox.wiki/) is a tool for running tests
 # in multiple virtualenvs. This configuration file will run the
 # test suite on all supported python versions. To use it, "pip install tox"
 # and then run "tox" from this directory.