unified undre postgres_tools

Author: pcrespov

packages/postgres-database/tests/conftest.py

@@ -18,6 +18,7 @@
 from aiopg.sa.engine import Engine
 from aiopg.sa.result import ResultProxy, RowProxy
 from faker import Faker
+from pytest_simcore.helpers import postgres_tools
 from pytest_simcore.helpers.faker_factories import (
     random_group,
     random_project,
@@ -165,10 +166,7 @@ def pg_sa_engine(
 
     yield sync_engine
 
-    # NOTE: ALL is deleted after
-    with sync_engine.begin() as conn:
-        conn.execute(sa.DDL("DROP TABLE IF EXISTS alembic_version"))
-    db_metadata.drop_all(sync_engine)
+    postgres_tools.drop_all_tables(sync_engine)
     sync_engine.dispose()
 
 

packages/postgres-database/tests/test_models_groups.py

@@ -1,6 +1,5 @@
 # pylint: disable=redefined-outer-name
-# pylint: disable=unused-argument
-# pylint: disable=unused-variable
+# pylint: disable=unused-argument# pylint: disable=unused-variable
 # pylint: disable=too-many-arguments
 
 
@@ -13,6 +12,7 @@
 from aiopg.sa.engine import Engine
 from aiopg.sa.result import ResultProxy, RowProxy
 from psycopg2.errors import ForeignKeyViolation, RaiseException, UniqueViolation
+from pytest_simcore.helpers import postgres_tools
 from pytest_simcore.helpers.faker_factories import random_user
 from simcore_postgres_database.models.base import metadata
 from simcore_postgres_database.webserver_models import (
@@ -36,7 +36,7 @@ async def connection(
     async with engine.acquire() as conn:
         yield conn
 
-    metadata.drop_all(sync_engine)
+    postgres_tools.drop_all_tables(sync_engine)
 
 
 async def test_user_group_uniqueness(

packages/postgres-database/tests/test_uniqueness_in_comp_tasks.py

@@ -4,10 +4,14 @@
 # pylint:disable=redefined-outer-name
 
 import json
+from collections.abc import Awaitable, Callable
 
+import aiopg.sa.engine
+import aiopg.sa.exc
 import pytest
 import sqlalchemy as sa
 from psycopg2.errors import UniqueViolation  # pylint: disable=no-name-in-module
+from pytest_simcore.helpers import postgres_tools
 from pytest_simcore.helpers.faker_factories import fake_pipeline, fake_task_factory
 from simcore_postgres_database.models.base import metadata
 from simcore_postgres_database.webserver_models import comp_pipeline, comp_tasks
@@ -16,9 +20,11 @@
 
 
 @pytest.fixture
-async def engine(make_engine):
-    engine = await make_engine()
-    sync_engine = make_engine(is_async=False)
+async def engine(
+    make_engine: Callable[[bool], Awaitable[aiopg.sa.engine.Engine] | sa.engine.Engine]
+):
+    engine: aiopg.sa.engine.Engine = await make_engine()
+    sync_engine: sa.engine.Engine = make_engine(is_async=False)
     metadata.drop_all(sync_engine)
     metadata.create_all(sync_engine)
 
@@ -32,6 +38,7 @@ async def engine(make_engine):
 
     yield engine
 
+    postgres_tools.drop_all_tables(sync_engine)
     engine.close()
     await engine.wait_closed()
 

packages/pytest-simcore/src/pytest_simcore/helpers/postgres_tools.py

@@ -17,6 +17,22 @@ class PostgresTestConfig(TypedDict):
     port: str
 
 
+def drop_all_tables(sa_sync_engine: sa.engine.Engine):
+    with sa_sync_engine.begin() as conn:
+        conn.execute(sa.DDL("DROP TABLE IF EXISTS alembic_version"))
+        conn.execute(
+            # NOTE: terminates all open transactions before droping all tables
+            # This solves https://github.com/ITISFoundation/osparc-simcore/issues/7008
+            sa.DDL(
+                "SELECT pg_terminate_backend(pid) "
+                "FROM pg_stat_activity "
+                "WHERE state = 'idle in transaction';"
+            )
+        )
+        # SEE https://github.com/ITISFoundation/osparc-simcore/issues/1776
+        metadata.drop_all(bind=conn)
+
+
 @contextmanager
 def migrated_pg_tables_context(
     postgres_config: PostgresTestConfig,
@@ -50,18 +66,8 @@ def migrated_pg_tables_context(
     simcore_postgres_database.cli.clean.callback()  # just cleans discover cache
 
     postgres_engine = sa.create_engine(dsn)
-    with postgres_engine.begin() as conn:
-        conn.execute(
-            # NOTE: terminates all open transactions before droping all tables
-            # This solves https://github.com/ITISFoundation/osparc-simcore/issues/7008
-            sa.DDL(
-                "SELECT pg_terminate_backend(pid) "
-                "FROM pg_stat_activity "
-                "WHERE state = 'idle in transaction';"
-            )
-        )
-        # SEE https://github.com/ITISFoundation/osparc-simcore/issues/1776
-        metadata.drop_all(bind=postgres_engine)
+    drop_all_tables(postgres_engine)
+    postgres_engine.dispose()
 
 
 def is_postgres_responsive(url) -> bool:

packages/service-library/tests/aiohttp/with_postgres/test_aiopg_utils.py

@@ -13,6 +13,7 @@
 import pytest
 import sqlalchemy as sa
 import sqlalchemy.exc as sa_exceptions
+from pytest_simcore.helpers import postgres_tools
 from servicelib.aiohttp.aiopg_utils import init_pg_tables, is_pg_responsive
 from servicelib.common_aiopg_utils import DataSourceName, create_pg_engine
 
@@ -64,7 +65,7 @@ def test_dsn_uri_with_query(postgres_service_with_fake_data: DataSourceName):
 
         # if url is wrong, these will fail
         metadata.create_all(sa_engine)
-        metadata.drop_all(sa_engine)
+        postgres_tools.drop_all_tables(sa_engine)
 
     except sa_exceptions.SQLAlchemyError as ee:
         pytest.fail(f"Cannot connect with {uri}: {ee}")

services/api-server/tests/unit/_with_db/conftest.py

@@ -24,6 +24,7 @@
 from models_library.api_schemas_api_server.api_keys import ApiKeyInDB
 from pydantic import PositiveInt
 from pytest_mock import MockerFixture
+from pytest_simcore.helpers import postgres_tools
 from pytest_simcore.helpers.faker_factories import (
     random_api_key,
     random_product,
@@ -32,7 +33,6 @@
 from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
 from pytest_simcore.helpers.typing_env import EnvVarsDict
 from simcore_postgres_database.models.api_keys import api_keys
-from simcore_postgres_database.models.base import metadata
 from simcore_postgres_database.models.products import products
 from simcore_postgres_database.models.users import users
 from simcore_service_api_server.core.application import init_app
@@ -141,8 +141,8 @@ def migrated_db(postgres_service: dict, make_engine: Callable):
     pg_cli.downgrade.callback("base")
     pg_cli.clean.callback()
     # FIXME: deletes all because downgrade is not reliable!
-    engine = make_engine(is_async=False)
-    metadata.drop_all(engine)
+    sync_engine = make_engine(is_async=False)
+    postgres_tools.drop_all_tables(sync_engine)
 
 
 @pytest.fixture

services/web/server/tests/unit/with_dbs/conftest.py

@@ -44,6 +44,7 @@
 from pydantic import ByteSize, TypeAdapter
 from pytest_docker.plugin import Services
 from pytest_mock import MockerFixture
+from pytest_simcore.helpers import postgres_tools
 from pytest_simcore.helpers.faker_factories import random_product
 from pytest_simcore.helpers.monkeypatch_envs import setenvs_from_dict
 from pytest_simcore.helpers.typing_env import EnvVarsDict
@@ -524,19 +525,7 @@ def postgres_db(
     # NOTE: we directly drop the table, that is faster
     # testing the upgrade/downgrade is already done in postgres-database.
     # there is no need to it here.
-    with engine.begin() as conn:
-        conn.execute(sa.DDL("DROP TABLE IF EXISTS alembic_version"))
-
-        conn.execute(
-            # NOTE: terminates all open transactions before droping all tables
-            # This solves https://github.com/ITISFoundation/osparc-simcore/issues/7008
-            sa.DDL(
-                "SELECT pg_terminate_backend(pid) "
-                "FROM pg_stat_activity "
-                "WHERE state = 'idle in transaction';"
-            )
-        )
-        orm.metadata.drop_all(bind=conn)
+    postgres_tools.drop_all_tables(engine)
 
     engine.dispose()