1from __future__ import annotations
  2
  3import importlib.metadata
  4import re
  5import time
  6import traceback
  7import weakref
  8from collections.abc import Callable, Generator
  9from contextlib import contextmanager
 10from typing import TYPE_CHECKING, Any
 11
 12from opentelemetry import context as otel_context
 13from opentelemetry import metrics, trace
 14from opentelemetry.metrics import CallbackOptions, Observation
 15from opentelemetry.semconv.metrics.db_metrics import DB_CLIENT_OPERATION_DURATION
 16
 17if TYPE_CHECKING:
 18    from opentelemetry.trace import Span
 19    from plain.postgres.connection import DatabaseConnection
 20    from plain.postgres.sources import PoolSource
 21    from psycopg import Connection as PsycopgConnection
 22
 23from opentelemetry.semconv._incubating.attributes.db_attributes import (
 24    DB_CLIENT_CONNECTION_POOL_NAME,
 25    DB_CLIENT_CONNECTION_STATE,
 26    DB_QUERY_PARAMETER_TEMPLATE,
 27    DbClientConnectionStateValues,
 28)
 29from opentelemetry.semconv._incubating.metrics.db_metrics import (
 30    DB_CLIENT_CONNECTION_COUNT,
 31    DB_CLIENT_CONNECTION_IDLE_MAX,
 32    DB_CLIENT_CONNECTION_IDLE_MIN,
 33    DB_CLIENT_CONNECTION_MAX,
 34    DB_CLIENT_CONNECTION_PENDING_REQUESTS,
 35    DB_CLIENT_CONNECTION_TIMEOUTS,
 36    DB_CLIENT_CONNECTION_USE_TIME,
 37    DB_CLIENT_CONNECTION_WAIT_TIME,
 38    DB_CLIENT_RESPONSE_RETURNED_ROWS,
 39)
 40from opentelemetry.semconv.attributes.code_attributes import (
 41    CODE_COLUMN_NUMBER,
 42    CODE_FILE_PATH,
 43    CODE_FUNCTION_NAME,
 44    CODE_LINE_NUMBER,
 45    CODE_STACKTRACE,
 46)
 47from opentelemetry.semconv.attributes.db_attributes import (
 48    DB_COLLECTION_NAME,
 49    DB_NAMESPACE,
 50    DB_OPERATION_NAME,
 51    DB_QUERY_SUMMARY,
 52    DB_QUERY_TEXT,
 53    DB_SYSTEM_NAME,
 54    DbSystemNameValues,
 55)
 56from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
 57from opentelemetry.semconv.attributes.network_attributes import (
 58    NETWORK_PEER_ADDRESS,
 59    NETWORK_PEER_PORT,
 60)
 61from opentelemetry.semconv.attributes.server_attributes import (
 62    SERVER_ADDRESS,
 63    SERVER_PORT,
 64)
 65from opentelemetry.trace import SpanKind
 66from plain.runtime import settings
 67from plain.utils.otel import format_exception_type
 68
 69# The public API of this module — everything else is instrumentation
 70# internals wired up by plain.postgres itself.
 71__all__ = ["suppress_db_tracing"]
 72
 73# Use a stable string key so OpenTelemetry context APIs receive the expected type.
 74_SUPPRESS_KEY = "plain.postgres.suppress_db_tracing"
 75
 76try:
 77    _package_version = importlib.metadata.version("plain.postgres")
 78except importlib.metadata.PackageNotFoundError:
 79    _package_version = "dev"
 80
 81tracer = trace.get_tracer("plain.postgres", _package_version)
 82
 83meter = metrics.get_meter("plain.postgres", version=_package_version)
 84query_duration_histogram = meter.create_histogram(
 85    name=DB_CLIENT_OPERATION_DURATION,
 86    unit="s",
 87    description="Duration of database client operations.",
 88)
 89returned_rows_histogram = meter.create_histogram(
 90    name=DB_CLIENT_RESPONSE_RETURNED_ROWS,
 91    unit="{row}",
 92    description="Number of rows returned by the operation.",
 93)
 94connection_wait_time_histogram = meter.create_histogram(
 95    name=DB_CLIENT_CONNECTION_WAIT_TIME,
 96    unit="s",
 97    description="The time it took to obtain an open connection from the pool.",
 98)
 99connection_use_time_histogram = meter.create_histogram(
100    name=DB_CLIENT_CONNECTION_USE_TIME,
101    unit="s",
102    description="The time between borrowing a connection and returning it to the pool.",
103)
104connection_timeouts_counter = meter.create_counter(
105    name=DB_CLIENT_CONNECTION_TIMEOUTS,
106    unit="{timeout}",
107    description="The number of connection timeouts that have occurred trying to obtain a connection from the pool.",
108)
109
110# WeakKeyDictionary prevents leaks if a conn is GC'd without explicit release().
111_use_start: weakref.WeakKeyDictionary[PsycopgConnection[Any], float] = (
112    weakref.WeakKeyDictionary()
113)
114
115DB_SYSTEM = DbSystemNameValues.POSTGRESQL.value
116
117
118def record_connection_acquire(
119    pool_name: str,
120    conn: PsycopgConnection[Any],
121    wait_seconds: float,
122    checkout_time: float,
123) -> None:
124    connection_wait_time_histogram.record(
125        wait_seconds, {DB_CLIENT_CONNECTION_POOL_NAME: pool_name}
126    )
127    _use_start[conn] = checkout_time
128
129
130def record_connection_release(
131    pool_name: str, conn: PsycopgConnection[Any], return_time: float
132) -> None:
133    start = _use_start.pop(conn, None)
134    if start is None:
135        return
136    connection_use_time_histogram.record(
137        return_time - start, {DB_CLIENT_CONNECTION_POOL_NAME: pool_name}
138    )
139
140
141def record_connection_timeout(pool_name: str) -> None:
142    connection_timeouts_counter.add(1, {DB_CLIENT_CONNECTION_POOL_NAME: pool_name})
143
144
145def register_pool_observables(pool_source: PoolSource) -> None:
146    """Register observable gauges that read `pool.get_stats()` at collection time.
147
148    Safe to call multiple times — the OTel SDK keeps one instrument per name.
149    """
150    pool_attrs = {DB_CLIENT_CONNECTION_POOL_NAME: pool_source.name}
151    idle_attrs = {
152        **pool_attrs,
153        DB_CLIENT_CONNECTION_STATE: DbClientConnectionStateValues.IDLE.value,
154    }
155    used_attrs = {
156        **pool_attrs,
157        DB_CLIENT_CONNECTION_STATE: DbClientConnectionStateValues.USED.value,
158    }
159
160    def _count(_options: CallbackOptions) -> list[Observation]:
161        stats = pool_source.get_stats()
162        if stats is None:
163            return []
164        size = stats.get("pool_size", 0)
165        available = stats.get("pool_available", 0)
166        used = max(size - available, 0)
167        return [
168            Observation(used, used_attrs),
169            Observation(available, idle_attrs),
170        ]
171
172    def _single(stats_key: str) -> Callable[[CallbackOptions], list[Observation]]:
173        def callback(_options: CallbackOptions) -> list[Observation]:
174            stats = pool_source.get_stats()
175            if stats is None:
176                return []
177            return [Observation(stats.get(stats_key, 0), pool_attrs)]
178
179        return callback
180
181    meter.create_observable_up_down_counter(
182        name=DB_CLIENT_CONNECTION_COUNT,
183        unit="{connection}",
184        description="The number of connections that are currently in state described by the state attribute.",
185        callbacks=[_count],
186    )
187    for name, unit, description, stats_key in (
188        (
189            DB_CLIENT_CONNECTION_MAX,
190            "{connection}",
191            "The maximum number of open connections allowed.",
192            "pool_max",
193        ),
194        (
195            DB_CLIENT_CONNECTION_IDLE_MIN,
196            "{connection}",
197            "The minimum number of idle open connections allowed.",
198            "pool_min",
199        ),
200        (
201            DB_CLIENT_CONNECTION_IDLE_MAX,
202            "{connection}",
203            "The maximum number of idle open connections allowed.",
204            "pool_max",
205        ),
206        (
207            DB_CLIENT_CONNECTION_PENDING_REQUESTS,
208            "{request}",
209            "The number of current pending requests for an open connection.",
210            "requests_waiting",
211        ),
212    ):
213        meter.create_observable_up_down_counter(
214            name=name,
215            unit=unit,
216            description=description,
217            callbacks=[_single(stats_key)],
218        )
219
220
221def extract_operation_and_target(sql: str) -> tuple[str, str | None, str | None]:
222    """Extract operation, table name, and collection from SQL.
223
224    Returns: (operation, summary, collection_name)
225    """
226    sql_upper = sql.upper().strip()
227
228    # Strip leading parentheses (e.g. UNION queries: "(SELECT ... UNION ...)")
229    operation = sql_upper.lstrip("(").split()[0] if sql_upper else "UNKNOWN"
230
231    # Pattern to match quoted and unquoted identifiers
232    # Matches: "quoted" (PostgreSQL), unquoted.name
233    identifier_pattern = r'("([^"]+)"|([\w.]+))'
234
235    # Map operations to the SQL keyword that precedes the table name.
236    keyword_by_operation = {
237        "SELECT": "FROM",
238        "DELETE": "FROM",
239        "INSERT": "INTO",
240        "UPDATE": "UPDATE",
241    }
242
243    # Extract table/collection name based on operation
244    collection_name = None
245    summary = operation
246
247    keyword = keyword_by_operation.get(operation)
248    if keyword:
249        match = re.search(rf"{keyword}\s+{identifier_pattern}", sql, re.IGNORECASE)
250        if match:
251            collection_name = _clean_identifier(match.group(1))
252            summary = f"{operation} {collection_name}"
253
254    # Detect UNION queries
255    if " UNION " in sql_upper and summary:
256        summary = f"{summary} UNION"
257
258    return operation, summary, collection_name
259
260
261def _clean_identifier(identifier: str) -> str:
262    """Remove quotes from SQL identifiers."""
263    if identifier.startswith('"') and identifier.endswith('"'):
264        return identifier[1:-1]
265    return identifier
266
267
268@contextmanager
269def db_span(
270    db: DatabaseConnection,
271    sql: Any,
272    *,
273    many: bool = False,
274    params: Any = None,
275    row_count_provider: Callable[[], int] | None = None,
276) -> Generator[Span | None]:
277    """Open an OpenTelemetry CLIENT span for a database query.
278
279    All common attributes (`db.*`, `network.*`, `server.*`, etc.) are set
280    automatically. Follows OpenTelemetry semantic conventions for database
281    instrumentation.
282
283    If `row_count_provider` is given, `db.client.response.returned_rows` is
284    recorded for SELECT operations using its return value (callable so the
285    final count is read after streaming consumers finish iterating).
286    """
287
288    # Fast-exit if instrumentation suppression flag set in context.
289    if otel_context.get_value(_SUPPRESS_KEY):
290        yield None
291        return
292
293    sql = str(sql)  # Ensure SQL is a string for span attributes.
294
295    # Extract operation and target information
296    operation, summary, collection_name = extract_operation_and_target(sql)
297
298    if many:
299        summary = f"{summary} many"
300
301    # Span name follows semantic conventions: {target} or {db.operation.name} {target}
302    if summary:
303        span_name = summary[:255]
304    else:
305        span_name = operation
306
307    # Single settings_dict read — the property delegates to source.config.
308    cfg = db.settings_dict
309
310    # Cheap attributes are passed at span creation so attribute-aware
311    # samplers can see them in should_sample(). Expensive ones (the
312    # per-query stack walk, DEBUG params) are added after, only when the
313    # span actually records.
314    attrs: dict[str, Any] = {
315        DB_SYSTEM_NAME: DB_SYSTEM,
316        DB_NAMESPACE: cfg.get("DATABASE"),
317        DB_QUERY_TEXT: sql,  # Already parameterized from Django/Plain
318        DB_QUERY_SUMMARY: summary,
319        DB_OPERATION_NAME: operation,
320    }
321
322    # Add collection name if detected
323    if collection_name:
324        attrs[DB_COLLECTION_NAME] = collection_name
325
326    # Server/network endpoint. `server.*` is the primary pair per current
327    # semconv; `network.peer.*` is recommended supplementary.
328    if host := cfg.get("HOST"):
329        attrs[SERVER_ADDRESS] = host
330        attrs[NETWORK_PEER_ADDRESS] = host
331
332    if port := cfg.get("PORT"):
333        try:
334            port_int = int(port)
335        except (TypeError, ValueError):
336            pass
337        else:
338            attrs[SERVER_PORT] = port_int
339            attrs[NETWORK_PEER_PORT] = port_int
340
341    with tracer.start_as_current_span(
342        span_name, kind=SpanKind.CLIENT, attributes=attrs
343    ) as span:
344        if span.is_recording():
345            expensive_attrs = _get_code_attributes()
346
347            # Add query parameters as attributes when DEBUG is True
348            if settings.DEBUG and params is not None:
349                # Convert params to appropriate format based on type
350                if isinstance(params, dict):
351                    # Dictionary params (e.g., for named placeholders)
352                    for key, value in params.items():
353                        expensive_attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.{key}"] = str(
354                            value
355                        )
356                elif isinstance(params, list | tuple):
357                    # Sequential params (e.g., for %s or ? placeholders)
358                    for i, value in enumerate(params):
359                        expensive_attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.{i + 1}"] = str(
360                            value
361                        )
362                else:
363                    # Single param (rare but possible)
364                    expensive_attrs[f"{DB_QUERY_PARAMETER_TEMPLATE}.1"] = str(params)
365
366            span.set_attributes(expensive_attrs)
367
368        start = time.perf_counter()
369        try:
370            yield span
371        except Exception as exc:
372            # record_exception + set_status(ERROR) handled by
373            # start_as_current_span when the exception propagates out.
374            if span.is_recording():
375                span.set_attribute(ERROR_TYPE, format_exception_type(exc))
376            raise
377        duration_s = time.perf_counter() - start
378
379        metric_attrs: dict[str, str] = {
380            DB_SYSTEM_NAME: DB_SYSTEM,
381            DB_OPERATION_NAME: operation,
382        }
383        if collection_name:
384            metric_attrs[DB_COLLECTION_NAME] = collection_name
385        query_duration_histogram.record(duration_s, metric_attrs)
386
387        # Scope returned_rows to SELECT; rowcount for INSERT/UPDATE/DELETE
388        # is rows-affected, which is a different semantic.
389        if row_count_provider is not None and operation == "SELECT":
390            count = row_count_provider()
391            if count >= 0:
392                returned_rows_histogram.record(count, metric_attrs)
393
394
395@contextmanager
396def suppress_db_tracing() -> Generator[None]:
397    token = otel_context.attach(otel_context.set_value(_SUPPRESS_KEY, True))
398    try:
399        yield
400    finally:
401        otel_context.detach(token)
402
403
404def _is_internal_frame(frame: traceback.FrameSummary) -> bool:
405    """Return True if the frame is internal to plain.postgres or contextlib."""
406    filepath = frame.filename
407    if not filepath:
408        return True
409    if "/plain/postgres/" in filepath:
410        return True
411    return filepath.endswith("contextlib.py")
412
413
414def _get_code_attributes() -> dict[str, Any]:
415    """Extract code context attributes for the current database query.
416
417    Returns a dict of OpenTelemetry code attributes.
418    """
419    stack = traceback.extract_stack()
420
421    # Find the first user code frame (outermost non-internal frame from the top of the call stack)
422    for frame in reversed(stack):
423        if _is_internal_frame(frame):
424            continue
425
426        attrs: dict[str, Any] = {
427            CODE_FILE_PATH: frame.filename,
428        }
429        if frame.lineno:
430            attrs[CODE_LINE_NUMBER] = frame.lineno
431        if frame.name:
432            attrs[CODE_FUNCTION_NAME] = frame.name
433        if frame.colno:
434            attrs[CODE_COLUMN_NUMBER] = frame.colno
435
436        # Add full stack trace only in DEBUG mode (expensive)
437        if settings.DEBUG:
438            filtered_stack = [f for f in stack if not _is_internal_frame(f)]
439            attrs[CODE_STACKTRACE] = "".join(traceback.format_list(filtered_stack))
440
441        return attrs
442
443    return {}