Skip to content

Commit 513acc9

Browse files
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>
1 parent 527fa21 commit 513acc9

File tree

3 files changed

+22
-17
lines changed

3 files changed

+22
-17
lines changed

ddtrace/contrib/psycopg/patch.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
# 3p
22
import psycopg2
3+
from psycopg2.sql import Composable
34

45
from ddtrace import Pin
56
from ddtrace import config
@@ -27,9 +28,6 @@
2728

2829
PSYCOPG2_VERSION = parse_version(psycopg2.__version__)
2930

30-
if PSYCOPG2_VERSION >= (2, 7):
31-
from psycopg2.sql import Composable
32-
3331

3432
def patch():
3533
"""Patch monkey patches psycopg's connection function
@@ -47,14 +45,15 @@ def unpatch():
4745
if getattr(psycopg2, "_datadog_patch", False):
4846
setattr(psycopg2, "_datadog_patch", False)
4947
psycopg2.connect = _connect
48+
_unpatch_extensions(_psycopg2_extensions)
5049

5150

5251
class Psycopg2TracedCursor(dbapi.TracedCursor):
5352
"""TracedCursor for psycopg2"""
5453

5554
def _trace_method(self, method, name, resource, extra_tags, *args, **kwargs):
5655
# treat psycopg2.sql.Composable resource objects as strings
57-
if PSYCOPG2_VERSION >= (2, 7) and isinstance(resource, Composable):
56+
if isinstance(resource, Composable):
5857
resource = resource.as_string(self.__wrapped__)
5958

6059
return super(Psycopg2TracedCursor, self)._trace_method(method, name, resource, extra_tags, *args, **kwargs)

ddtrace/contrib/sqlalchemy/engine.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -144,10 +144,10 @@ def _set_tags_from_url(span, url):
144144
def _set_tags_from_cursor(span, vendor, cursor):
145145
"""attempt to set db connection tags by introspecting the cursor."""
146146
if "postgres" == vendor:
147-
if hasattr(cursor, "connection") and hasattr(cursor.connection, "dsn"):
147+
if hasattr(cursor, "connection"):
148148
dsn = getattr(cursor.connection, "dsn", None)
149149
if dsn:
150150
d = sqlx.parse_pg_dsn(dsn)
151-
span.set_tag(sqlx.DB, d.get("dbname"))
152-
span.set_tag(netx.TARGET_HOST, d.get("host"))
153-
span.set_tag(netx.TARGET_PORT, d.get("port"))
151+
span._set_str_tag(sqlx.DB, d.get("dbname"))
152+
span._set_str_tag(netx.TARGET_HOST, d.get("host"))
153+
span.set_metric(netx.TARGET_PORT, int(d.get("port")))

ddtrace/ext/sql.py

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
1+
from typing import Dict
2+
3+
14
# tags
25
QUERY = "sql.query" # the query text
36
ROWS = "sql.rows" # number of rows returned by a query
47
DB = "sql.db" # the name of the database
58

69

710
def normalize_vendor(vendor):
11+
# type: (str) -> str
812
"""Return a canonical name for a type of database."""
913
if not vendor:
1014
return "db" # should this ever happen?
@@ -16,13 +20,15 @@ def normalize_vendor(vendor):
1620
return vendor
1721

1822

19-
def parse_pg_dsn(dsn):
20-
"""
21-
Return a dictionary of the components of a postgres DSN.
23+
try:
24+
from psycopg2.extensions import parse_dsn as parse_pg_dsn # type: ignore[import]
25+
except ImportError:
2226

23-
>>> parse_pg_dsn('user=dog port=1543 dbname=dogdata')
24-
{'user':'dog', 'port':'1543', 'dbname':'dogdata'}
25-
"""
26-
# FIXME: replace by psycopg2.extensions.parse_dsn when available
27-
# https://github.com/psycopg/psycopg2/pull/321
28-
return {c.split("=")[0]: c.split("=")[1] for c in dsn.split() if "=" in c}
27+
def parse_pg_dsn(dsn):
28+
# type: (str) -> Dict[str, str]
29+
"""
30+
Return a dictionary of the components of a postgres DSN.
31+
>>> parse_pg_dsn('user=dog port=1543 dbname=dogdata')
32+
{'user':'dog', 'port':'1543', 'dbname':'dogdata'}
33+
"""
34+
return dict(_.split("=", maxsplit=1) for _ in dsn.split())

0 commit comments

Comments
 (0)