Skip to content

Commit fd87333

Browse files
emmettbutleremmettbutler
authored
chore(psycopg): use the core api in the psycopg integration (#14375)
This change removes logic specific to Tracing from the `psycopg` integration. It also deletes a block of dead code in `extensions.py` discovered in the course of the refactor. ## 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)
1 parent 9f80d9d commit fd87333

File tree

4 files changed

+29
-112
lines changed

4 files changed

+29
-112
lines changed

‎ddtrace/_trace/trace_handlers.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,9 @@ def _start_span(ctx: core.ExecutionContext, call_trace: bool = True, **kwargs) -
138138

139139
for tk, tv in ctx.get_item("tags", dict()).items():
140140
span.set_tag_str(tk, tv)
141+
if ctx.get_item("measured"):
142+
# PERF: avoid setting via Span.set_tag
143+
span.set_metric(_SPAN_MEASURED_KEY, 1)
141144

142145
ctx.span = span
143146

@@ -991,6 +994,7 @@ def listen():
991994
"azure.functions.patched_service_bus",
992995
"azure.functions.patched_timer",
993996
"azure.servicebus.patched_producer",
997+
"psycopg.patched_connect",
994998
):
995999
core.on(f"context.started.{context_name}", _start_span)
9961000

@@ -1003,6 +1007,7 @@ def listen():
10031007
"django.middleware.process_template_response",
10041008
"django.middleware.process_view",
10051009
"django.template.render",
1010+
"psycopg.patched_connect",
10061011
):
10071012
core.on(f"context.ended.{name}", _finish_span)
10081013

‎ddtrace/contrib/internal/psycopg/async_connection.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from ddtrace import config
2-
from ddtrace.constants import _SPAN_MEASURED_KEY
32
from ddtrace.constants import SPAN_KIND
43
from ddtrace.contrib import dbapi_async
54
from ddtrace.contrib.internal.psycopg.async_cursor import Psycopg3FetchTracedAsyncCursor
@@ -9,6 +8,7 @@
98
from ddtrace.ext import SpanKind
109
from ddtrace.ext import SpanTypes
1110
from ddtrace.ext import db
11+
from ddtrace.internal import core
1212
from ddtrace.internal.constants import COMPONENT
1313
from ddtrace.trace import Pin
1414

@@ -48,18 +48,19 @@ async def patched_connect_async(connect_func, _, args, kwargs):
4848
if not pin or not pin.enabled() or not pin._config.trace_connect:
4949
conn = await connect_func(*args, **kwargs)
5050
else:
51-
with pin.tracer.trace(
52-
"{}.{}".format(connect_func.__module__, connect_func.__name__),
51+
with core.context_with_data(
52+
"psycopg.patched_connect",
53+
span_name="{}.{}".format(connect_func.__module__, connect_func.__name__),
5354
service=ext_service(pin, pin._config),
5455
span_type=SpanTypes.SQL,
55-
) as span:
56-
span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
57-
span.set_tag_str(COMPONENT, pin._config.integration_name)
58-
if span.get_tag(db.SYSTEM) is None:
59-
span.set_tag_str(db.SYSTEM, pin._config.dbms_name)
60-
61-
# PERF: avoid setting via Span.set_tag
62-
span.set_metric(_SPAN_MEASURED_KEY, 1)
56+
pin=pin,
57+
tags={
58+
SPAN_KIND: SpanKind.CLIENT,
59+
COMPONENT: pin._config.integration_name,
60+
db.SYSTEM: pin._config.dbms_name,
61+
},
62+
measured=True,
63+
):
6364
conn = await connect_func(*args, **kwargs)
6465

6566
return patch_conn(conn, pin=pin, traced_conn_cls=traced_conn_cls)

‎ddtrace/contrib/internal/psycopg/connection.py‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from ddtrace import config
2-
from ddtrace.constants import _SPAN_MEASURED_KEY
32
from ddtrace.constants import SPAN_KIND
43
from ddtrace.contrib import dbapi
54
from ddtrace.contrib.internal.psycopg.cursor import Psycopg2FetchTracedCursor
@@ -13,6 +12,7 @@
1312
from ddtrace.ext import db
1413
from ddtrace.ext import net
1514
from ddtrace.ext import sql
15+
from ddtrace.internal import core
1616
from ddtrace.internal.constants import COMPONENT
1717
from ddtrace.trace import Pin
1818

@@ -92,18 +92,19 @@ def patched_connect(connect_func, _, args, kwargs):
9292
if not pin or not pin.enabled() or not pin._config.trace_connect:
9393
conn = connect_func(*args, **kwargs)
9494
else:
95-
with pin.tracer.trace(
96-
"{}.{}".format(connect_func.__module__, connect_func.__name__),
95+
with core.context_with_data(
96+
"psycopg.patched_connect",
97+
span_name="{}.{}".format(connect_func.__module__, connect_func.__name__),
9798
service=ext_service(pin, pin._config),
9899
span_type=SpanTypes.SQL,
99-
) as span:
100-
span.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
101-
span.set_tag_str(COMPONENT, pin._config.integration_name)
102-
if span.get_tag(db.SYSTEM) is None:
103-
span.set_tag_str(db.SYSTEM, pin._config.dbms_name)
104-
105-
# PERF: avoid setting via Span.set_tag
106-
span.set_metric(_SPAN_MEASURED_KEY, 1)
100+
pin=pin,
101+
tags={
102+
SPAN_KIND: SpanKind.CLIENT,
103+
COMPONENT: pin._config.integration_name,
104+
db.SYSTEM: pin._config.dbms_name,
105+
},
106+
measured=True,
107+
):
107108
conn = connect_func(*args, **kwargs)
108109

