Add driver parameter to get_uri() to facilitate psycopg3 support (#10)

Also updates the tests to check that SQLAlchemy works with both psycopg2
and psycopg3.

Author: aaron-siegel

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pixeltable-pgserver"
-version = "0.2.6"
+version = "0.2.7"
 description = "Embedded Postgres Server for Pixeltable"
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,11 +26,15 @@ dev = [
 ]
 test = [
     "pytest",
-    "psycopg2-binary",
+    "psycopg2-binary==2.9.9",
+    "psycopg[binary]==3.1.18",
     "sqlalchemy>=2",
     "sqlalchemy-utils"
 ]
 
+[tool.isort]
+line_length = 120
+
 [tool.setuptools.packages.find]
 where = ["src"]  # list of folders that contain the packages (["."] by default)
 include = ["pixeltable_pgserver*"]  # package names should match these glob patterns (["*"] by default)

src/pixeltable_pgserver/postgres_server.py

@@ -1,29 +1,30 @@
-from pathlib import Path
-from typing import Optional, Dict, Union
-import shutil
 import atexit
-import subprocess
-import os
 import logging
+import os
 import platform
-import psutil
+import shutil
+import subprocess
 import time
+from pathlib import Path
+from typing import Optional, Union
+
+import psutil
 
 from ._commands import POSTGRES_BIN_PATH, initdb, pg_ctl
-from .utils import find_suitable_port, find_suitable_socket_dir, DiskList, PostmasterInfo, process_is_running
+from .utils import DiskList, PostmasterInfo, find_suitable_port, find_suitable_socket_dir
 
 if platform.system() != 'Windows':
-    from .utils import ensure_user_exists, ensure_prefix_permissions, ensure_folder_permissions
+    from .utils import ensure_folder_permissions, ensure_prefix_permissions, ensure_user_exists
 
 _logger = logging.getLogger('pixeltable_pgserver')
 
 class PostgresServer:
     """ Provides a common interface for interacting with a server.
     """
-    import platformdirs
     import fasteners
+    import platformdirs
 
-    _instances : Dict[Path, 'PostgresServer'] = {}
+    _instances : dict[Path, 'PostgresServer'] = {}
 
     # NB home does not always support locking, eg NFS or LUSTRE (eg some clusters)
     # so, use user_runtime_path instead, which seems to be in a local filesystem
@@ -75,10 +76,10 @@ def get_pid(self) -> Optional[int]:
         """
         return self.get_postmaster_info().pid
 
-    def get_uri(self, database : Optional[str] = None) -> str:
+    def get_uri(self, database: Optional[str] = None, driver: Optional[str] = None) -> str:
         """ Returns a connection string for the postgresql server.
         """
-        return self.get_postmaster_info().get_uri(database=database)
+        return self.get_postmaster_info().get_uri(database=database, driver=driver)
 
     def ensure_pgdata_inited(self) -> None:
         """ Initializes the pgdata directory if it is not already initialized.

src/pixeltable_pgserver/utils.py

@@ -1,16 +1,15 @@
-from pathlib import Path
-import typing
-from typing import Optional, List, Dict
-import subprocess
+import datetime
+import hashlib
 import json
 import logging
-import hashlib
-import socket
 import platform
+import socket
 import stat
+import subprocess
+from pathlib import Path
+from typing import Dict, List, Optional
+
 import psutil
-import datetime
-import shutil
 
 _logger = logging.getLogger('pixeltable_pgserver')
 
@@ -87,16 +86,20 @@ def read_from_pgdata(cls, pgdata : Path) -> Optional['PostmasterInfo']:
         lines = postmaster_file.read_text().splitlines()
         return cls(lines)
 
-    def get_uri(self, user : str = 'postgres', database : Optional[str] = None) -> str:
+    def get_uri(self, user: str = 'postgres', database: Optional[str] = None, driver: Optional[str] = None) -> str:
         """ Returns a connection uri string for the postgresql server using the information in postmaster.pid"""
         if database is None:
             database = user
+        if driver is None:
+            driver_suffix = ''
+        else:
+            driver_suffix = f'+{driver}'
 
         if self.socket_dir is not None:
-            return f"postgresql://{user}:@/{database}?host={self.socket_dir}"
+            return f"postgresql{driver_suffix}://{user}:@/{database}?host={self.socket_dir}"
         elif self.port is not None:
             assert self.hostname is not None
-            return f"postgresql://{user}:@{self.hostname}:{self.port}/{database}"
+            return f"postgresql{driver_suffix}://{user}:@{self.hostname}:{self.port}/{database}"
         else:
             raise RuntimeError("postmaster.pid does not contain port or socket information")
 

tests/test_pgserver.py

@@ -1,25 +1,27 @@
-import pytest
-import pixeltable_pgserver
-import subprocess
-import tempfile
-from typing import Optional, Union
+import logging
 import multiprocessing as mp
+import os
+import platform
 import shutil
-from pathlib import Path
-import pixeltable_pgserver.utils
 import socket
-from pixeltable_pgserver.utils import find_suitable_port, process_is_running
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Optional, Union
+
 import psutil
-import platform
+import pytest
 import sqlalchemy as sa
-import datetime
-from sqlalchemy_utils import database_exists, create_database
-import logging
-import os
+from sqlalchemy_utils import create_database, database_exists
 
-def _check_sqlalchemy_works(srv : pixeltable_pgserver.PostgresServer):
+import pixeltable_pgserver
+import pixeltable_pgserver.utils
+from pixeltable_pgserver.utils import find_suitable_port, process_is_running
+
+
+def _check_sqlalchemy_works(srv: pixeltable_pgserver.PostgresServer, driver: Optional[str] = None):
     database_name = 'testdb'
-    uri = srv.get_uri(database_name)
+    uri = srv.get_uri(database_name, driver)
 
     if not database_exists(uri):
         create_database(uri)
@@ -39,7 +41,7 @@ def _check_sqlalchemy_works(srv : pixeltable_pgserver.PostgresServer):
         assert result
         assert result[0] == 1
 
-def _check_time_zones(srv : pixeltable_pgserver.PostgresServer):
+def _check_time_zones(srv: pixeltable_pgserver.PostgresServer):
     # Check that time zone information was properly compiled
     database_name = 'testdb'
     uri = srv.get_uri(database_name)
@@ -78,7 +80,8 @@ def _check_server(pg : pixeltable_pgserver.PostgresServer) -> int:
     # parse second row (first two are headers)
     ret_path = Path(ret.splitlines()[2].strip())
     assert pg.pgdata == ret_path
-    _check_sqlalchemy_works(pg)
+    _check_sqlalchemy_works(pg, None)       # Test with psycopg2 (default)
+    _check_sqlalchemy_works(pg, 'psycopg')  # Test with psycopg3
     _check_time_zones(pg)
     return postmaster_info.pid