Add flake8 plugin to ensure psycopg is imported from the right place

backend/.flake8

@@ -7,6 +7,7 @@ per-file-ignores =
     ../enterprise/backend/tests/*: F841
     src/baserow/contrib/database/migrations/*: X1
     src/baserow/core/migrations/*: X1
+    src/baserow/core/psycopg.py: BRP001
 exclude =
     .git,
     __pycache__,
@@ -16,4 +17,5 @@ exclude =
 [flake8:local-plugins]
 extension =
     X1 = flake8_baserow:DocstringPlugin
+    BRP001 = flake8_baserow:BaserowPsycopgChecker
 paths = ./flake8_plugins

backend/flake8_plugins/__init__.py

@@ -1 +1 @@
-from .flake8_baserow import DocstringPlugin
+from .flake8_baserow import DocstringPlugin, BaserowPsycopgChecker

backend/flake8_plugins/flake8_baserow/__init__.py

@@ -1,3 +1,4 @@
 from .docstring import Plugin as DocstringPlugin
+from .psycopg import BaserowPsycopgChecker
 
-__all__ = ["DocstringPlugin"]
+__all__ = ["DocstringPlugin", "BaserowPsycopgChecker"]

backend/flake8_plugins/flake8_baserow/psycopg.py

@@ -0,0 +1,30 @@
+import ast
+from typing import Iterator, Tuple, Any
+
+class BaserowPsycopgChecker:
+    name = 'flake8-baserow-psycopg'
+    version = '0.1.0'
+
+    def __init__(self, tree: ast.AST, filename: str):
+        self.tree = tree
+        self.filename = filename
+
+    def run(self) -> Iterator[Tuple[int, int, str, Any]]:
+        for node in ast.walk(self.tree):
+            if isinstance(node, ast.Import):
+                for alias in node.names:
+                    if alias.name in ('psycopg', 'psycopg2'):
+                        yield (
+                            node.lineno,
+                            node.col_offset,
+                            'BRP001 Import psycopg/psycopg2 from baserow.core.psycopg instead',
+                            type(self)
+                        )
+            elif isinstance(node, ast.ImportFrom):
+                if node.module in ('psycopg', 'psycopg2'):
+                    yield (
+                        node.lineno,
+                        node.col_offset,
+                        'BRP001 Import psycopg/psycopg2 from baserow.core.psycopg instead',
+                        type(self)
+                    )

backend/flake8_plugins/tests/test_flake8_baserow_psycopg.py

@@ -0,0 +1,38 @@
+import ast
+from flake8_baserow.psycopg import BaserowPsycopgChecker
+
+
+def run_checker(code: str):
+    tree = ast.parse(code)
+    checker = BaserowPsycopgChecker(tree, 'test.py')
+    return list(checker.run())
+
+def test_direct_import():
+    code = '''
+import psycopg
+import psycopg2
+from psycopg import connect
+from psycopg2 import connect as pg_connect
+    '''
+    errors = run_checker(code)
+    assert len(errors) == 4
+    assert all(error[2].startswith('BRP001') for error in errors)
+
+def test_allowed_import():
+    code = '''
+from baserow.core.psycopg import connect
+from baserow.core.psycopg import psycopg2
+    '''
+    errors = run_checker(code)
+    assert len(errors) == 0
+
+def test_mixed_imports():
+    code = '''
+import psycopg
+from baserow.core.psycopg import connect
+from psycopg2 import connect as pg_connect
+    '''
+    errors = run_checker(code)
+    assert len(errors) == 2
+    assert errors[0][2].startswith('BRP001')
+    assert errors[1][2].startswith('BRP001')

backend/src/baserow/contrib/database/fields/utils/duration.py

@@ -708,8 +708,9 @@ def text_value_sql_to_duration(field: "DurationField") -> str:
 
 
 if is_psycopg3:
-    import psycopg
-    from psycopg.types.datetime import IntervalLoader
+    from psycopg.types.datetime import IntervalLoader  # noqa: BRP001
+
+    from baserow.core.psycopg import psycopg
 
     class BaserowIntervalLoader(IntervalLoader):
         """

backend/src/baserow/core/db.py

@@ -27,7 +27,8 @@ from django.db.transaction import Atomic, get_connection
 
 from loguru import logger
 
-from .psycopg import sql
+from baserow.core.psycopg import sql
+
 from .utils import find_intermediate_order
 
 ModelInstance = TypeVar("ModelInstance", bound=object)

changelog/entries/unreleased/refactor/1689_upgrade_to_psycopg3.json

@@ -1,7 +0,0 @@
-{
-    "type": "refactor",
-    "message": "Upgrade to psycopg3",
-    "issue_number": 1689,
-    "bullet_points": [],
-    "created_at": "2025-01-08"
-}