Add support for psycopg3

.pre-commit-config.yaml

@@ -1,9 +1,9 @@
 repos:
   - repo: https://github.com/pycqa/isort
-    rev: 5.10.1
+    rev: 5.11.5
     hooks:
       - id: isort
-  - repo: https://gitlab.com/pycqa/flake8
+  - repo: https://github.com/pycqa/flake8
     rev: 4.0.1
     hooks:
       - id: flake8

README.md

@@ -96,15 +96,15 @@ Here's the raw data:
 
 ```
 
-pip install DSLR psycopg2 # or psycopg2-binary
+pip install DSLR psycopg2 # or psycopg2-binary, or psycopg
 
 ```
 
 **Install using pipx**
 
 ```
 
-pipx install DSLR[psycopg2] # or psycopg2-binary
+pipx install DSLR[psycopg2] # or psycopg2-binary, or psycopg
 
 ````
 

dslr/operations.py

@@ -2,7 +2,10 @@
 from datetime import datetime
 from typing import List, Optional
 
-from psycopg2 import sql
+try:
+    from psycopg import sql
+except ImportError:
+    from psycopg2 import sql
 
 from .config import settings
 from .runner import exec_shell, exec_sql

dslr/pg_client.py

@@ -1,7 +1,10 @@
 from typing import Any, List, Optional, Tuple
 
-import psycopg2
-import psycopg2.extensions
+try:
+    import psycopg as psycopg
+except ImportError:
+    import psycopg2 as psycopg
+    import psycopg2.extensions
 
 from dslr.console import console
 
@@ -20,17 +23,24 @@ def __init__(self, host, port, user, password, dbname):
         self.password = password
         self.dbname = dbname
 
-        self.conn = psycopg2.connect(
+        self.conn = psycopg.connect(
             host=host,
             port=port,
             user=user,
             password=password,
             dbname=dbname,
         )
-        self.conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
-
+        self._set_autocommit()
         self.cur = self.conn.cursor()
 
+    def _set_autocommit(self):
+        if hasattr(self.conn, "set_autocommit"):
+            self.conn.set_autocommit(True)  # type: ignore
+        else:
+            self.conn.set_isolation_level(  # type: ignore
+                psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT
+            )
+
     def execute(self, sql, data) -> Optional[List[Tuple[Any, ...]]]:
         if settings.debug:
             console.log(f"SQL: {sql}")
@@ -40,7 +50,7 @@ def execute(self, sql, data) -> Optional[List[Tuple[Any, ...]]]:
 
         try:
             result = self.cur.fetchall()
-        except psycopg2.ProgrammingError:
+        except psycopg.ProgrammingError:
             result = None
 
         return result

dslr/runner.py

@@ -3,7 +3,10 @@
 from collections import namedtuple
 from typing import Any, List, Optional, Tuple, Union
 
-from psycopg2 import sql
+try:
+    from psycopg import sql
+except ImportError:
+    from psycopg2 import sql
 
 from dslr.pg_client import PGClient
 

pyproject.toml

@@ -27,13 +27,15 @@ timeago = "^1.0.15"
 tomli = "^2.0.1"
 psycopg2 = { version = "^2.9.3", optional = true }
 psycopg2-binary = { version = "^2.9.3", optional = true }
+psycopg = { version = "^3.1.14", optional = true }
 
 [tool.poetry.extras]
 psycopg2 = ["psycopg2"]
 psycopg2-binary = ["psycopg2-binary"]
+psycopg = ["psycopg"]
 
 [tool.poetry.dev-dependencies]
-isort = "^5.10.1"
+isort = "^5.11.5"
 flake8 = "^4.0.1"
 black = "^22.6.0"
 ipdb = "^0.13.9"
@@ -43,6 +45,9 @@ twine = "^4.0.1"
 [tool.poetry.scripts]
 dslr = "dslr.cli:cli"
 
+[tool.pyright]
+reportMissingImports = false
+
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"

tox.ini

@@ -5,23 +5,32 @@ envlist =
     black
     isort
     pyright
-    py37
-    py38
-    py39
-    py310
+    py37-psycopg2
+    py37-psycopg3
+    py38-psycopg2
+    py38-psycopg3
+    py39-psycopg2
+    py39-psycopg3
+    py310-psycopg2
+    py310-psycopg3
 
 [gh-actions]
 python =
-    3.7: py37
-    3.8: py38
-    3.9: py39
-    3.10: py310, flake8, black, isort, pyright
+    3.7: py37-psycopg2, py37-psycopg3
+    3.8: py38-psycopg2, py38-psycopg3
+    3.9: py39-psycopg2, py39-psycopg3
+    3.10: py310-psycopg2, py310-psycopg3, flake8, black, isort, pyright
 
 [testenv]
-deps = psycopg2-binary
 commands =
     python -m unittest
 
+[testenv:{py37,py38,py39,py310}-psycopg2]
+deps = psycopg2-binary
+
+[testenv:{py37,py38,py39,py310}-psycopg3]
+deps = psycopg
+
 [testenv:flake8]
 deps = flake8 >=4.0.1, <5.0.0
 commands = flake8