async-orm-kit 1.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (92) hide show
  1. async_orm_kit/__init__.py +151 -0
  2. async_orm_kit/_logging.py +58 -0
  3. async_orm_kit/_version.py +1 -0
  4. async_orm_kit/config.py +31 -0
  5. async_orm_kit/exceptions.py +353 -0
  6. async_orm_kit/integrations/__init__.py +15 -0
  7. async_orm_kit/integrations/alembic.py +56 -0
  8. async_orm_kit/integrations/fastapi.py +134 -0
  9. async_orm_kit/integrations/storage.py +213 -0
  10. async_orm_kit/manager/__init__.py +3 -0
  11. async_orm_kit/manager/manager.py +568 -0
  12. async_orm_kit/models/__init__.py +3 -0
  13. async_orm_kit/models/base.py +568 -0
  14. async_orm_kit/models/fields/__init__.py +23 -0
  15. async_orm_kit/models/fields/exceptions.py +10 -0
  16. async_orm_kit/models/fields/typed/__init__.py +11 -0
  17. async_orm_kit/models/fields/typed/base.py +155 -0
  18. async_orm_kit/models/fields/typed/file.py +697 -0
  19. async_orm_kit/models/fields/typed/string.py +38 -0
  20. async_orm_kit/models/fields/validators.py +110 -0
  21. async_orm_kit/models/meta.py +181 -0
  22. async_orm_kit/models/mixins.py +19 -0
  23. async_orm_kit/observers/__init__.py +4 -0
  24. async_orm_kit/observers/mixin.py +26 -0
  25. async_orm_kit/observers/model_observer.py +20 -0
  26. async_orm_kit/observers/queryset_observer.py +38 -0
  27. async_orm_kit/py.typed +0 -0
  28. async_orm_kit/query/__init__.py +0 -0
  29. async_orm_kit/query/aggregations/__init__.py +8 -0
  30. async_orm_kit/query/aggregations/avg.py +71 -0
  31. async_orm_kit/query/aggregations/base.py +317 -0
  32. async_orm_kit/query/aggregations/count.py +115 -0
  33. async_orm_kit/query/aggregations/max.py +76 -0
  34. async_orm_kit/query/aggregations/min.py +76 -0
  35. async_orm_kit/query/aggregations/sum.py +73 -0
  36. async_orm_kit/query/compiler.py +574 -0
  37. async_orm_kit/query/debugger.py +269 -0
  38. async_orm_kit/query/expressions/__init__.py +62 -0
  39. async_orm_kit/query/expressions/analysis.py +56 -0
  40. async_orm_kit/query/expressions/case.py +321 -0
  41. async_orm_kit/query/expressions/collections.py +251 -0
  42. async_orm_kit/query/expressions/f.py +362 -0
  43. async_orm_kit/query/expressions/functions.py +202 -0
  44. async_orm_kit/query/expressions/subquery.py +206 -0
  45. async_orm_kit/query/expressions/value.py +101 -0
  46. async_orm_kit/query/expressions/window.py +95 -0
  47. async_orm_kit/query/lookups/__init__.py +121 -0
  48. async_orm_kit/query/lookups/conditions.py +315 -0
  49. async_orm_kit/query/lookups/correlated.py +86 -0
  50. async_orm_kit/query/lookups/json_paths.py +167 -0
  51. async_orm_kit/query/lookups/parser.py +276 -0
  52. async_orm_kit/query/lookups/resolver/__init__.py +5 -0
  53. async_orm_kit/query/lookups/resolver/context.py +150 -0
  54. async_orm_kit/query/lookups/resolver/resolver.py +572 -0
  55. async_orm_kit/query/lookups/resolver/types.py +150 -0
  56. async_orm_kit/query/mixins/__init__.py +0 -0
  57. async_orm_kit/query/mixins/combinable.py +376 -0
  58. async_orm_kit/query/mixins/orderable.py +208 -0
  59. async_orm_kit/query/observability.py +346 -0
  60. async_orm_kit/query/planning/__init__.py +0 -0
  61. async_orm_kit/query/planning/directives.py +117 -0
  62. async_orm_kit/query/planning/factory.py +40 -0
  63. async_orm_kit/query/planning/passes/__init__.py +15 -0
  64. async_orm_kit/query/planning/passes/base.py +47 -0
  65. async_orm_kit/query/planning/passes/cartesian_detection.py +74 -0
  66. async_orm_kit/query/planning/passes/depth_splitting.py +78 -0
  67. async_orm_kit/query/planning/passes/join_promotion.py +81 -0
  68. async_orm_kit/query/planning/passes/path_conflicts.py +103 -0
  69. async_orm_kit/query/planning/passes/redundant_path.py +67 -0
  70. async_orm_kit/query/planning/passes/strategy_selection.py +49 -0
  71. async_orm_kit/query/planning/planner.py +160 -0
  72. async_orm_kit/query/planning/protocols.py +78 -0
  73. async_orm_kit/query/postgres/__init__.py +0 -0
  74. async_orm_kit/query/postgres/optimizer.py +189 -0
  75. async_orm_kit/query/q.py +596 -0
  76. async_orm_kit/query/relations/__init__.py +0 -0
  77. async_orm_kit/query/relations/enums.py +16 -0
  78. async_orm_kit/query/relations/graph.py +908 -0
  79. async_orm_kit/query/relations/node.py +421 -0
  80. async_orm_kit/queryset/__init__.py +4 -0
  81. async_orm_kit/queryset/builder.py +1259 -0
  82. async_orm_kit/queryset/compound.py +379 -0
  83. async_orm_kit/queryset/executor.py +718 -0
  84. async_orm_kit/queryset/queryset.py +1726 -0
  85. async_orm_kit/queryset/state.py +164 -0
  86. async_orm_kit/queryset/validators.py +361 -0
  87. async_orm_kit/queryset/writer.py +755 -0
  88. async_orm_kit/types.py +23 -0
  89. async_orm_kit-1.0.0.dist-info/METADATA +318 -0
  90. async_orm_kit-1.0.0.dist-info/RECORD +92 -0
  91. async_orm_kit-1.0.0.dist-info/WHEEL +4 -0
  92. async_orm_kit-1.0.0.dist-info/licenses/LICENSE +21 -0
