chore(psycopg2): update integration implementation (#2591)

The earier supported version of psycopg2 is 2.7, which means that the
checks for this version can be removed. Moreover, the parse_dsn method
is already available on this version so we can use that instead of a
custom implementation.

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: P403n1x87

ddtrace/contrib/psycopg/patch.py

@@ -1,5 +1,6 @@
 # 3p
 import psycopg2
+from psycopg2.sql import Composable
 
 from ddtrace import Pin
 from ddtrace import config
@@ -27,9 +28,6 @@
 
 PSYCOPG2_VERSION = parse_version(psycopg2.__version__)
 
-if PSYCOPG2_VERSION >= (2, 7):
-    from psycopg2.sql import Composable
-
 
 def patch():
     """Patch monkey patches psycopg's connection function
@@ -47,14 +45,15 @@ def unpatch():
     if getattr(psycopg2, "_datadog_patch", False):
         setattr(psycopg2, "_datadog_patch", False)
         psycopg2.connect = _connect
+        _unpatch_extensions(_psycopg2_extensions)
 
 
 class Psycopg2TracedCursor(dbapi.TracedCursor):
     """TracedCursor for psycopg2"""
 
     def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
         # treat psycopg2.sql.Composable resource objects as strings
-        if PSYCOPG2_VERSION >= (2, 7) and isinstance(resource, Composable):
+        if isinstance(resource, Composable):
             resource = resource.as_string(self.__wrapped__)
 
         return super(Psycopg2TracedCursor, self)._trace_method(method, name, resource, extra_tags, *args, **kwargs)

ddtrace/contrib/sqlalchemy/engine.py

@@ -144,10 +144,10 @@ def _set_tags_from_url(span, url):
 def _set_tags_from_cursor(span, vendor, cursor):
     """attempt to set db connection tags by introspecting the cursor."""
     if "postgres" == vendor:
-        if hasattr(cursor, "connection") and hasattr(cursor.connection, "dsn"):
+        if hasattr(cursor, "connection"):
             dsn = getattr(cursor.connection, "dsn", None)
             if dsn:
                 d = sqlx.parse_pg_dsn(dsn)
-                span.set_tag(sqlx.DB, d.get("dbname"))
-                span.set_tag(netx.TARGET_HOST, d.get("host"))
-                span.set_tag(netx.TARGET_PORT, d.get("port"))
+                span._set_str_tag(sqlx.DB, d.get("dbname"))
+                span._set_str_tag(netx.TARGET_HOST, d.get("host"))
+                span.set_metric(netx.TARGET_PORT, int(d.get("port")))

ddtrace/ext/sql.py

@@ -1,10 +1,14 @@
+from typing import Dict
+
+
 # tags
 QUERY = "sql.query"  # the query text
 ROWS = "sql.rows"  # number of rows returned by a query
 DB = "sql.db"  # the name of the database
 
 
 def normalize_vendor(vendor):
+    # type: (str) -> str
     """Return a canonical name for a type of database."""
     if not vendor:
         return "db"  # should this ever happen?
@@ -16,13 +20,15 @@ def normalize_vendor(vendor):
         return vendor
 
 
-def parse_pg_dsn(dsn):
-    """
-    Return a dictionary of the components of a postgres DSN.
+try:
+    from psycopg2.extensions import parse_dsn as parse_pg_dsn  # type: ignore[import]
+except ImportError:
 
-    >>> parse_pg_dsn('user=dog port=1543 dbname=dogdata')
-    {'user':'dog', 'port':'1543', 'dbname':'dogdata'}
-    """
-    # FIXME: replace by psycopg2.extensions.parse_dsn when available
-    # https://github.com/psycopg/psycopg2/pull/321
-    return {c.split("=")[0]: c.split("=")[1] for c in dsn.split() if "=" in c}
+    def 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("=", maxsplit=1) for _ in dsn.split())