fix(sql): use psycopg's parse_dsn when available (backport #5927 to 1.13) (#5964)

Backports #5927 + backports a
[refactor](082e1e1)
required by this fix

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/contributing.html#Release-Note-Guidelines)
are followed.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.

---------

Co-authored-by: Gabriele N. Tornetta <[email protected]>

Author: mabdinur

ddtrace/bootstrap/sitecustomize.py

@@ -4,7 +4,6 @@
 """
 from ddtrace import LOADED_MODULES  # isort:skip
 
-from functools import partial  # noqa
 import logging  # noqa
 import os  # noqa
 import sys
@@ -148,7 +147,7 @@ def drop(module_name):
     # to the newly imported threading module to allow it to retrieve the correct
     # thread object information, like the thread name. We register a post-import
     # hook on the threading module to perform this update.
-    @partial(ModuleWatchdog.register_module_hook, "threading")
+    @ModuleWatchdog.after_module_imported("threading")
     def _(threading):
         logging.threading = threading
 

ddtrace/ext/sql.py

@@ -1,9 +1,13 @@
 from typing import Dict
 
 from ddtrace.internal.compat import ensure_pep562
+from ddtrace.internal.logger import get_logger
+from ddtrace.internal.module import ModuleWatchdog
 from ddtrace.vendor.debtcollector import deprecate
 
 
+log = get_logger(__name__)
+
 # tags
 QUERY = "sql.query"  # the query text
 ROWS = "sql.rows"  # number of rows returned by a query
@@ -23,14 +27,56 @@ def normalize_vendor(vendor):
         return vendor
 
 
-def parse_pg_dsn(dsn):
+def _dd_parse_pg_dsn(dsn):
     # type: (str) -> Dict[str, str]
     """
     Return a dictionary of the components of a postgres DSN.
     >>> parse_pg_dsn('user=dog port=1543 dbname=dogdata')
     {'user':'dog', 'port':'1543', 'dbname':'dogdata'}
     """
-    return dict(_.split("=", 1) for _ in dsn.split())
+    dsn_dict = dict()
+    try:
+        # Provides a default implementation for parsing DSN strings.
+        # The following is an example of a valid DSN string that fails to be parsed:
+        # "db=moon user=ears options='-c statement_timeout=1000 -c lock_timeout=250'"
+        dsn_dict = dict(_.split("=", 1) for _ in dsn.split())
+    except Exception:
+        log.debug("Failed to parse postgres dsn connection", exc_info=True)
+    return dsn_dict
+
+
+# Do not import from psycopg directly! This reference will be updated at runtime to use
+# a better implementation that is provided by the psycopg library.
+# This is done to avoid circular imports.
+parse_pg_dsn = _dd_parse_pg_dsn
+
+
+@ModuleWatchdog.after_module_imported("psycopg2")
+def use_psycopg2_parse_dsn(psycopg_module):
+    """Replaces parse_pg_dsn with the helper function defined in psycopg2"""
+    global parse_pg_dsn
+
+    try:
+        from psycopg2.extensions import parse_dsn
+
+        parse_pg_dsn = parse_dsn
+    except ImportError:
+        # Best effort, we'll use our own parser: _dd_parse_pg_dsn
+        pass
+
+
+@ModuleWatchdog.after_module_imported("psycopg")
+def use_psycopg3_parse_dsn(psycopg_module):
+    """Replaces parse_pg_dsn with the helper function defined in psycopg3"""
+    global parse_pg_dsn
+
+    try:
+        from psycopg.conninfo import conninfo_to_dict
+
+        parse_pg_dsn = conninfo_to_dict
+    except ImportError:
+        # Best effort, we'll use our own parser: _dd_parse_pg_dsn
+        pass
 
 
 def __getattr__(name):

ddtrace/internal/module.py

@@ -525,6 +525,15 @@ def unregister_module_hook(cls, module, hook):
         except ValueError:
             raise ValueError("Hook %r not registered for module %r" % (hook, module))
 
+    @classmethod
+    def after_module_imported(cls, module):
+        # type: (str) -> Callable[[ModuleHookType], None]
+        def _(hook):
+            # type: (ModuleHookType) -> None
+            cls.register_module_hook(module, hook)
+
+        return _
+
     @classmethod
     def register_pre_exec_module_hook(cls, cond, hook):
         # type: (Type[ModuleWatchdog], PreExecHookCond, PreExecHookType) -> None

docs/spelling_wordlist.txt

@@ -72,6 +72,7 @@ doctest
 dogpile
 dogpile.cache
 dogstatsd
+dsn
 elasticsearch
 elasticsearch1
 elasticsearch7

releasenotes/notes/fix-tracer-parse-pg-dsn-take2-7ef89362545b2721.yaml

@@ -0,0 +1,3 @@
+fixes:
+  - |
+    psycopg: Fixes ``ValueError`` raised when dsn connection strings are parsed. This was fixed in ddtrace v1.9.0 and was re-introduced in v1.13.0.

tests/contrib/psycopg/test_psycopg.py

@@ -136,6 +136,15 @@ def assert_conn_is_traced(self, db, service):
         self.assertIsNone(root.get_tag("sql.query"))
         self.reset()
 
+    def test_psycopg3_connection_with_string(self):
+        # Regression test for DataDog/dd-trace-py/issues/5926
+        configs_arr = ["{}={}".format(k, v) for k, v in POSTGRES_CONFIG.items()]
+        configs_arr.append("options='-c statement_timeout=1000 -c lock_timeout=250'")
+        conn = psycopg.connect(" ".join(configs_arr))
+
+        Pin.get_from(conn).clone(service="postgres", tracer=self.tracer).onto(conn)
+        self.assert_conn_is_traced(conn, "postgres")
+
     def test_opentracing_propagation(self):
         # ensure OpenTracing plays well with our integration
         query = """SELECT 'tracing'"""

tests/contrib/psycopg2/test_psycopg.py

@@ -142,6 +142,15 @@ def assert_conn_is_traced(self, db, service):
         self.assertIsNone(root.get_tag("sql.query"))
         self.reset()
 
+    def test_psycopg2_connection_with_string(self):
+        # Regression test for DataDog/dd-trace-py/issues/5926
+        configs_arr = ["{}={}".format(k, v) for k, v in POSTGRES_CONFIG.items()]
+        configs_arr.append("options='-c statement_timeout=1000 -c lock_timeout=250'")
+        conn = psycopg2.connect(" ".join(configs_arr))
+
+        Pin.get_from(conn).clone(service="postgres", tracer=self.tracer).onto(conn)
+        self.assert_conn_is_traced(conn, "postgres")
+
     def test_opentracing_propagation(self):
         # ensure OpenTracing plays well with our integration
         query = """SELECT 'tracing'"""

tests/internal/test_module.py

@@ -73,6 +73,14 @@ def test_import_module_hook_for_imported_module(module_watchdog):
     hook.assert_called_once_with(module)
 
 
+def test_after_module_imported_decorator(module_watchdog):
+    hook = mock.Mock()
+    module = sys.modules[__name__]
+    module_watchdog.after_module_imported(module.__name__)(hook)
+
+    hook.assert_called_once_with(module)
+
+
 def test_register_hook_without_install():
     with pytest.raises(RuntimeError):
         ModuleWatchdog.register_origin_hook(__file__, mock.Mock())