109110
return patch_conn(conn, pin=pin, traced_conn_cls=traced_conn_cls)

‎ddtrace/contrib/internal/psycopg/extensions.py‎

Lines changed: 0 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -1,97 +1,7 @@
1-
"""
2-
Tracing utilities for the psycopg2 potgres client library.
3-
"""
4-
import functools
5-
61
import wrapt
72

8-
from ddtrace import config
9-
from ddtrace.constants import _SPAN_MEASURED_KEY
10-
from ddtrace.constants import SPAN_KIND
11-
from ddtrace.ext import SpanKind
12-
from ddtrace.ext import SpanTypes
13-
from ddtrace.ext import db
14-
from ddtrace.ext import net
15-
from ddtrace.internal.constants import COMPONENT
16-
from ddtrace.internal.schema import schematize_database_operation
17-
183

194
def get_psycopg2_extensions(psycopg_module):
20-
class TracedCursor(psycopg_module.extensions.cursor):
21-
"""Wrapper around cursor creating one span per query"""
22-
23-
def __init__(self, *args, **kwargs):
24-
self._datadog_tracer = kwargs.pop("datadog_tracer", None)
25-
self._datadog_service = kwargs.pop("datadog_service", None)
26-
self._datadog_tags = kwargs.pop("datadog_tags", None)
27-
super(TracedCursor, self).__init__(*args, **kwargs)
28-
29-
def execute(self, query, vars=None): # noqa: A002
30-
"""just wrap the cursor execution in a span"""
31-
if not self._datadog_tracer:
32-
return psycopg_module.extensions.cursor.execute(self, query, vars)
33-
34-
with self._datadog_tracer.trace(
35-
schematize_database_operation("postgres.query", database_provider="postgresql"),
36-
service=self._datadog_service,
37-
span_type=SpanTypes.SQL,
38-
) as s:
39-
s.set_tag_str(COMPONENT, config.psycopg.integration_name)
40-
s.set_tag_str(db.SYSTEM, config.psycopg.dbms_name)
41-
42-
# set span.kind to the type of operation being performed
43-
s.set_tag_str(SPAN_KIND, SpanKind.CLIENT)
44-
45-
# PERF: avoid setting via Span.set_tag
46-
s.set_metric(_SPAN_MEASURED_KEY, 1)
47-
if s.context.sampling_priority is None or s.context.sampling_priority <= 0:
48-
return super(TracedCursor, self).execute(query, vars)
49-
50-
s.resource = query
51-
s.set_tags(self._datadog_tags)
52-
try:
53-
return super(TracedCursor, self).execute(query, vars)
54-
finally:
55-
s.set_metric(db.ROWCOUNT, self.rowcount)
56-
57-
def callproc(self, procname, vars=None): # noqa: A002
58-
"""just wrap the execution in a span"""
59-
return psycopg_module.extensions.cursor.callproc(self, procname, vars)
60-
61-
class TracedConnection(psycopg_module.extensions.connection):
62-
"""Wrapper around psycopg2 for tracing"""
63-
64-
def __init__(self, *args, **kwargs):
65-
self._datadog_tracer = kwargs.pop("datadog_tracer", None)
66-
self._datadog_service = kwargs.pop("datadog_service", None)
67-
68-
super(TracedConnection, self).__init__(*args, **kwargs)
69-
70-
# add metadata (from the connection, string, etc)
71-
dsn = psycopg_module.extensions.parse_dsn(self.dsn)
72-
self._datadog_tags = {
73-
net.TARGET_HOST: dsn.get("host"),
74-
net.TARGET_PORT: dsn.get("port"),
75-
net.SERVER_ADDRESS: dsn.get("host"),
76-
db.NAME: dsn.get("dbname"),
77-
db.USER: dsn.get("user"),
78-
db.SYSTEM: config.psycopg.dbms_name,
79-
"db.application": dsn.get("application_name"),
80-
}
81-
82-
self._datadog_cursor_class = functools.partial(
83-
TracedCursor,
84-
datadog_tracer=self._datadog_tracer,
85-
datadog_service=self._datadog_service,
86-
datadog_tags=self._datadog_tags,
87-
)
88-
89-
def cursor(self, *args, **kwargs):
90-
"""register our custom cursor factory"""
91-
kwargs.setdefault("cursor_factory", self._datadog_cursor_class)
92-
return super(TracedConnection, self).cursor(*args, **kwargs)
93-
94-
# extension hooks
955
_extensions = [
966
(
977
psycopg_module.extensions.register_type,

0 commit comments

Comments
 (0)