fix: remove potential circular import psycopg [backport 3.12] (#14340)

Backport bea1bd2 from #14329 to 3.12.

Signed-off-by: Juanjo Alvarez <[email protected]>##
Checklist
- [X] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)
Signed-off-by: Juanjo Alvarez
<[email protected]>Signed-off-by: Juanjo Alvarez
<[email protected]>Signed-off-by: Juanjo Alvarez
<[email protected]>Signed-off-by: Juanjo Alvarez
<[email protected]>

Signed-off-by: Juanjo Alvarez <[email protected]>
Co-authored-by: Juanjo Alvarez Martinez <[email protected]>

Author: dd-octo-sts[bot]

ddtrace/contrib/internal/psycopg/patch.py

@@ -26,17 +26,22 @@
 from ddtrace.trace import Pin
 
 
-try:
-    psycopg_import = import_module("psycopg")
+# These will be initialized lazily to avoid circular imports
+_original_connect = None
+_original_async_connect = None
 
-    # must get the original connect class method from the class __dict__ to use later in unpatch
-    # Python 3.11 and wrapt result in the class method being rebinded as an instance method when
-    # using unwrap
-    _original_connect = psycopg_import.Connection.__dict__["connect"]
-    _original_async_connect = psycopg_import.AsyncConnection.__dict__["connect"]
-# AttributeError can happen due to circular imports under certain integration methods
-except (ImportError, AttributeError):
-    pass
+
+def _get_psycopg3_original_methods():
+    """Get psycopg3 original method references, avoiding top-level evaluation to avoid circular imports"""
+    global _original_connect, _original_async_connect
+    if _original_connect is None or _original_async_connect is None:
+        try:
+            psycopg_import = import_module("psycopg")
+            _original_connect = psycopg_import.Connection.__dict__["connect"]
+            _original_async_connect = psycopg_import.AsyncConnection.__dict__["connect"]
+        # AttributeError can happen due to circular imports under certain integration methods
+        except (ImportError, AttributeError):
+            pass
 
 
 def _psycopg_sql_injector(dbm_comment, sql_statement):
@@ -131,6 +136,8 @@ def _patch(psycopg_module):
 
         config.psycopg["_patched_modules"].add(psycopg_module)
     else:
+        _get_psycopg3_original_methods()
+
         _w(psycopg_module, "connect", patched_connect_factory(psycopg_module))
         _w(psycopg_module, "Cursor", init_cursor_from_connection_factory(psycopg_module))
         _w(psycopg_module, "AsyncCursor", init_cursor_from_connection_factory(psycopg_module))
@@ -162,8 +169,10 @@ def _unpatch(psycopg_module):
 
             # _u throws an attribute error for Python 3.11, no __get__ on the BoundFunctionWrapper
             # unlike Python Class Methods which implement __get__
-            psycopg_module.Connection.connect = _original_connect
-            psycopg_module.AsyncConnection.connect = _original_async_connect
+            if _original_connect is not None:
+                psycopg_module.Connection.connect = _original_connect
+            if _original_async_connect is not None:
+                psycopg_module.AsyncConnection.connect = _original_async_connect
 
         pin = Pin.get_from(psycopg_module)
         if pin:

releasenotes/notes/circular-psycopg3-import-c27df66875895fc7.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    psycopg: This fix resolves a potential circular import with the psycopg3 contrib.

tests/contrib/psycopg/test_psycopg_patch.py

@@ -2,7 +2,6 @@
 # script. If you want to make changes to it, you should make sure that you have
 # removed the ``_generated`` suffix from the file name, to prevent the content
 # from being overwritten by future re-generations.
-
 from ddtrace.contrib.internal.psycopg.patch import get_version
 from ddtrace.contrib.internal.psycopg.patch import get_versions
 from ddtrace.contrib.internal.psycopg.patch import patch

tests/contrib/psycopg2/fixtures/__init__.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python3
+"""
+Reproduction script for psycopg circular import issue.
+Run with: ddtrace-run python reproduce_psycopg_cyclic_import_error.py
+"""
+
+import sys
+
+
+def clear_psycopg_modules():
+    """Clear all psycopg-related modules from sys.modules"""
+    modules_to_remove = [name for name in list(sys.modules.keys()) if ("psycopg" in name and "ddtrace" not in name)]
+    for module in modules_to_remove:
+        del sys.modules[module]
+
+
+def main():
+    clear_psycopg_modules()
+    import psycopg2  # noqa:F401
+
+
+if __name__ == "__main__":
+    main()

tests/contrib/psycopg2/fixtures/reproduce_psycopg_cyclic_import_error.py

@@ -2,6 +2,9 @@
 # script. If you want to make changes to it, you should make sure that you have
 # removed the ``_generated`` suffix from the file name, to prevent the content
 # from being overwritten by future re-generations.
+import os
+import subprocess
+import sys
 
 from ddtrace.contrib.internal.psycopg.patch import get_version
 from ddtrace.contrib.internal.psycopg.patch import get_versions
@@ -40,3 +43,32 @@ def test_and_emit_get_version(self):
         assert versions.get("psycopg2")
         emit_integration_and_version_to_test_agent("psycopg", versions["psycopg2"], "psycopg2")
         unpatch()
+
+    def test_psycopg_circular_import_fix(self):
+        fixture_path = os.path.join(os.path.dirname(__file__), "fixtures", "reproduce_psycopg_cyclic_import_error.py")
+
+        env = os.environ.copy()
+        codebase_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+        env["PYTHONPATH"] = os.pathsep.join([codebase_root, env.get("PYTHONPATH", "")])
+
+        # Run the reproduction script with ddtrace-run
+        # Note: tried with @run_in_subprocess but that fails to reproduce the error in the affected version
+        # (v3.10.2) for some reason while running the separate script works.
+        result = subprocess.run(
+            [sys.executable, "-m", "ddtrace.commands.ddtrace_run", "python", fixture_path],
+            capture_output=True,
+            text=True,
+            env=env,
+            timeout=30,
+        )
+
+        full_output = result.stdout + result.stderr
+        if "failed to enable ddtrace support for psycopg: partially initialized module" in full_output:
+            self.fail("Circular import issue detected: the original bug is present and should be fixed")
+        elif "psycopg2" in full_output and "ImportError" in full_output:
+            self.skipTest("psycopg2 not available for testing")
+        elif result.returncode != 0:
+            self.fail(
+                f"Reproduction script failed unexpectedly: "
+                f"returncode={result.returncode}, stdout={result.stdout}, stderr={result.stderr}"
+            )