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

Backports #5927 

## 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/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):

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

@@ -137,6 +137,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

@@ -143,6 +143,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'"""