Upgrade to psycopg v3

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 3.0.0
+current_version = 3.1.0rc1
 commit = False
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(rc(?P<build>\d+))?

.github/workflows/run-codspeed-tests.yml

@@ -61,7 +61,7 @@ jobs:
 
       - uses: CodSpeedHQ/action@v3
         with:
-          run: CACHE_URI=redis://redis DATABASE_URI=postgresql://$POSTGRES_USER:$POSTGRES_PASSWORD@$POSTGRES_HOST/$POSTGRES_DB pytest test/unit_tests --codspeed
+          run: CACHE_URI=redis://redis DATABASE_URI=postgresql+psycopg://$POSTGRES_USER:$POSTGRES_PASSWORD@$POSTGRES_HOST/$POSTGRES_DB pytest test/unit_tests --codspeed
           token: ${{ secrets.CODSPEED_TOKEN }}
         env:
           POSTGRES_DB: orchestrator-core-test

.github/workflows/run-unit-tests.yml

@@ -51,7 +51,7 @@ jobs:
         env:
           FLIT_ROOT_INSTALL: 1
       - name: Run Unit tests
-        run: CACHE_URI=redis://redis DATABASE_URI=postgresql://$POSTGRES_USER:$POSTGRES_PASSWORD@$POSTGRES_HOST/$POSTGRES_DB  pytest --cov-branch --cov=orchestrator --cov-report=xml --ignore=test --ignore=orchestrator/devtools --ignore=examples --ignore=docs --ignore=orchestrator/vendor
+        run: CACHE_URI=redis://redis DATABASE_URI=postgresql+psycopg://$POSTGRES_USER:$POSTGRES_PASSWORD@$POSTGRES_HOST/$POSTGRES_DB  pytest --cov-branch --cov=orchestrator --cov-report=xml --ignore=test --ignore=orchestrator/devtools --ignore=examples --ignore=docs --ignore=orchestrator/vendor
         env:
           POSTGRES_DB: orchestrator-core-test
           POSTGRES_USER: nwa

orchestrator/__init__.py

@@ -13,7 +13,7 @@
 
 """This is the orchestrator workflow engine."""
 
-__version__ = "3.0.0"
+__version__ = "3.1.0rc1"
 
 from orchestrator.app import OrchestratorCore
 from orchestrator.settings import app_settings

orchestrator/cli/generator/generator/migration.py

@@ -31,13 +31,16 @@
     sort_product_blocks_by_dependencies,
 )
 from orchestrator.cli.generator.generator.settings import product_generator_settings as settings
+from orchestrator.settings import convert_database_uri
 
 logger = structlog.getLogger(__name__)
 
 
 def create_migration_file(message: str, head: str) -> Path | None:
-    if not environ.get("DATABASE_URI"):
-        environ.update({"DATABASE_URI": "postgresql://nwa:nwa@localhost/orchestrator-core"})
+    if environ.get("DATABASE_URI"):
+        environ.update({"DATABASE_URI": convert_database_uri(environ["DATABASE_URI"])})
+    else:
+        environ.update({"DATABASE_URI": "postgresql+psycopg://nwa:nwa@localhost/orchestrator-core"})
     if not environ.get("PYTHONPATH"):
         environ.update({"PYTHONPATH": "."})
     logger.info(

orchestrator/migrations/helpers.py

@@ -880,10 +880,10 @@ def delete_product(conn: sa.engine.Connection, name: str) -> None:
                 RETURNING  product_id
             ),
             deleted_p_pb AS (
-                DELETE FROM product_product_blocks WHERE product_id IN (SELECT product_id FROM deleted_p)
+                DELETE FROM product_product_blocks WHERE product_id = ANY(SELECT product_id FROM deleted_p)
             ),
             deleted_pb_rt AS (
-                DELETE FROM products_workflows WHERE product_id IN (SELECT product_id FROM deleted_p)
+                DELETE FROM products_workflows WHERE product_id = ANY(SELECT product_id FROM deleted_p)
             )
             SELECT * from deleted_p;
             """
@@ -911,10 +911,10 @@ def delete_product_block(conn: sa.engine.Connection, name: str) -> None:
                 RETURNING  product_block_id
             ),
             deleted_p_pb AS (
-                DELETE FROM product_product_blocks WHERE product_block_id IN (SELECT product_block_id FROM deleted_pb)
+                DELETE FROM product_product_blocks WHERE product_block_id =ANY(SELECT product_block_id FROM deleted_pb)
             ),
             deleted_pb_rt AS (
-                DELETE FROM product_block_resource_types WHERE product_block_id IN (SELECT product_block_id FROM deleted_pb)
+                DELETE FROM product_block_resource_types WHERE product_block_id =ANY(SELECT product_block_id FROM deleted_pb)
             )
             SELECT * from deleted_pb;
             """
