fix: remove potential circular import psycopg [backport 3.10] (#14338)

Backport bea1bd2 from #14329 to 3.10.

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]

.github/workflows/system-tests.yml

@@ -27,22 +27,18 @@ jobs:
     env:
       TEST_LIBRARY: python
       WEBLOG_VARIANT: ${{ matrix.weblog-variant }}
-      # system-tests requires an API_KEY, but it does not have to be a valid key, as long as we don't run a scenario
-      # that make assertion on backend data. Using a fake key allow to run system tests on PR originating from forks.
-      # If ever it's needed, a valid key exists in the repo, using ${{ secrets.DD_API_KEY }}
-      DD_API_KEY: ${{ secrets.FAKE_DD_API_KEY }}
       CMAKE_BUILD_PARALLEL_LEVEL: 12
       SYSTEM_TESTS_AWS_ACCESS_KEY_ID: ${{ secrets.IDM_AWS_ACCESS_KEY_ID }}
       SYSTEM_TESTS_AWS_SECRET_ACCESS_KEY: ${{ secrets.IDM_AWS_SECRET_ACCESS_KEY }}
+      DD_API_KEY: ${{ secrets.FAKE_DD_API_KEY }}
     steps:
-
       - name: Checkout system tests
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '13b08f0faf7673656bd68dad5c142ad04eddee69'
+          ref: 'a29091199a21a3a83082cd64b2658a5c59453f33'
 
       - name: Checkout dd-trace-py
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -81,10 +77,6 @@ jobs:
     env:
       TEST_LIBRARY: python
       WEBLOG_VARIANT: ${{ matrix.weblog-variant }}
-      # system-tests requires an API_KEY, but it does not have to be a valid key, as long as we don't run a scenario
-      # that make assertion on backend data. Using a fake key allow to run system tests on PR originating from forks.
-      # If ever it's needed, a valid key exists in the repo, using ${{ secrets.DD_API_KEY }}
-      DD_API_KEY: ${{ secrets.FAKE_DD_API_KEY }}
       CMAKE_BUILD_PARALLEL_LEVEL: 12
       SYSTEM_TESTS_AWS_ACCESS_KEY_ID: ${{ secrets.IDM_AWS_ACCESS_KEY_ID }}
       SYSTEM_TESTS_AWS_SECRET_ACCESS_KEY: ${{ secrets.IDM_AWS_SECRET_ACCESS_KEY }}
@@ -96,7 +88,7 @@ jobs:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '13b08f0faf7673656bd68dad5c142ad04eddee69'
+          ref: 'a29091199a21a3a83082cd64b2658a5c59453f33'
 
       - name: Build runner
         uses: ./.github/actions/install_runner
@@ -277,7 +269,7 @@ jobs:
           persist-credentials: false
           repository: 'DataDog/system-tests'
           # Automatically managed, use scripts/update-system-tests-version to update
-          ref: '13b08f0faf7673656bd68dad5c142ad04eddee69'
+          ref: 'a29091199a21a3a83082cd64b2658a5c59453f33'
       - name: Checkout dd-trace-py
         uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:

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}"
+            )