psycopg: upgrade from psycopg2 to psycopg (v3)

CHANGELOG.md

@@ -1,4 +1,7 @@
 # Yellowbox Changelog
+## 0.11.0
+### Changed
+* psycopg: upgrade from psycopg2 to psycopg (v3)
 ## 0.10.0
 ### Removed
 * dropped support for python 3.8, 3.9

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "yellowbox"
-version = "0.9.0"
+version = "0.11.0"
 description = ""
 authors = ["biocatch ltd"]
 license = "MIT"
@@ -25,8 +25,8 @@
 azure-storage-blob = { version = ">=12.0.0", optional = true }
 
 cffi = { version = ">=1.14.0", optional = true }
-sqlalchemy = { version = ">=1.3.0", optional = true }
-psycopg2 = { version = ">=2.8.6", optional = true }
+sqlalchemy = { version = ">=2.0.0", optional = true }
+psycopg = { version = "^3.2.12", optional = true, extras = ["binary"] }
 simple_websocket_server = { version = "*", optional = true }
 websockets = { version = "*", optional = true }
 hvac = { version = "*", optional = true }
@@ -86,7 +86,7 @@
 [tool.poetry.extras]
 azure = ["azure-storage-blob", "cffi"]
 kafka = ["kafka-python", "confluent-kafka"]
-postgresql = ["sqlalchemy", "psycopg2", "SQLAlchemy-Utils"]
+postgresql = ["sqlalchemy", "psycopg", "SQLAlchemy-Utils"]
 mssql = ['sqlalchemy', "SQLAlchemy-Utils"]
 rabbit = ["pika"]
 redis = ["redis", "async-timeout"]
@@ -96,7 +96,7 @@
 gcs = []  # empty for now, but we want to maintain forwards compatiblity in case we even want to add extras here
 aerospike = ["aerospike"]
 