@@ -968,7 +968,7 @@ def delete_resource_type(conn: sa.engine.Connection, resource_type: str) -> None
                 RETURNING  resource_type_id
             ),
             deleted_pb_rt AS (
-                DELETE FROM product_block_resource_types WHERE resource_type_id IN (SELECT resource_type_id FROM deleted_pb)
+                DELETE FROM product_block_resource_types WHERE resource_type_id =ANY(SELECT resource_type_id FROM deleted_pb)
             )
             SELECT * from deleted_pb;
             """

orchestrator/settings.py

@@ -13,6 +13,7 @@
 
 import secrets
 import string
+import warnings
 from pathlib import Path
 from typing import Literal
 
@@ -23,6 +24,10 @@
 from pydantic_forms.types import strEnum
 
 
+class OrchestratorDeprecationWarning(DeprecationWarning):
+    pass
+
+
 class ExecutorType(strEnum):
     WORKER = "celery"
     THREADPOOL = "threadpool"
@@ -49,7 +54,7 @@ class AppSettings(BaseSettings):
     EXECUTOR: str = ExecutorType.THREADPOOL
     WORKFLOWS_SWAGGER_HOST: str = "localhost"
     WORKFLOWS_GUI_URI: str = "http://localhost:3000"
-    DATABASE_URI: PostgresDsn = "postgresql://nwa:nwa@localhost/orchestrator-core"  # type: ignore
+    DATABASE_URI: PostgresDsn = "postgresql+psycopg://nwa:nwa@localhost/orchestrator-core"  # type: ignore
     MAX_WORKERS: int = 5
     MAIL_SERVER: str = "localhost"
     MAIL_PORT: int = 25
@@ -88,6 +93,22 @@ class AppSettings(BaseSettings):
     VALIDATE_OUT_OF_SYNC_SUBSCRIPTIONS: bool = False
     FILTER_BY_MODE: Literal["partial", "exact"] = "exact"
 
+    def __init__(self) -> None:
+        super(AppSettings, self).__init__()
+        self.DATABASE_URI = PostgresDsn(convert_database_uri(str(self.DATABASE_URI)))
+
+
+def convert_database_uri(db_uri: str) -> str:
+    if db_uri.startswith(("postgresql://", "postgresql+psycopg2://")):
+        db_uri = "postgresql+psycopg" + db_uri[db_uri.find("://") :]
+        warnings.filterwarnings("always", category=OrchestratorDeprecationWarning)
+        warnings.warn(
+            "DATABASE_URI converted to postgresql+psycopg:// format, please update your enviroment variable",
+            OrchestratorDeprecationWarning,
+            stacklevel=2,
+        )
+    return db_uri
+
 
 app_settings = AppSettings()
 

pyproject.toml

@@ -47,7 +47,7 @@ dependencies = [
     "itsdangerous",
     "Jinja2==3.1.5",
     "orjson==3.10.15",
-    "psycopg2-binary==2.9.10",
+    "psycopg[binary]==3.2.5",
     "pydantic[email]~=2.8.2",
     "pydantic-settings~=2.8.0",
     "python-dateutil==2.8.2",

test/unit_tests/conftest.py

@@ -174,7 +174,7 @@ def db_uri(worker_id):
         Database uri to be used in the test thread
 
     """
-    database_uri = os.environ.get("DATABASE_URI", "postgresql://nwa:nwa@localhost/orchestrator-core-test")
+    database_uri = os.environ.get("DATABASE_URI", "postgresql+psycopg://nwa:nwa@localhost/orchestrator-core-test")
     if worker_id == "master":
         # pytest is being run without any workers
         return database_uri
@@ -205,9 +205,9 @@ def database(db_uri):
         url.database = "postgres"
     engine = create_engine(url)
     with closing(engine.connect()) as conn:
-        conn.execute(text("COMMIT;"))
-        conn.execute(text(f'DROP DATABASE IF EXISTS "{db_to_create}";'))
-        conn.execute(text("COMMIT;"))
+        conn.commit()
+        conn.execution_options(isolation_level="AUTOCOMMIT").execute(text(f'DROP DATABASE IF EXISTS "{db_to_create}";'))
+        conn.commit()
         conn.execute(text(f'CREATE DATABASE "{db_to_create}";'))
 
     run_migrations(db_uri)
@@ -218,8 +218,10 @@ def database(db_uri):
     finally:
         db.wrapped_database.engine.dispose()
         with closing(engine.connect()) as conn:
-            conn.execute(text("COMMIT;"))
-            conn.execute(text(f'DROP DATABASE IF EXISTS "{db_to_create}";'))
+            conn.commit()
+            conn.execution_options(isolation_level="AUTOCOMMIT").execute(
+                text(f'DROP DATABASE IF EXISTS "{db_to_create}";')
+            )
 
 
 @pytest.fixture(autouse=True)