@@ -0,0 +1,208 @@
1
+ from typing import Any
2
+
3
+ from async_orm_kit.query.lookups.resolver import ResolverContext
4
+
5
+
6
+ class OrderedExpression:
7
+ """Wrapper that adds sort direction to ORM expressions for use in ``order_by()``.
8
+
9
+ Encapsulates any ORM expression (field reference, computed expression, aggregate,
10
+ or conditional) with an ascending or descending sort direction. Created by
11
+ calling ``.asc()`` or ``.desc()`` on any :class:`OrderableMixin` subclass.
12
+
13
+ During query compilation, the wrapped expression is resolved to a SQLAlchemy
14
+ column/expression, then the appropriate ``.asc()`` or ``.desc()`` method is
15
+ applied to generate the final ORDER BY clause element.
16
+
17
+ Args:
18
+ expression: Any ORM expression with a ``resolve(field_resolver, context)``
19
+ method. Typically a :class:`F`, :class:`FExpression`, :class:`Aggregate`,
20
+ or :class:`Case` instance.
21
+ ascending: If ``True``, sort in ascending order (smallest to largest).
22
+ If ``False``, sort in descending order (largest to smallest).Defaults to ``True``.Attributes:
23
+ expression: The wrapped ORM expression to be sorted.
24
+ ascending (bool): Sort direction flag.
25
+
26
+ Example:
27
+ >>> from async_orm_kit.query.expressions.f import F
28
+ >>> from async_orm_kit.query.aggregations import Sum
29
+ >>> from async_orm_kit.query.expressions.case import Case, When
30
+ >>> from async_orm_kit.query.expressions.value import Value
31
+ >>>
32
+ >>> # Simple field ordering
33
+ >>> Product.objects.order_by(F("price").desc())
34
+ >>> # SQL: ORDER BY "price" DESC
35
+ >>>
36
+ >>> # Computed expression ordering
37
+ >>> Product.objects.order_by((F("price") * F("quantity")).asc())
38
+ >>> # SQL: ORDER BY ("price" * "quantity") ASC
39
+ >>>
40
+ >>> # Aggregate ordering (requires annotation)
41
+ >>> Product.objects.annotate(
42
+ ... total_sales=Sum("order_items__quantity")
43
+ ... ).order_by(F("total_sales").desc())
44
+ >>> # SQL: ORDER BY SUM("order_items"."quantity") DESC
45
+ >>>
46
+ >>> # Conditional ordering
47
+ >>> priority = Case(
48
+ ... When(status="urgent", then=Value(0)),
49
+ ... When(status="high", then=Value(1)),
50
+ ... default=Value(2)
51
+ ... )
52
+ >>> Task.objects.order_by(priority.asc(), F("created_at").desc())
53
+ >>> # SQL: ORDER BY CASE ... END ASC, "created_at" DESC
54
+
55
+ Note:
56
+ - The wrapped expression is **not resolved** until query compilation time,
57
+ allowing the same ``OrderedExpression`` instance to be reused across
58
+ multiple queries.
59
+ - Multiple ``OrderedExpression`` instances can be passed to ``order_by()``
60
+ to create multi-column sorts with mixed directions.
61
+ - Database NULL handling varies by backend and can be controlled with
62
+ ``NULLS FIRST`` / ``NULLS LAST`` (not currently supported by this class,
63
+ but available via raw SQLAlchemy if needed).
64
+
65
+ See Also:
66
+ - :class:`OrderableMixin` for the ``.asc()`` / ``.desc()`` factory methods.
67
+ - :meth:`QuerySet.order_by` for usage in query construction.
68
+ """
69
+
70
+ def __init__(self, expression: Any, ascending: bool = True):
71
+ self.expression = expression
72
+ self.ascending = ascending
73
+
74
+ def resolve(self, field_resolver, context: ResolverContext = ResolverContext()):
75
+ """Resolve the wrapped expression and apply sort direction.Delegates resolution to the wrapped expression's ``resolve()`` method,
76
+ then applies SQLAlchemy's ``.asc()`` or ``.desc()`` method to the
77
+ resulting column/expression to produce the final ORDER BY element.
78
+
79
+ Args:
80
+ field_resolver: FieldResolver instance for resolving field paths
81
+ and managing JOIN registration.
82
+ context: Resolution context for JOIN optimization. Defaults to
83
+ empty context (no special flags).
84
+
85
+ Returns:
86
+ SQLAlchemy column expression with ``.asc()`` or ``.desc()`` applied,
87
+ ready for inclusion in an ORDER BY clause.
88
+
89
+ Example:
90
+ >>> resolver = FieldResolver(Product, join_graph)
91
+ >>> ordered = F("price").desc()
92
+ >>> sql_expr = ordered.resolve(resolver)
93
+ >>> # Returns: Product.price.desc() (SQLAlchemy UnaryExpression)
94
+
95
+ Note:
96
+ The ``context`` parameter is propagated to the wrapped expression's
97
+ resolution, ensuring consistent JOIN strategy if the expression
98
+ references relationships.
99
+ """
100
+ resolved = self.expression.resolve(field_resolver, context)
101
+ if self.ascending:
102
+ return resolved.asc()
103
+ return resolved.desc()
104
+
105
+ def __repr__(self) -> str:
106
+ """Return developer-friendly string representation.
107
+
108
+ Returns:
109
+ String showing the wrapped expression and sort direction.
110
+
111
+ Example:
112
+ >>> repr(F("price").desc())
113
+ >>> # Returns: "OrderedExpression(F('price'), DESC)"
114
+ """
115
+ direction = "ASC" if self.ascending else "DESC"
116
+ return f"OrderedExpression({self.expression!r}, {direction})"
117
+
118
+
119
+ class OrderableMixin:
120
+ """Mixin that adds ``.asc()`` and ``.desc()`` methods for sort direction control.
121
+
122
+ Enables any ORM expression class to produce :class:`OrderedExpression` instances
123
+ for use in ``order_by()`` clauses. Classes inheriting this mixin can be sorted
124
+ in ascending or descending order by calling the provided methods.
125
+
126
+ Classes Using This Mixin:
127
+ - :class:`F` — field references (``F("created_at").desc()``)
128
+ - :class:`FExpression` — computed expressions (``(F("a") + F("b")).asc()``)
129
+ - :class:`Aggregate` — aggregate results (``Sum("amount").desc()``)
130
+ - :class:`Case` — conditional expressions (``Case(...).asc()``)
131
+ - :class:`Value` — literal values (``Value(1).desc()`` — rare but valid)
132
+
133
+ Example:
134
+ >>> from async_orm_kit.query.expressions.f import F
135
+ >>> from async_orm_kit.query.aggregations import Count
136
+ >>> # Ascending sort (default direction)
137
+ >>> Product.objects.order_by(F("name").asc())
138
+ >>> # SQL: ORDER BY "name" ASC
139
+ >>>
140
+ >>> # Descending sort
141
+ >>> Product.objects.order_by(F("price").desc())
142
+ >>> # SQL: ORDER BY "price" DESC
143
+ >>>
144
+ >>> # Multi-column sort with mixed directions
145
+ >>> Product.objects.order_by(
146
+ ... F("category").asc(),
147
+ ... F("price").desc(),
148
+ ... F("name").asc()
149
+ ... )
150
+ >>> # SQL: ORDER BY "category" ASC, "price" DESC, "name" ASC
151
+ >>>
152
+ >>> # Sorting by computed expression
153
+ >>> Product.objects.order_by((F("price") * F("quantity")).desc())
154
+ >>> # SQL: ORDER BY ("price" * "quantity") DESC
155
+ >>>
156
+ >>> # Sorting by aggregate (requires annotation)
157
+ >>> Category.objects.annotate(
158
+ ... product_count=Count("products")
159
+ ... ).order_by(F("product_count").desc())
160
+ >>> # SQL: ORDER BY COUNT("products"."id") DESC
161
+
162
+ Note:
163
+ - Both methods return new :class:`OrderedExpression` instances; the
164
+ original expression object is **never mutated** (immutable pattern).
165
+ - The ``__slots__ = ()`` declaration prevents instance attribute creation,
166
+ ensuring this mixin adds no memory overhead to inheriting classes.
167
+ - Calling ``.asc()`` or ``.desc()`` multiple times on the same expression
168
+ creates independent ``OrderedExpression`` wrappers — the last one used
169
+ in ``order_by()`` determines the actual sort direction.
170
+
171
+ See Also:
172
+ - :class:`OrderedExpression` for the wrapper class structure.
173
+ - :meth:`QuerySet.order_by` for usage in query construction.
174
+ """
175
+
176
+ __slots__ = ()
177
+
178
+ def asc(self) -> OrderedExpression:
179
+ """Create an ascending sort wrapper for this expression.
180
+
181
+ Returns:
182
+ OrderedExpression with ``ascending=True``, suitable for passing
183
+ to ``order_by()``.
184
+
185
+ Example:
186
+ >>> F("price").asc()
187
+ >>> # Returns: OrderedExpression(F('price'), ASC)
188
+ >>>
189
+ >>> Product.objects.order_by(F("created_at").asc())
190
+ >>> # SQL: ORDER BY "created_at" ASC
191
+ """
192
+ return OrderedExpression(self, True)
193
+
194
+ def desc(self) -> OrderedExpression:
195
+ """Create a descending sort wrapper for this expression.
196
+
197
+ Returns:
198
+ OrderedExpression with ``ascending=False``, suitable for passing
199
+ to ``order_by()``.
200
+
201
+ Example:
202
+ >>> F("price").desc()
203
+ >>> # Returns: OrderedExpression(F('price'), DESC)
204
+ >>>
205
+ >>> Product.objects.order_by(F("price").desc())
206
+ >>> # SQL: ORDER BY "price" DESC
207
+ """
208
+ return OrderedExpression(self, False)
@@ -0,0 +1,346 @@
1
+ """Opt-in SQLAlchemy execution diagnostics, scoped to the current async task.
2
+
3
+ The profiler measures DBAPI cursor-execution time, not complete QuerySet
4
+ materialization, network connection acquisition, or request latency. It works
5
+ with SQLAlchemy ``Engine`` and ``AsyncEngine`` without wrapping/monkey-patching
6
+ SQLAlchemy internals and introduces no new mandatory dependencies.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import re
12
+ import threading
13
+ import time
14
+ from collections import Counter
15
+ from collections.abc import AsyncIterator, Callable, Iterator
16
+ from contextlib import asynccontextmanager, contextmanager
17
+ from contextvars import ContextVar
18
+ from dataclasses import dataclass, field
19
+ from typing import Any
20
+
21
+ from sqlalchemy import event
22
+ from sqlalchemy.engine import Engine
23
+ from sqlalchemy.ext.asyncio import AsyncEngine
24
+
25
+ from async_orm_kit._logging import get_logger
26
+
27
+ logger = get_logger(__name__)
28
+
29
+ # Only SQL templates are retained. Redact SQL literals and comments even when
30
+ # raw SQL is executed through exec_driver_sql(). Parameter values are never read.
31
+ _SQL_LITERAL = re.compile(r"'(?:''|[^'])*'", re.DOTALL)
32
+ _SQL_BLOCK_COMMENT = re.compile(r"/\*.*?\*/", re.DOTALL)
33
+ _SQL_LINE_COMMENT = re.compile(r"--[^\n]*")
34
+ _SQL_NUMBER = re.compile(r"(?<![\w:$])\b\d+(?:\.\d+)?\b(?![\w])")
35
+ _SQL_WHITESPACE = re.compile(r"\s+")
36
+
37
+
38
+ def sanitized_sql(sql: str, *, max_length: int = 2048) -> str:
39
+ """Return a bounded, best-effort redacted statement (never bind values).
40
+
41
+ For absolute secrecy in production, do not enable SQL recording: arbitrary
42
+ driver SQL and SQL identifiers may contain confidential information.
43
+ """
44
+ text = _SQL_LITERAL.sub("'?'", sql)
45
+ text = _SQL_BLOCK_COMMENT.sub(" ", text)
46
+ text = _SQL_LINE_COMMENT.sub(" ", text)
47
+ text = _SQL_NUMBER.sub("?", text)
48
+ text = _SQL_WHITESPACE.sub(" ", text).strip()
49
+ return text[:max_length]
50
+
51
+
52
+ def sql_fingerprint(sql: str) -> str:
53
+ """Normalize a statement for *heuristic* repetition detection."""
54
+ return sanitized_sql(sql, max_length=4096)
55
+
56
+
57
+ @dataclass(frozen=True, slots=True)
58
+ class QueryRecord:
59
+ """One database cursor execution, without SQL bind parameters."""
60
+
61
+ sql: str
62
+ duration_ms: float
63
+ timestamp: float
64
+ row_count: int | None
65
+ success: bool
66
+ error_type: str | None = None
67
+
68
+
69
+ @dataclass(slots=True)
70
+ class QueryProfile:
71
+ """A bounded task-local capture, including totals for evicted records."""
72
+
73
+ max_entries: int
74
+ slow_threshold_ms: float
75
+ records: list[QueryRecord] = field(default_factory=list)
76
+ total_queries: int = 0
77
+ total_duration_ms: float = 0.0
78
+ slow_queries: int = 0
79
+ failed_queries: int = 0
80
+ dropped_entries: int = 0
81
+ untracked_patterns: int = 0
82
+ patterns: Counter[str] = field(default_factory=Counter)
83
+ active: bool = True
84
+ _slowest: QueryRecord | None = field(default=None, repr=False)
85
+
86
+ def _record(self, record: QueryRecord) -> None:
87
+ if not self.active:
88
+ return
89
+ self.total_queries += 1
90
+ self.total_duration_ms += record.duration_ms
91
+ self.slow_queries += record.duration_ms >= self.slow_threshold_ms
92
+ self.failed_queries += not record.success
93
+ if self._slowest is None or record.duration_ms > self._slowest.duration_ms:
94
+ self._slowest = record
95
+ if record.success:
96
+ fingerprint = sql_fingerprint(record.sql)
97
+ if fingerprint in self.patterns or len(self.patterns) < self.max_entries:
98
+ self.patterns[fingerprint] += 1
99
+ else:
100
+ # Query patterns are bounded too; distinct raw SQL must not
101
+ # grow the profile indefinitely in a long-running request.
102
+ self.untracked_patterns += 1
103
+ if len(self.records) >= self.max_entries:
104
+ self.records.pop(0)
105
+ self.dropped_entries += 1
106
+ self.records.append(record)
107
+
108
+ def stats(self) -> dict[str, Any]:
109
+ """Return full execution totals, even if older records were evicted."""
110
+ return {
111
+ "total_queries": self.total_queries,
112
+ "total_duration_ms": round(self.total_duration_ms, 3),
113
+ "avg_duration_ms": round(self.total_duration_ms / self.total_queries, 3)
114
+ if self.total_queries
115
+ else 0.0,
116
+ "slow_queries": self.slow_queries,
117
+ "failed_queries": self.failed_queries,
118
+ "dropped_entries": self.dropped_entries,
119
+ "untracked_patterns": self.untracked_patterns,
120
+ "slowest_query": (
121
+ {"sql": self._slowest.sql, "duration_ms": self._slowest.duration_ms}
122
+ if self._slowest is not None
123
+ else None
124
+ ),
125
+ }
126
+
127
+ def n_plus_one_report(self, *, threshold: int = 10) -> dict[str, Any]:
128
+ """Report suspected repeated SELECT-by-key patterns, not guaranteed N+1.
129
+
130
+ Repeated WHERE queries are only a heuristic; deliberate polling and
131
+ bulk jobs may also match. This does not trace object relationships.
132
+ """
133
+ if (
134
+ isinstance(threshold, bool)
135
+ or not isinstance(threshold, int)
136
+ or threshold < 2
137
+ ):
138
+ raise ValueError("threshold must be an integer >= 2")
139
+ suspects: list[dict[str, Any]] = [
140
+ {
141
+ "pattern": sql[:200],
142
+ "count": count,
143
+ "severity": "high" if count > 100 else "medium",
144
+ }
145
+ for sql, count in self.patterns.items()
146
+ if count >= threshold
147
+ and sql.lstrip().upper().startswith("SELECT ")
148
+ and " WHERE " in f" {sql.upper()} "
149
+ ]
150
+ suspects.sort(key=lambda entry: (-entry["count"], entry["pattern"]))
151
+ return {
152
+ "total_queries": self.total_queries,
153
+ "unique_patterns": len(self.patterns),
154
+ "untracked_patterns": self.untracked_patterns,
155
+ "potential_n_plus_one": suspects,
156
+ }
157
+
158
+
159
+ class QueryProfiler:
160
+ """Capture cursor executions from a SQLAlchemy Engine or AsyncEngine.
161
+
162
+ A profiler's ``capture()`` is scoped via ``ContextVar`` to the current
163
+ execution context; captures in concurrent async tasks do not mix. Nested
164
+ captures both receive the query. Lifecycle registration is reference-counted
165
+ so one exiting request cannot detach another active request's listener.
166
+
167
+ Example::
168
+
169
+ profiler = QueryProfiler(async_engine)
170
+ with profiler.capture() as profile:
171
+ results = await session.execute(statement)
172
+ print(profile.stats(), profile.n_plus_one_report(threshold=3))
173
+
174
+ ``capture()`` automatically registers and removes listeners. Use ``attach``
175
+ and ``detach`` for persistent monitoring, though uncaptured executions are
176
+ ignored unless an ``on_query`` callback is configured.
177
+ """
178
+
179
+ def __init__(
180
+ self,
181
+ engine: Engine | AsyncEngine,
182
+ *,
183
+ max_entries: int = 1000,
184
+ slow_threshold_ms: float = 100.0,
185
+ on_query: Callable[[QueryRecord], None] | None = None,
186
+ ) -> None:
187
+ sync_engine = engine.sync_engine if isinstance(engine, AsyncEngine) else engine
188
+ if not isinstance(sync_engine, Engine):
189
+ raise TypeError("engine must be a SQLAlchemy Engine or AsyncEngine")
190
+ if (
191
+ isinstance(max_entries, bool)
192
+ or not isinstance(max_entries, int)
193
+ or max_entries < 1
194
+ ):
195
+ raise ValueError("max_entries must be an integer >= 1")
196
+ if slow_threshold_ms < 0:
197
+ raise ValueError("slow_threshold_ms must be nonnegative")
198
+ if on_query is not None and not callable(on_query):
199
+ raise TypeError("on_query must be callable")
200
+ self._engine = sync_engine
201
+ self._max_entries = max_entries
202
+ self._slow_threshold_ms = float(slow_threshold_ms)
203
+ self._on_query = on_query
204
+ self._captures: ContextVar[tuple[QueryProfile, ...]] = ContextVar(
205
+ f"ormkit_query_profiles_{id(self)}", default=()
206
+ )
207
+ self._registered = False
208
+ self._permanent = False
209
+ self._capture_count = 0
210
+ self._lock = threading.RLock()
211
+ self._clock_key = f"_ormkit_query_start_{id(self)}"
212
+ # Hold exact callback objects for event.remove(), including on Python
213
+ # versions that create a fresh bound method object on attribute access.
214
+ self._before_cb = self._before_cursor_execute
215
+ self._after_cb = self._after_cursor_execute
216
+ self._error_cb = self._handle_error
217
+
218
+ def _install(self) -> None:
219
+ if self._registered:
220
+ return
221
+ event.listen(self._engine, "before_cursor_execute", self._before_cb)
222
+ event.listen(self._engine, "after_cursor_execute", self._after_cb)
223
+ event.listen(self._engine, "handle_error", self._error_cb)
224
+ self._registered = True
225
+
226
+ def _uninstall(self) -> None:
227
+ if not self._registered:
228
+ return
229
+ event.remove(self._engine, "before_cursor_execute", self._before_cb)
230
+ event.remove(self._engine, "after_cursor_execute", self._after_cb)
231
+ event.remove(self._engine, "handle_error", self._error_cb)
232
+ self._registered = False
233
+
234
+ def attach(self) -> None:
235
+ """Keep listeners installed until ``detach()`` is called."""
236
+ with self._lock:
237
+ self._permanent = True
238
+ self._install()
239
+
240
+ def detach(self) -> None:
241
+ """Remove persistent monitoring without disrupting active captures."""
242
+ with self._lock:
243
+ self._permanent = False
244
+ if not self._capture_count:
245
+ self._uninstall()
246
+
247
+ @contextmanager
248
+ def capture(self) -> Iterator[QueryProfile]:
249
+ """Capture queries in this task; usable across ``await`` boundaries."""
250
+ profile = QueryProfile(
251
+ max_entries=self._max_entries,
252
+ slow_threshold_ms=self._slow_threshold_ms,
253
+ )
254
+ with self._lock:
255
+ self._install()
256
+ self._capture_count += 1
257
+ token = self._captures.set((*self._captures.get(), profile))
258
+ try:
259
+ yield profile
260
+ finally:
261
+ profile.active = False
262
+ self._captures.reset(token)
263
+ with self._lock:
264
+ self._capture_count -= 1
265
+ if not self._permanent and not self._capture_count:
266
+ self._uninstall()
267
+
268
+ @asynccontextmanager
269
+ async def acapture(self) -> AsyncIterator[QueryProfile]:
270
+ """``async with`` equivalent of ``capture()``."""
271
+ with self.capture() as profile:
272
+ yield profile
273
+
274
+ def _before_cursor_execute(
275
+ self,
276
+ connection: Any,
277
+ cursor: Any,
278
+ statement: str,
279
+ parameters: Any,
280
+ context: Any,
281
+ executemany: bool,
282
+ ) -> None:
283
+ if self._captures.get() or self._on_query is not None:
284
+ setattr(context, self._clock_key, time.perf_counter_ns())
285
+
286
+ def _finish(
287
+ self,
288
+ statement: str,
289
+ context: Any,
290
+ *,
291
+ row_count: int | None = None,
292
+ error_type: str | None = None,
293
+ ) -> None:
294
+ start = getattr(context, self._clock_key, None)
295
+ if start is None:
296
+ return
297
+ # Avoid double counting if an adapter emits more than one callback.
298
+ delattr(context, self._clock_key)
299
+ duration_ms = max(0.0, (time.perf_counter_ns() - start) / 1_000_000)
300
+ record = QueryRecord(
301
+ sql=sanitized_sql(statement),
302
+ duration_ms=duration_ms,
303
+ timestamp=time.time(),
304
+ row_count=row_count if row_count is not None and row_count >= 0 else None,
305
+ success=error_type is None,
306
+ error_type=error_type,
307
+ )
308
+ for profile in self._captures.get():
309
+ profile._record(record)
310
+ if self._on_query is not None:
311
+ try:
312
+ self._on_query(record)
313
+ except Exception:
314
+ # Diagnostics must never change the result of a DB operation.
315
+ logger.exception("Query monitoring callback failed")
316
+
317
+ def _after_cursor_execute(
318
+ self,
319
+ connection: Any,
320
+ cursor: Any,
321
+ statement: str,
322
+ parameters: Any,
323
+ context: Any,
324
+ executemany: bool,
325
+ ) -> None:
326
+ self._finish(statement, context, row_count=getattr(cursor, "rowcount", None))
327
+
328
+ def _handle_error(self, exception_context: Any) -> None:
329
+ context = getattr(exception_context, "execution_context", None)
330
+ if context is None:
331
+ return # failure before cursor execution is outside this profiler
332
+ error = getattr(exception_context, "original_exception", None)
333
+ self._finish(
334
+ getattr(exception_context, "statement", "") or "",
335
+ context,
336
+ error_type=type(error).__name__ if error is not None else "UnknownError",
337
+ )
338
+
339
+
340
+ __all__ = [
341
+ "QueryProfiler",
342
+ "QueryProfile",
343
+ "QueryRecord",
344
+ "sanitized_sql",
345
+ "sql_fingerprint",
346
+ ]
File without changes
@@ -0,0 +1,117 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass, field
4
+ from enum import Enum, auto
5
+
6
+
7
+ class LoadStrategy(Enum):
8
+ JOINEDLOAD = auto()
9
+ SELECTINLOAD = auto()
10
+ SUBQUERYLOAD = auto()
11
+ SKIP = auto()
12
+
13
+
14
+ @dataclass(frozen=True, slots=True)
15
+ class LoadDirective:
16
+ """
17
+ Immutable output of the planning process.
18
+
19
+ One directive per loading path, carrying the chosen strategy
20
+ and a human-readable reason for debugging / EXPLAIN output.
21
+ """
22
+
23
+ path: str
24
+ strategy: LoadStrategy
25
+ reason: str
26
+ original_strategy: LoadStrategy | None = None # before any demotion
27
+ segment_strategies: tuple[LoadStrategy, ...] = ()
28
+
29
+ def __post_init__(self) -> None:
30
+ if not isinstance(self.path, str) or not self.path:
31
+ raise ValueError("LoadDirective.path must be a non-empty string.")
32
+ path_parts = self.path.split("__")
33
+ if any(not part or part.strip() != part for part in path_parts):
34
+ raise ValueError("LoadDirective.path contains an invalid segment.")
35
+ if not isinstance(self.strategy, LoadStrategy):
36
+ raise TypeError("LoadDirective.strategy must be a LoadStrategy.")
37
+ if not isinstance(self.reason, str) or not self.reason:
38
+ raise ValueError("LoadDirective.reason must be a non-empty string.")
39
+ if self.original_strategy is not None and not isinstance(
40
+ self.original_strategy, LoadStrategy
41
+ ):
42
+ raise TypeError("original_strategy must be a LoadStrategy or None.")
43
+ try:
44
+ segment_strategies = tuple(self.segment_strategies)
45
+ except TypeError as exc:
46
+ raise TypeError(
47
+ "segment_strategies must be an iterable of strategies."
48
+ ) from exc
49
+ object.__setattr__(self, "segment_strategies", segment_strategies)
50
+ if segment_strategies:
51
+ if len(segment_strategies) != len(path_parts):
52
+ raise ValueError(
53
+ "segment_strategies must contain one strategy per path segment."
54
+ )
55
+ if any(
56
+ not isinstance(strategy, LoadStrategy)
57
+ for strategy in segment_strategies
58
+ ):
59
+ raise TypeError("segment_strategies must contain LoadStrategy values.")
60
+
61
+ @property
62
+ def was_demoted(self) -> bool:
63
+ return (
64
+ self.original_strategy is not None
65
+ and self.original_strategy != self.strategy
66
+ )
67
+
68
+ def __repr__(self) -> str:
69
+ demoted = " (demoted)" if self.was_demoted else ""
70
+ return (
71
+ f"LoadDirective({self.path!r}, "
72
+ f"{self.strategy.name}{demoted}, "
73
+ f"{self.reason!r})"
74
+ )
75
+
76
+
77
+ @dataclass(frozen=True, slots=True)
78
+ class JoinPromotion:
79
+ """Record of a LEFT → INNER promotion for audit trail."""
80
+
81
+ path: str
82
+ old_type: str
83
+ new_type: str
84
+ reason: str
85
+
86
+
87
+ @dataclass(frozen=True, slots=True)
88
+ class QueryContext:
89
+ """
90
+ Immutable snapshot of query-level facts that planning
91
+ passes need to make decisions.
92
+
93
+ Built ONCE by the QuerySet compiler, then passed through
94
+ read-only. No pass mutates this.
95
+ """
96
+
97
+ has_aggregates: bool = False
98
+ is_distinct: bool = False
99
+ limit: int | None = None
100
+
101
+ # Paths the user explicitly requested
102
+ select_related_paths: frozenset[str] = field(default_factory=frozenset)
103
+ prefetch_paths: frozenset[str] = field(default_factory=frozenset)
104
+
105
+ def __post_init__(self) -> None:
106
+ if type(self.has_aggregates) is not bool or type(self.is_distinct) is not bool:
107
+ raise TypeError("QueryContext flags must be bool values.")
108
+ if self.limit is not None and (type(self.limit) is not int or self.limit < 0):
109
+ raise ValueError(
110
+ "QueryContext.limit must be a non-negative integer or None."
111
+ )
112
+
113
+ for attribute in ("select_related_paths", "prefetch_paths"):
114
+ paths = frozenset(getattr(self, attribute))
115
+ if any(not isinstance(path, str) or not path for path in paths):
116
+ raise ValueError(f"{attribute} must contain non-empty strings.")
117
+ object.__setattr__(self, attribute, paths)