Upgrade to SQLAlchemy 2

Signed-off-by: Aurélien Bompard <[email protected]>

datanommer.commands/datanommer/commands/__init__.py

@@ -16,11 +16,11 @@
 import json
 import logging
 import time
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 
 import click
 from fedora_messaging import config as fedora_messaging_config
-from sqlalchemy import func
+from sqlalchemy import func, select
 
 import datanommer.models as m
 
@@ -157,23 +157,17 @@ def stats(config_path, topic, category):
     )
 
     if topic:
+        query = select(m.Message.topic, func.count(m.Message.topic))
         if category:
-            query = m.session.query(
-                m.Message.topic, func.count(m.Message.topic)
-            ).filter(m.Message.category == category)
-        else:
-            query = m.session.query(m.Message.topic, func.count(m.Message.topic))
+            query = query.where(m.Message.category == category)
         query = query.group_by(m.Message.topic)
     else:
+        query = select(m.Message.category, func.count(m.Message.category))
         if category:
-            query = m.session.query(
-                m.Message.category, func.count(m.Message.category)
-            ).filter(m.Message.category == category)
-        else:
-            query = m.session.query(m.Message.category, func.count(m.Message.category))
+            query = query.where(m.Message.category == category)
         query = query.group_by(m.Message.category)
 
-    results = query.all()
+    results = m.session.execute(query).all()
 
     if topic:
         for topic, count in results:
@@ -295,30 +289,28 @@ def latest(config_path, topic, category, overall, timestamp, timesince, human):
     )
 
     if topic:
-        queries = [m.Message.query.filter(m.Message.topic == topic)]
+        queries = [select(m.Message).where(m.Message.topic == topic)]
 
     elif category:
-        queries = [m.Message.query.filter(m.Message.category == category)]
+        queries = [select(m.Message).where(m.Message.category == category)]
     elif not overall:
         # If no args..
-        categories = [
-            c[0]
-            for c in m.session.query(m.Message.category)
-            .distinct()
-            .order_by(m.Message.category)
-        ]
+        categories_query = (
+            select(m.Message.category).distinct().order_by(m.Message.category)
+        )
+        categories = m.session.execute(categories_query).scalars()
         queries = [
-            m.Message.query.filter(m.Message.category == category)
+            select(m.Message).where(m.Message.category == category)
             for category in categories
         ]
     else:
         # Show only the single latest message, regardless of type.
-        queries = [m.Message.query]
+        queries = [select(m.Message)]
 
     # Only check messages from the last year to speed up queries
     a_year = timedelta(days=365)
-    earliest = datetime.utcnow() - a_year
-    queries = [q.filter(m.Message.timestamp > earliest) for q in queries]
+    earliest = datetime.now(tz=timezone.utc) - a_year
+    queries = [q.where(m.Message.timestamp > earliest) for q in queries]
 
     # Order and limit to the latest.
     queries = [q.order_by(m.Message.timestamp.desc()).limit(1) for q in queries]
@@ -337,7 +329,9 @@ def formatter(key, val):
             return f'{{"{key}": {json.dumps(val.as_fedora_message_dict())}}}'
 
     results = []
-    for result in sum((query.all() for query in queries), []):
+    for result in sum(
+        (list(m.session.execute(query).scalars()) for query in queries), []
+    ):
         results.append(formatter(result.category, result))
 
     click.echo(f"[{','.join(results)}]")

datanommer.commands/tests/test_commands.py