-dev = ["redis", "async-timeout", "pika", "kafka-python", "azure-storage-blob", "cffi", "sqlalchemy", "psycopg2",
+dev = ["redis", "async-timeout", "pika", "kafka-python", "azure-storage-blob", "cffi", "sqlalchemy", "psycopg",
     "simple_websocket_server", "websocket_client", "starlette", "python-igraph", "uvicorn", "websockets", "hvac",
     'pyodbc', 'SQLAlchemy-Utils', "aerospike"]
 

tests/conftest.py

@@ -7,6 +7,14 @@
 from yellowbox.containers import create_and_pull as _create_and_pull, is_removed
 
 
+@fixture
+def sqlalchemy_version():
+    """Provide the installed SQLAlchemy version for tests that need to install it in containers"""
+    import sqlalchemy  # noqa: PLC0415
+
+    return sqlalchemy.__version__
+
+
 @fixture
 def create_and_pull():
     """A wrapper around yellowbox's create_and_pull, to ensure that all created containers are removed"""

tests/extras/test_mssql.py

@@ -183,7 +183,7 @@ def test_alchemy_usage(service, engine):
     assert vals == [2, 3, 10]
 
 
-def test_remote_connection_string(service, db, engine, create_and_pull, docker_client):
+def test_remote_connection_string(service, db, engine, create_and_pull, docker_client, sqlalchemy_version: str):
     with temp_network(docker_client) as network, connect(network, service) as service_alias:
         with engine.begin() as connection:
             connection.execute(
@@ -204,16 +204,17 @@ def test_remote_connection_string(service, db, engine, create_and_pull, docker_c
         container = create_and_pull(
             docker_client,
             "laudio/pyodbc:latest",
-            'sh -c "pip install sqlalchemy==1.4.46 pyodbc && python ./main.py"',
+            f'sh -c "pip install sqlalchemy=={sqlalchemy_version} pyodbc && python ./main.py"',
             detach=True,
         )
         upload_file(
             container,
             "./main.py",
             bytes(
-                "import sqlalchemy as sa;"
-                f"e = sa.create_engine('{conn_string}');"
-                "e.execute('DELETE FROM foo WHERE x < 3');",
+                "import sqlalchemy as sa\n"
+                f"e = sa.create_engine('{conn_string}')\n"
+                "with e.begin() as conn:\n"
+                "    conn.execute(sa.text('DELETE FROM foo WHERE x < 3'))\n",
                 "ascii",
             ),
         )
@@ -228,7 +229,7 @@ def test_remote_connection_string(service, db, engine, create_and_pull, docker_c
         assert vals == ["three", "ten"]
 
 
-def test_remote_connection_string_host(service, db, engine, create_and_pull, docker_client):
+def test_remote_connection_string_host(service, db, engine, create_and_pull, docker_client, sqlalchemy_version: str):
     with engine.begin() as connection:
         connection.execute(
             text(
@@ -247,14 +248,17 @@ def test_remote_connection_string_host(service, db, engine, create_and_pull, doc
     container = create_and_pull(
         docker_client,
         "laudio/pyodbc:latest",
-        'sh -c "pip install sqlalchemy==1.4.46 pyodbc && python ./main.py"',
+        f'sh -c "pip install sqlalchemy=={sqlalchemy_version} pyodbc && python ./main.py"',
         detach=True,
     )
     upload_file(
         container,
         "./main.py",
         bytes(
-            f"import sqlalchemy as sa;e = sa.create_engine('{conn_string}');e.execute('DELETE FROM foo WHERE x < 3');",
+            "import sqlalchemy as sa\n"
+            f"e = sa.create_engine('{conn_string}')\n"
+            "with e.begin() as conn:\n"
+            "    conn.execute(sa.text('DELETE FROM foo WHERE x < 3'))\n",
             "ascii",
         ),
     )

tests/extras/test_postgres.py

@@ -177,7 +177,7 @@ def test_alchemy_usage(docker_client, engine):
     assert vals == [2, 3, 10]
 
 
-def test_remote_connection_string(docker_client, create_and_pull, service, engine, db):
+def test_remote_connection_string(docker_client, create_and_pull, service, engine, db, sqlalchemy_version: str):
     with temp_network(docker_client) as network, connect(network, service) as service_alias:
         with engine.begin() as connection:
             connection.execute(
@@ -192,16 +192,17 @@ def test_remote_connection_string(docker_client, create_and_pull, service, engin
         container = create_and_pull(
             docker_client,
             "python:latest",
-            'sh -c "pip install sqlalchemy==1.4.46 psycopg2 && python ./main.py"',
+            f"sh -c \"pip install sqlalchemy=={sqlalchemy_version} 'psycopg[binary]' && python ./main.py\"",
             detach=True,
         )
         upload_file(
             container,
             "./main.py",
             bytes(
-                "import sqlalchemy as sa;"
-                f"e = sa.create_engine('{conn_string}');"
-                "e.execute('DELETE FROM foo WHERE x < 3');",
+                "import sqlalchemy as sa\n"
+                f"e = sa.create_engine('{conn_string}')\n"
+                "with e.begin() as conn:\n"
+                "    conn.execute(sa.text('DELETE FROM foo WHERE x < 3'))\n",
                 "ascii",
             ),
         )
@@ -216,7 +217,7 @@ def test_remote_connection_string(docker_client, create_and_pull, service, engin
         assert vals == ["three", "ten"]
 
 
-def test_remote_connection_string_host(docker_client, create_and_pull, service, engine, db):
+def test_remote_connection_string_host(docker_client, create_and_pull, service, engine, db, sqlalchemy_version: str):
     with engine.begin() as connection:
         connection.execute(
             text(
@@ -230,14 +231,17 @@ def test_remote_connection_string_host(docker_client, create_and_pull, service,
     container = create_and_pull(
         docker_client,
         "python:latest",
-        'sh -c "pip install sqlalchemy==1.4.46 psycopg2 && python ./main.py"',
+        f"sh -c \"pip install sqlalchemy=={sqlalchemy_version} 'psycopg[binary]' && python ./main.py\"",
         detach=True,
     )
     upload_file(
         container,
         "./main.py",
         bytes(
-            f"import sqlalchemy as sa;e = sa.create_engine('{conn_string}');e.execute('DELETE FROM foo WHERE x < 3');",
+            f"import sqlalchemy as sa\n"
+            f"e = sa.create_engine('{conn_string}')\n"
+            f"with e.begin() as conn:\n"
+            f"    conn.execute(sa.text('DELETE FROM foo WHERE x < 3'))\n",
             "ascii",
         ),
     )

yellowbox/extras/postgresql.py

@@ -25,6 +25,7 @@ class PostgreSQLService(SQLService, SingleContainerService):
 
     INTERNAL_PORT = POSTGRES_INTERNAL_PORT
     DIALECT = "postgresql"
+    DEFAULT_DRIVER = "psycopg"  # Use psycopg (v3) instead of deprecated psycopg2
 
     def __init__(
         self,
@@ -49,6 +50,10 @@ def __init__(
         self.user = user
         self.password = password
         self.default_db = default_db
+
+        if "local_driver" not in kwargs:
+            kwargs["local_driver"] = self.DEFAULT_DRIVER
+
         super().__init__(
             create_and_pull_with_defaults(
                 docker_client,
@@ -79,7 +84,7 @@ def container_connection_string(
         self,
         hostname: str,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         database: str | None = None,
         options: ConnectionOptions = None,
     ):
@@ -89,7 +94,7 @@ def container_connection_string(
     def host_connection_string(
         self,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         database: str | None = None,
         options: ConnectionOptions = None,
     ):
@@ -109,7 +114,7 @@ def engine(self, **kwargs):
     )
     def connection(self, **kwargs):
         """
-        Create an sqlalchemy Connection connected to the service's default db.
+        Create a sqlalchemy Connection connected to the service's default db.
         """
         return self.engine().connect(**kwargs)
 

yellowbox/extras/sql_base.py

@@ -35,7 +35,7 @@ def __init__(self, name: str, owner: SQLService):
     def local_connection_string(
         self,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         options: ConnectionOptions | None | AsDefault = as_default,
     ):
         return self.owner.local_connection_string(dialect, driver, database=self.name, options=options)
@@ -44,15 +44,15 @@ def container_connection_string(
         self,
         hostname: str,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         options: ConnectionOptions | None = None,
     ):
         return self.owner.container_connection_string(hostname, dialect, driver, database=self.name, options=options)
 
     def host_connection_string(
         self,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         options: ConnectionOptions | None = None,
     ):
         return self.owner.host_connection_string(dialect, driver, database=self.name, options=options)
@@ -125,7 +125,7 @@ def local_connection_string(
         options: ConnectionOptions | AsDefault | None = as_default,
     ) -> str:
         """
-        Generate an sqlalchemy-style connection string to the database in the service from the docker host.
+        Generate a sqlalchemy-style connection string to the database in the service from the docker host.
         Args:
             dialect: The dialect of the sql server.
             driver: additional driver for sqlalchemy to use.
@@ -153,13 +153,13 @@ def container_connection_string(
         self,
         hostname: str,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         *,
         database: str,
         options: ConnectionOptions | None = None,
     ) -> str:
         """
-        Generate an sqlalchemy-style connection string to the database in the service from another container on a
+        Generate a sqlalchemy-style connection string to the database in the service from another container on a
          common network.
         Args:
             hostname: the alias of the container.
@@ -171,6 +171,8 @@ def container_connection_string(
         if dialect is as_default:
             dialect = self.DIALECT
 
+        if driver is as_default:
+            driver = self.local_driver
         if driver is not None:
             dialect += "+" + driver
 
@@ -181,13 +183,13 @@ def container_connection_string(
     def host_connection_string(
         self,
         dialect: str | AsDefault = as_default,
-        driver: str | None = None,
+        driver: str | AsDefault | None = as_default,
         *,
         database: str,
         options: ConnectionOptions | None = None,
     ) -> str:
         """
-        Generate an sqlalchemy-style connection string to the database in the service from another container.
+        Generate a sqlalchemy-style connection string to the database in the service from another container.
         Args:
             dialect: The dialect of the sql server.
             driver: additional driver for sqlalchemy to use.
@@ -197,6 +199,8 @@ def host_connection_string(
         if dialect is as_default:
             dialect = self.DIALECT
 
+        if driver is as_default:
+            driver = self.local_driver
         if driver is not None:
             dialect += "+" + driver