@@ -64,6 +64,7 @@ def test_create(mocker):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.create, [])
+    assert result.exit_code == 0, result.output
 
     assert result.output == "Creating Datanommer database and tables\n"
     mock_model_init.assert_called_once_with(
@@ -91,6 +92,7 @@ def test_stats(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.stats, [])
+    assert result.exit_code == 0, result.output
 
     assert "git has 2 entries" in result.output
     assert "fas has 1 entries" in result.output
@@ -116,6 +118,7 @@ def test_stats_topics(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.stats, ["--topic"])
+    assert result.exit_code == 0, result.output
 
     assert (
         "org.fedoraproject.prod.git.receive.valgrind.master has 1 entries"
@@ -148,6 +151,7 @@ def test_stats_category_topics(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.stats, ["--topic", "--category", "git"])
+    assert result.exit_code == 0, result.output
 
     assert (
         "org.fedoraproject.prod.git.receive.valgrind.master has 1 entries"
@@ -180,6 +184,7 @@ def test_stats_category(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.stats, ["--category", "git"])
+    assert result.exit_code == 0, result.output
 
     assert result.output == "git has 2 entries\n"
 
@@ -196,6 +201,7 @@ def test_dump(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.dump, [])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -219,6 +225,7 @@ def test_dump_before(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.dump, ["--before", "2013-02-16"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -246,6 +253,7 @@ def test_dump_since(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.dump, ["--since", "2013-02-14T08:00:00"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -276,6 +284,7 @@ def test_dump_timespan(datanommer_models, mock_config, mock_init):
         datanommer.commands.dump,
         ["--before", "2013-02-16", "--since", "2013-02-14T08:00:00"],
     )
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -288,9 +297,11 @@ def test_dump_timespan(datanommer_models, mock_config, mock_init):
 def test_dump_invalid_dates(datanommer_models, mock_config, mock_init):
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.dump, ["--before", "2013-02-16asdasd"])
+    assert result.exit_code > 0, result.output
     assert result.output == "Error: Invalid date format\n"
 
     result = runner.invoke(datanommer.commands.dump, ["--since", "2013-02-16asdasd"])
+    assert result.exit_code > 0, result.output
     assert result.output == "Error: Invalid date format\n"
 
 
@@ -314,6 +325,7 @@ def test_latest_overall(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--overall"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -343,6 +355,7 @@ def test_latest_topic(datanommer_models, mock_config, mock_init):
     result = runner.invoke(
         datanommer.commands.latest, ["--topic", "org.fedoraproject.stg.fas.user.create"]
     )
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -370,6 +383,7 @@ def test_latest_category(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--category", "fas"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -402,6 +416,7 @@ def test_latest_timestamp_human(datanommer_models, mocker, mock_config, mock_ini
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--timestamp", "--human"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -431,6 +446,7 @@ def test_latest_timestamp(datanommer_models, mocker, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--timestamp"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -465,6 +481,7 @@ def test_latest_timesince(datanommer_models, mocker, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--timesince"])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 
@@ -501,6 +518,7 @@ def test_latest_timesince_human(datanommer_models, mock_config, mock_init, mocke
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, ["--timesince", "--human"])
+    assert result.exit_code == 0, result.output
 
     assert json.loads(result.output) == ["1 day, 0:00:00", "0:00:01"]
 
@@ -527,6 +545,7 @@ def test_latest(datanommer_models, mock_config, mock_init):
 
     runner = CliRunner()
     result = runner.invoke(datanommer.commands.latest, [])
+    assert result.exit_code == 0, result.output
 
     json_object = json.loads(result.output)
 

datanommer.commands/tests/test_extract_users.py

@@ -80,7 +80,6 @@ def test_extract_users_topic_and_category(mock_config, mock_init):
     result = runner.invoke(
         extract_users, ["--category", "bodhi", "--topic", "some.topic"]
     )
-    print(result.output)
     assert result.exit_code != 0, result.output
     assert "Error: can't use both --topic and --category, choose one." in result.output
 

datanommer.commands/tox.ini

@@ -12,6 +12,8 @@ skip_install = true
 #     poetry
 deps =
     poetry>=1.2
+env =
+    SQLALCHEMY_WARN_20=1
 commands_pre =
     poetry install --all-extras
     poetry run {toxinidir}/../tools/install-models-as-editable.sh

datanommer.consumer/tox.ini

@@ -12,6 +12,8 @@ skip_install = true
 #     poetry
 deps =
     poetry>=1.2
+env =
+    SQLALCHEMY_WARN_20=1
 commands_pre =
     poetry install --all-extras
     poetry run {toxinidir}/../tools/install-models-as-editable.sh

datanommer.models/datanommer/models/__init__.py

@@ -81,7 +81,7 @@ def init(uri=None, alembic_ini=None, engine=None, create=False):
         raise ValueError("One of uri or engine must be specified")
 
     if uri and not engine:
-        engine = create_engine(uri)
+        engine = create_engine(uri, future=True)
 
     # We need to hang our own attribute on the sqlalchemy session to stop
     # ourselves from initializing twice.  That is only a problem if the code
@@ -95,7 +95,8 @@ def init(uri=None, alembic_ini=None, engine=None, create=False):
     DeclarativeBase.query = session.query_property()
 
     if create:
-        session.execute(text("CREATE EXTENSION IF NOT EXISTS timescaledb"))
+        with engine.begin() as connection:
+            connection.execute(text("CREATE EXTENSION IF NOT EXISTS timescaledb"))
         DeclarativeBase.metadata.create_all(engine)
         # Loads the alembic configuration and generates the version table, with
         # the most recent revision stamped as head
@@ -509,7 +510,7 @@ def get_or_create(cls, name):
         if name in cls._cache:
             # If we cache the instance, SQLAlchemy will run this query anyway because the instance
             # will be from a different transaction. So just cache the id.
-            return cls.query.get(cls._cache[name])
+            return session.get(cls, cls._cache[name])
         obj = cls.query.filter_by(name=name).one_or_none()
         if obj is None:
             obj = cls(name=name)

datanommer.models/datanommer/models/testing/__init__.py

@@ -32,16 +32,14 @@ def datanommer_db(postgresql_proc, datanommer_db_url):
         # template_dbname=postgresql_proc.template_dbname,
         version=postgresql_proc.version,
     ):
-        engine = sa.create_engine(datanommer_db_url, poolclass=sa.pool.NullPool)
-        # dm.init will also run this, but somehow it fails then at creating tables
-        # as if timescaledb was not loaded.
-        with engine.connect() as connection:
-            connection.execute(sa.text("CREATE EXTENSION IF NOT EXISTS timescaledb"))
+        engine = sa.create_engine(
+            datanommer_db_url, future=True, poolclass=sa.pool.NullPool
+        )
         # Renew the global object, dm.init checks a custom attribute
         dm.session = scoped_session(dm.maker)
         dm.init(engine=engine, create=True)
         yield engine
-        sa.orm.close_all_sessions()
+        dm.session.close()
 
 
 @pytest.fixture()