altcodepro-polydb-python 2.3.20__py3-none-any.whl → 2.3.21__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.
- {altcodepro_polydb_python-2.3.20.dist-info → altcodepro_polydb_python-2.3.21.dist-info}/METADATA +1 -1
- {altcodepro_polydb_python-2.3.20.dist-info → altcodepro_polydb_python-2.3.21.dist-info}/RECORD +13 -8
- polydb/__init__.py +2 -0
- polydb/adapters/PostgreSQLAdapter.py +296 -189
- polydb/errors.py +6 -0
- polydb/observability/__init__.py +3 -0
- polydb/observability/logging.py +124 -0
- polydb/services/__init__.py +9 -0
- polydb/services/compliance_service.py +141 -0
- polydb/services/security_service.py +133 -0
- {altcodepro_polydb_python-2.3.20.dist-info → altcodepro_polydb_python-2.3.21.dist-info}/WHEEL +0 -0
- {altcodepro_polydb_python-2.3.20.dist-info → altcodepro_polydb_python-2.3.21.dist-info}/licenses/LICENSE +0 -0
- {altcodepro_polydb_python-2.3.20.dist-info → altcodepro_polydb_python-2.3.21.dist-info}/top_level.txt +0 -0
|
@@ -1,6 +1,7 @@
|
|
|
1
1
|
# src/polydb/adapters/postgres.py
|
|
2
2
|
import os
|
|
3
3
|
import threading
|
|
4
|
+
import time
|
|
4
5
|
from typing import Any, Iterator, List, Optional, Tuple, Union
|
|
5
6
|
import hashlib
|
|
6
7
|
from contextlib import contextmanager
|
|
@@ -9,9 +10,10 @@ from decimal import Decimal
|
|
|
9
10
|
from datetime import datetime, date
|
|
10
11
|
|
|
11
12
|
import psycopg2.extensions
|
|
13
|
+
from psycopg2 import sql as pg_sql
|
|
12
14
|
from psycopg2.extras import Json
|
|
13
15
|
|
|
14
|
-
from ..errors import DatabaseError, ConnectionError
|
|
16
|
+
from ..errors import DatabaseError, ConnectionError, InsufficientBalanceError
|
|
15
17
|
from ..retry import retry
|
|
16
18
|
from ..utils import validate_table_name, validate_column_name
|
|
17
19
|
from ..query import QueryBuilder
|
|
@@ -25,6 +27,7 @@ class PostgreSQLAdapter:
|
|
|
25
27
|
from ..utils import setup_logger
|
|
26
28
|
|
|
27
29
|
self.logger = setup_logger(__name__)
|
|
30
|
+
self._slow_query_ms: float = float(os.getenv("POLYDB_SLOW_QUERY_MS", "1000"))
|
|
28
31
|
self.connection_string = connection_string or os.getenv(
|
|
29
32
|
"POSTGRES_CONNECTION_STRING",
|
|
30
33
|
os.getenv("POSTGRES_URL", ""),
|
|
@@ -40,45 +43,29 @@ class PostgreSQLAdapter:
|
|
|
40
43
|
# ---------------------------------------------------------------------
|
|
41
44
|
|
|
42
45
|
def _is_idle(self, conn) -> bool:
|
|
43
|
-
"""True iff the connection is not inside a (possibly aborted) transaction."""
|
|
44
46
|
try:
|
|
45
47
|
return conn.info.transaction_status == psycopg2.extensions.TRANSACTION_STATUS_IDLE
|
|
46
48
|
except Exception:
|
|
47
49
|
return False
|
|
48
50
|
|
|
49
51
|
def _drain_transaction(self, conn) -> None:
|
|
50
|
-
"""Force the connection back to IDLE so it's safe to change session
|
|
51
|
-
settings (e.g. autocommit). Safe to call when already idle."""
|
|
52
52
|
if self._is_idle(conn):
|
|
53
53
|
return
|
|
54
54
|
try:
|
|
55
55
|
conn.rollback()
|
|
56
56
|
except Exception:
|
|
57
|
-
# Last resort: caller will surface a real error on next use.
|
|
58
57
|
pass
|
|
59
58
|
|
|
60
59
|
def _ping_connection(self, conn) -> bool:
|
|
61
|
-
"""Test if connection is still alive.
|
|
62
|
-
|
|
63
|
-
psycopg2 auto-starts a transaction on the first statement after
|
|
64
|
-
commit/rollback, so we MUST rollback after the SELECT 1 — otherwise
|
|
65
|
-
the next caller inherits an active transaction and any attempt to
|
|
66
|
-
toggle autocommit raises ``set_session cannot be used inside a
|
|
67
|
-
transaction``.
|
|
68
|
-
"""
|
|
69
60
|
try:
|
|
70
61
|
with conn.cursor() as cur:
|
|
71
62
|
cur.execute("SELECT 1")
|
|
72
|
-
# Clean up the auto-started transaction so the connection
|
|
73
|
-
# leaves this method in IDLE state.
|
|
74
63
|
try:
|
|
75
64
|
conn.rollback()
|
|
76
65
|
except Exception:
|
|
77
66
|
pass
|
|
78
67
|
return True
|
|
79
68
|
except Exception:
|
|
80
|
-
# Best-effort cleanup on a failed ping. The connection is
|
|
81
|
-
# almost certainly broken; the caller will close it.
|
|
82
69
|
try:
|
|
83
70
|
conn.rollback()
|
|
84
71
|
except Exception:
|
|
@@ -90,7 +77,6 @@ class PostgreSQLAdapter:
|
|
|
90
77
|
import psycopg2.pool
|
|
91
78
|
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
|
|
92
79
|
|
|
93
|
-
# Enhance connection string with better Azure settings
|
|
94
80
|
dsn = self.connection_string
|
|
95
81
|
if "postgresql://" in dsn:
|
|
96
82
|
parsed = urlparse(dsn)
|
|
@@ -105,17 +91,37 @@ class PostgreSQLAdapter:
|
|
|
105
91
|
parsed = parsed._replace(query=new_query)
|
|
106
92
|
dsn = urlunparse(parsed)
|
|
107
93
|
|
|
94
|
+
_maxconn = int(os.getenv("POSTGRES_MAX_CONNECTIONS", "100"))
|
|
108
95
|
with self._lock:
|
|
109
96
|
if not self._pool:
|
|
110
97
|
self._pool = psycopg2.pool.ThreadedConnectionPool(
|
|
111
98
|
minconn=int(os.getenv("POSTGRES_MIN_CONNECTIONS", "5")),
|
|
112
|
-
maxconn=
|
|
99
|
+
maxconn=_maxconn,
|
|
113
100
|
dsn=dsn,
|
|
114
101
|
)
|
|
115
|
-
self.logger.info(
|
|
102
|
+
self.logger.info(
|
|
103
|
+
"PostgreSQL pool initialized: min=%s max=%s",
|
|
104
|
+
os.getenv("POSTGRES_MIN_CONNECTIONS", "5"),
|
|
105
|
+
_maxconn,
|
|
106
|
+
)
|
|
116
107
|
except Exception as e:
|
|
117
108
|
raise ConnectionError(f"Failed to initialize PostgreSQL pool: {str(e)}")
|
|
118
109
|
|
|
110
|
+
def _log_pool_utilization(self) -> None:
|
|
111
|
+
if not self._pool:
|
|
112
|
+
return
|
|
113
|
+
try:
|
|
114
|
+
maxconn = int(os.getenv("POSTGRES_MAX_CONNECTIONS", "100"))
|
|
115
|
+
used_count = len(getattr(self._pool, "_used", {}))
|
|
116
|
+
pct = (used_count / maxconn * 100) if maxconn else 0
|
|
117
|
+
if pct > 80:
|
|
118
|
+
self.logger.warning(
|
|
119
|
+
"PostgreSQL pool utilization: %d/%d connections in use (%.0f%%)",
|
|
120
|
+
used_count, maxconn, pct,
|
|
121
|
+
)
|
|
122
|
+
except Exception:
|
|
123
|
+
pass
|
|
124
|
+
|
|
119
125
|
def _get_connection(self) -> Any:
|
|
120
126
|
if not self._pool:
|
|
121
127
|
self._initialize_pool()
|
|
@@ -123,16 +129,13 @@ class PostgreSQLAdapter:
|
|
|
123
129
|
try:
|
|
124
130
|
conn = self._pool.getconn() # type: ignore
|
|
125
131
|
|
|
126
|
-
# Critical: Validate connection for Azure transient issues.
|
|
127
|
-
# _ping_connection rolls back its own SELECT 1 so the connection
|
|
128
|
-
# is returned to callers in IDLE state.
|
|
129
132
|
if not self._ping_connection(conn):
|
|
130
133
|
self.logger.warning("Stale connection detected from pool, closing and retrying")
|
|
131
134
|
self._pool.putconn(conn, close=True) # type: ignore
|
|
132
|
-
conn = self._pool.getconn() # type: ignore
|
|
133
|
-
# New connection might still have been used before — make sure it's idle.
|
|
135
|
+
conn = self._pool.getconn() # type: ignore
|
|
134
136
|
self._drain_transaction(conn)
|
|
135
137
|
|
|
138
|
+
self._log_pool_utilization()
|
|
136
139
|
return conn
|
|
137
140
|
|
|
138
141
|
except Exception as e:
|
|
@@ -140,20 +143,45 @@ class PostgreSQLAdapter:
|
|
|
140
143
|
raise ConnectionError(f"Could not obtain database connection: {e}") from e
|
|
141
144
|
|
|
142
145
|
def _return_connection(self, conn: Any):
|
|
143
|
-
"""Return a connection to the pool, defensively draining any
|
|
144
|
-
leftover transaction state. Belt-and-suspenders: every operation
|
|
145
|
-
in this adapter already commits/rolls-back before returning, but
|
|
146
|
-
if any path ever forgets, the pool still gets a clean connection.
|
|
147
|
-
"""
|
|
148
146
|
if self._pool and conn:
|
|
149
147
|
self._drain_transaction(conn)
|
|
150
148
|
self._pool.putconn(conn)
|
|
151
149
|
|
|
150
|
+
# ---------------------------------------------------------------------
|
|
151
|
+
# QUERY TIMING HELPER
|
|
152
|
+
# ---------------------------------------------------------------------
|
|
153
|
+
|
|
154
|
+
def _timed_execute(
|
|
155
|
+
self,
|
|
156
|
+
cursor: Any,
|
|
157
|
+
sql: str,
|
|
158
|
+
params: Any,
|
|
159
|
+
*,
|
|
160
|
+
operation: str = "execute",
|
|
161
|
+
table: str = "",
|
|
162
|
+
) -> float:
|
|
163
|
+
t0 = time.perf_counter()
|
|
164
|
+
cursor.execute(sql, params)
|
|
165
|
+
duration_ms = (time.perf_counter() - t0) * 1000.0
|
|
166
|
+
|
|
167
|
+
self.logger.debug(
|
|
168
|
+
"SQL executed",
|
|
169
|
+
extra={"operation": operation, "table": table, "duration_ms": round(duration_ms, 3)},
|
|
170
|
+
)
|
|
171
|
+
|
|
172
|
+
if duration_ms > self._slow_query_ms:
|
|
173
|
+
self.logger.warning(
|
|
174
|
+
"Slow query detected: operation=%s table=%s duration_ms=%.1f threshold=%.1f",
|
|
175
|
+
operation, table, duration_ms, self._slow_query_ms,
|
|
176
|
+
)
|
|
177
|
+
|
|
178
|
+
return duration_ms
|
|
179
|
+
|
|
152
180
|
# ---------------------------------------------------------------------
|
|
153
181
|
# TRANSACTIONS
|
|
154
182
|
# ---------------------------------------------------------------------
|
|
183
|
+
|
|
155
184
|
def reset_pool(self):
|
|
156
|
-
"""Reset the entire pool (call during startup or after major failures)"""
|
|
157
185
|
with self._lock:
|
|
158
186
|
if self._pool:
|
|
159
187
|
try:
|
|
@@ -164,27 +192,17 @@ class PostgreSQLAdapter:
|
|
|
164
192
|
self._initialize_pool()
|
|
165
193
|
|
|
166
194
|
def begin_transaction(self) -> Any:
|
|
167
|
-
"""Begin a transaction and return the connection handle.
|
|
168
|
-
|
|
169
|
-
Defensively drains any leftover transaction state on the pooled
|
|
170
|
-
connection before toggling autocommit. Without this, a connection
|
|
171
|
-
that's still in TRANSACTION_STATUS_INTRANS (from a previous user
|
|
172
|
-
or from the pool's connection check) causes psycopg2 to raise
|
|
173
|
-
``set_session cannot be used inside a transaction``.
|
|
174
|
-
"""
|
|
175
195
|
conn = self._get_connection()
|
|
176
196
|
self._drain_transaction(conn)
|
|
177
197
|
conn.autocommit = False
|
|
178
198
|
return conn
|
|
179
199
|
|
|
180
200
|
def commit(self, tx: Any):
|
|
181
|
-
"""Commit the transaction using the provided connection."""
|
|
182
201
|
if tx:
|
|
183
202
|
tx.commit()
|
|
184
203
|
self._return_connection(tx)
|
|
185
204
|
|
|
186
205
|
def rollback(self, tx: Any):
|
|
187
|
-
"""Rollback the transaction using the provided connection."""
|
|
188
206
|
if tx:
|
|
189
207
|
tx.rollback()
|
|
190
208
|
self._return_connection(tx)
|
|
@@ -192,12 +210,8 @@ class PostgreSQLAdapter:
|
|
|
192
210
|
# ---------------------------------------------------------------------
|
|
193
211
|
# JSON HELPERS
|
|
194
212
|
# ---------------------------------------------------------------------
|
|
213
|
+
|
|
195
214
|
def _json_safe(self, obj: Any):
|
|
196
|
-
"""
|
|
197
|
-
Ensure JSON serialization never fails. Used only for Json() wrapping.
|
|
198
|
-
Recurses into dicts, lists, AND tuples so nested datetime/Decimal
|
|
199
|
-
values are made safe at any depth.
|
|
200
|
-
"""
|
|
201
215
|
if isinstance(obj, datetime):
|
|
202
216
|
return obj.isoformat()
|
|
203
217
|
if isinstance(obj, Decimal):
|
|
@@ -211,21 +225,6 @@ class PostgreSQLAdapter:
|
|
|
211
225
|
return obj
|
|
212
226
|
|
|
213
227
|
def _serialize_value(self, v: Any) -> Any:
|
|
214
|
-
"""
|
|
215
|
-
Serialize a value being WRITTEN to a column (insert / update SET /
|
|
216
|
-
upsert data values).
|
|
217
|
-
|
|
218
|
-
This platform provisions every ARRAY and OBJECT field as a JSONB
|
|
219
|
-
column (SchemaProvisioner._FIELD_TO_SQL_TYPE never emits a native
|
|
220
|
-
text[]/int[] column), so dicts AND lists/tuples are ALWAYS
|
|
221
|
-
JSON-encoded via psycopg2's Json adapter.
|
|
222
|
-
|
|
223
|
-
Json() handles empty {} and [] correctly ('{}'::jsonb / '[]'::jsonb),
|
|
224
|
-
handles nesting, and never produces a quoted string — which is what
|
|
225
|
-
broke the earlier native-list, json.dumps, and NULL-ify-empties
|
|
226
|
-
attempts. Empties are NOT turned into NULL: an empty JSONB array/
|
|
227
|
-
object is a valid, meaningful value.
|
|
228
|
-
"""
|
|
229
228
|
if v is None:
|
|
230
229
|
return None
|
|
231
230
|
if isinstance(v, (dict, list, tuple)):
|
|
@@ -237,15 +236,6 @@ class PostgreSQLAdapter:
|
|
|
237
236
|
return v
|
|
238
237
|
|
|
239
238
|
def _serialize_param(self, v: Any) -> Any:
|
|
240
|
-
"""
|
|
241
|
-
Serialize a value used as a QUERY PARAMETER (WHERE values, IN / ANY
|
|
242
|
-
lists, LIMIT/OFFSET, raw execute() params).
|
|
243
|
-
|
|
244
|
-
Unlike _serialize_value, primitive lists/tuples are kept NATIVE so
|
|
245
|
-
that ``IN %s`` / ``= ANY(%s)`` parameters expand correctly. A dict
|
|
246
|
-
(or a list that contains dicts) is JSON-encoded so it can be compared
|
|
247
|
-
against a JSONB column.
|
|
248
|
-
"""
|
|
249
239
|
if v is None:
|
|
250
240
|
return None
|
|
251
241
|
if isinstance(v, dict):
|
|
@@ -254,7 +244,7 @@ class PostgreSQLAdapter:
|
|
|
254
244
|
seq = list(v)
|
|
255
245
|
if any(isinstance(x, dict) for x in seq):
|
|
256
246
|
return Json(self._json_safe(seq))
|
|
257
|
-
return seq
|
|
247
|
+
return seq
|
|
258
248
|
if isinstance(v, (datetime, date)):
|
|
259
249
|
return v
|
|
260
250
|
if isinstance(v, Decimal):
|
|
@@ -265,14 +255,9 @@ class PostgreSQLAdapter:
|
|
|
265
255
|
return [self._serialize_param(p) for p in params]
|
|
266
256
|
|
|
267
257
|
def _deserialize_row(self, row: JsonDict) -> JsonDict:
|
|
268
|
-
"""
|
|
269
|
-
Postgres JSON/JSONB often comes back as dict/list already depending on
|
|
270
|
-
driver config. If it comes as a string, try json.loads safely.
|
|
271
|
-
"""
|
|
272
258
|
for k, v in list(row.items()):
|
|
273
259
|
if isinstance(v, str):
|
|
274
260
|
s = v.strip()
|
|
275
|
-
# quick cheap check to avoid parsing normal strings
|
|
276
261
|
if (s.startswith("{") and s.endswith("}")) or (
|
|
277
262
|
s.startswith("[") and s.endswith("]")
|
|
278
263
|
):
|
|
@@ -300,24 +285,18 @@ class PostgreSQLAdapter:
|
|
|
300
285
|
|
|
301
286
|
try:
|
|
302
287
|
cursor = conn.cursor()
|
|
303
|
-
|
|
304
288
|
columns = ", ".join(data.keys())
|
|
305
289
|
placeholders = ", ".join(["%s"] * len(data))
|
|
306
290
|
query = f"INSERT INTO {table} ({columns}) VALUES ({placeholders}) RETURNING *"
|
|
307
|
-
|
|
308
291
|
values = [self._serialize_value(v) for v in data.values()]
|
|
309
|
-
|
|
310
|
-
|
|
292
|
+
self._timed_execute(cursor, query, values, operation="insert", table=table)
|
|
311
293
|
result_row = cursor.fetchone()
|
|
312
294
|
columns_list = [desc[0] for desc in cursor.description]
|
|
313
295
|
result = dict(zip(columns_list, result_row))
|
|
314
|
-
|
|
315
296
|
if own_conn:
|
|
316
297
|
conn.commit()
|
|
317
|
-
|
|
318
298
|
cursor.close()
|
|
319
299
|
return self._deserialize_row(result)
|
|
320
|
-
|
|
321
300
|
except Exception as e:
|
|
322
301
|
if own_conn:
|
|
323
302
|
conn.rollback()
|
|
@@ -348,7 +327,6 @@ class PostgreSQLAdapter:
|
|
|
348
327
|
|
|
349
328
|
try:
|
|
350
329
|
cursor = conn.cursor()
|
|
351
|
-
|
|
352
330
|
sql = f"SELECT * FROM {table}"
|
|
353
331
|
params: List[Any] = []
|
|
354
332
|
|
|
@@ -356,8 +334,6 @@ class PostgreSQLAdapter:
|
|
|
356
334
|
where_parts: List[str] = []
|
|
357
335
|
for k, v in query.items():
|
|
358
336
|
validate_column_name(k)
|
|
359
|
-
|
|
360
|
-
# IMPORTANT: None must be "IS NULL" not "= %s"
|
|
361
337
|
if v is None:
|
|
362
338
|
where_parts.append(f"{k} IS NULL")
|
|
363
339
|
elif isinstance(v, (list, tuple)):
|
|
@@ -367,7 +343,6 @@ class PostgreSQLAdapter:
|
|
|
367
343
|
else:
|
|
368
344
|
where_parts.append(f"{k} = %s")
|
|
369
345
|
params.append(v)
|
|
370
|
-
|
|
371
346
|
if where_parts:
|
|
372
347
|
sql += " WHERE " + " AND ".join(where_parts)
|
|
373
348
|
|
|
@@ -378,14 +353,14 @@ class PostgreSQLAdapter:
|
|
|
378
353
|
sql += " OFFSET %s"
|
|
379
354
|
params.append(offset)
|
|
380
355
|
|
|
381
|
-
|
|
356
|
+
self._timed_execute(
|
|
357
|
+
cursor, sql, self._serialize_params(params), operation="select", table=table
|
|
358
|
+
)
|
|
382
359
|
columns = [desc[0] for desc in cursor.description]
|
|
383
360
|
results = [self._deserialize_row(dict(zip(columns, row))) for row in cursor.fetchall()]
|
|
384
361
|
cursor.close()
|
|
385
|
-
|
|
386
362
|
if own_conn:
|
|
387
363
|
conn.commit()
|
|
388
|
-
|
|
389
364
|
return results
|
|
390
365
|
except Exception as e:
|
|
391
366
|
if own_conn:
|
|
@@ -413,11 +388,9 @@ class PostgreSQLAdapter:
|
|
|
413
388
|
) -> Tuple[List[JsonDict], Optional[str]]:
|
|
414
389
|
offset = int(continuation_token) if continuation_token else 0
|
|
415
390
|
results = self.select(table, query, limit=page_size + 1, offset=offset, tx=tx)
|
|
416
|
-
|
|
417
391
|
has_more = len(results) > page_size
|
|
418
392
|
if has_more:
|
|
419
393
|
results = results[:page_size]
|
|
420
|
-
|
|
421
394
|
next_token = str(offset + page_size) if has_more else None
|
|
422
395
|
return results, next_token
|
|
423
396
|
|
|
@@ -445,9 +418,7 @@ class PostgreSQLAdapter:
|
|
|
445
418
|
|
|
446
419
|
try:
|
|
447
420
|
cursor = conn.cursor()
|
|
448
|
-
|
|
449
421
|
set_clause = ", ".join([f"{k} = %s" for k in data.keys()])
|
|
450
|
-
# SET values are written into columns -> write serializer (JSONB).
|
|
451
422
|
params: List[Any] = [self._serialize_value(v) for v in data.values()]
|
|
452
423
|
|
|
453
424
|
if isinstance(entity_id, dict):
|
|
@@ -458,7 +429,6 @@ class PostgreSQLAdapter:
|
|
|
458
429
|
where_parts.append(f"{k} IS NULL")
|
|
459
430
|
else:
|
|
460
431
|
where_parts.append(f"{k} = %s")
|
|
461
|
-
# WHERE values are query params -> param serializer.
|
|
462
432
|
params.append(self._serialize_param(v))
|
|
463
433
|
where_clause = " AND ".join(where_parts)
|
|
464
434
|
else:
|
|
@@ -466,18 +436,14 @@ class PostgreSQLAdapter:
|
|
|
466
436
|
params.append(entity_id)
|
|
467
437
|
|
|
468
438
|
query = f"UPDATE {table} SET {set_clause} WHERE {where_clause} RETURNING *"
|
|
469
|
-
|
|
470
|
-
|
|
439
|
+
self._timed_execute(cursor, query, params, operation="update", table=table)
|
|
471
440
|
result_row = cursor.fetchone()
|
|
472
441
|
if not result_row:
|
|
473
442
|
raise DatabaseError("No rows updated")
|
|
474
|
-
|
|
475
443
|
columns = [desc[0] for desc in cursor.description]
|
|
476
444
|
result = dict(zip(columns, result_row))
|
|
477
|
-
|
|
478
445
|
if own_conn:
|
|
479
446
|
conn.commit()
|
|
480
|
-
|
|
481
447
|
cursor.close()
|
|
482
448
|
return self._deserialize_row(result)
|
|
483
449
|
except Exception as e:
|
|
@@ -506,10 +472,8 @@ class PostgreSQLAdapter:
|
|
|
506
472
|
|
|
507
473
|
try:
|
|
508
474
|
cursor = conn.cursor()
|
|
509
|
-
|
|
510
475
|
columns = ", ".join(data.keys())
|
|
511
476
|
placeholders = ", ".join(["%s"] * len(data))
|
|
512
|
-
|
|
513
477
|
conflict_columns = ["id"] if "id" in data else list(data.keys())[:1]
|
|
514
478
|
update_fields = [k for k in data.keys() if k not in conflict_columns]
|
|
515
479
|
|
|
@@ -526,14 +490,10 @@ class PostgreSQLAdapter:
|
|
|
526
490
|
{on_conflict}
|
|
527
491
|
RETURNING *
|
|
528
492
|
"""
|
|
529
|
-
|
|
530
493
|
values = [self._serialize_value(v) for v in data.values()]
|
|
531
|
-
|
|
532
|
-
|
|
494
|
+
self._timed_execute(cursor, query, values, operation="upsert", table=table)
|
|
533
495
|
result_row = cursor.fetchone()
|
|
534
496
|
if not result_row:
|
|
535
|
-
# DO NOTHING case: fetch existing row (best effort)
|
|
536
|
-
# If conflict is on id, we can read it back.
|
|
537
497
|
if "id" in conflict_columns and "id" in data:
|
|
538
498
|
cursor.execute(f"SELECT * FROM {table} WHERE id = %s", [data["id"]])
|
|
539
499
|
result_row = cursor.fetchone()
|
|
@@ -544,10 +504,8 @@ class PostgreSQLAdapter:
|
|
|
544
504
|
|
|
545
505
|
columns_list = [desc[0] for desc in cursor.description]
|
|
546
506
|
result = dict(zip(columns_list, result_row))
|
|
547
|
-
|
|
548
507
|
if own_conn:
|
|
549
508
|
conn.commit()
|
|
550
|
-
|
|
551
509
|
cursor.close()
|
|
552
510
|
return self._deserialize_row(result)
|
|
553
511
|
except Exception as e:
|
|
@@ -575,7 +533,6 @@ class PostgreSQLAdapter:
|
|
|
575
533
|
|
|
576
534
|
try:
|
|
577
535
|
cursor = conn.cursor()
|
|
578
|
-
|
|
579
536
|
params: List[Any] = []
|
|
580
537
|
if isinstance(entity_id, dict):
|
|
581
538
|
where_parts: List[str] = []
|
|
@@ -585,7 +542,6 @@ class PostgreSQLAdapter:
|
|
|
585
542
|
where_parts.append(f"{k} IS NULL")
|
|
586
543
|
else:
|
|
587
544
|
where_parts.append(f"{k} = %s")
|
|
588
|
-
# WHERE values are query params -> param serializer.
|
|
589
545
|
params.append(self._serialize_param(v))
|
|
590
546
|
where_clause = " AND ".join(where_parts)
|
|
591
547
|
else:
|
|
@@ -593,18 +549,14 @@ class PostgreSQLAdapter:
|
|
|
593
549
|
params.append(entity_id)
|
|
594
550
|
|
|
595
551
|
query = f"DELETE FROM {table} WHERE {where_clause} RETURNING *"
|
|
596
|
-
|
|
552
|
+
self._timed_execute(cursor, query, params, operation="delete", table=table)
|
|
597
553
|
result_row = cursor.fetchone()
|
|
598
|
-
|
|
599
554
|
if not result_row:
|
|
600
555
|
raise DatabaseError("No rows deleted")
|
|
601
|
-
|
|
602
556
|
columns = [desc[0] for desc in cursor.description]
|
|
603
557
|
result = dict(zip(columns, result_row))
|
|
604
|
-
|
|
605
558
|
if own_conn:
|
|
606
559
|
conn.commit()
|
|
607
|
-
|
|
608
560
|
cursor.close()
|
|
609
561
|
return self._deserialize_row(result)
|
|
610
562
|
except Exception as e:
|
|
@@ -623,9 +575,7 @@ class PostgreSQLAdapter:
|
|
|
623
575
|
def query_linq(
|
|
624
576
|
self, table: str, builder: QueryBuilder, tx: Optional[Any] = None
|
|
625
577
|
) -> Union[List[JsonDict], int]:
|
|
626
|
-
|
|
627
578
|
table = validate_table_name(table)
|
|
628
|
-
|
|
629
579
|
conn = tx
|
|
630
580
|
own_conn = False
|
|
631
581
|
if not conn:
|
|
@@ -634,117 +584,67 @@ class PostgreSQLAdapter:
|
|
|
634
584
|
|
|
635
585
|
try:
|
|
636
586
|
cursor = conn.cursor()
|
|
637
|
-
|
|
638
|
-
# ------------------------------------------------
|
|
639
|
-
# SELECT clause
|
|
640
|
-
# ------------------------------------------------
|
|
641
|
-
|
|
642
587
|
if builder.count_only:
|
|
643
588
|
sql = f"SELECT COUNT(*) FROM {table}"
|
|
644
|
-
|
|
645
589
|
elif builder.selected_fields:
|
|
646
590
|
for f in builder.selected_fields:
|
|
647
591
|
validate_column_name(f)
|
|
648
|
-
|
|
649
592
|
fields = ", ".join(builder.selected_fields)
|
|
650
|
-
|
|
651
593
|
if builder.distinct_flag:
|
|
652
594
|
sql = f"SELECT DISTINCT {fields} FROM {table}"
|
|
653
595
|
else:
|
|
654
596
|
sql = f"SELECT {fields} FROM {table}"
|
|
655
|
-
|
|
656
597
|
else:
|
|
657
598
|
sql = f"SELECT * FROM {table}"
|
|
658
599
|
|
|
659
600
|
params: List[Any] = []
|
|
660
|
-
|
|
661
|
-
# ------------------------------------------------
|
|
662
|
-
# WHERE
|
|
663
|
-
# ------------------------------------------------
|
|
664
|
-
|
|
665
601
|
where_clause, where_params = builder.to_sql_where()
|
|
666
|
-
|
|
667
602
|
if where_clause:
|
|
668
603
|
sql += f" WHERE {where_clause}"
|
|
669
604
|
params.extend(where_params)
|
|
670
605
|
|
|
671
|
-
# ------------------------------------------------
|
|
672
|
-
# GROUP BY
|
|
673
|
-
# ------------------------------------------------
|
|
674
|
-
|
|
675
606
|
if builder.group_by_fields:
|
|
676
|
-
|
|
677
607
|
for f in builder.group_by_fields:
|
|
678
608
|
validate_column_name(f)
|
|
679
|
-
|
|
680
609
|
sql += f" GROUP BY {', '.join(builder.group_by_fields)}"
|
|
681
610
|
|
|
682
|
-
# ------------------------------------------------
|
|
683
|
-
# ORDER BY
|
|
684
|
-
# ------------------------------------------------
|
|
685
|
-
|
|
686
611
|
if builder.order_by_fields:
|
|
687
|
-
|
|
688
612
|
order_parts = []
|
|
689
|
-
|
|
690
613
|
for field, desc in builder.order_by_fields:
|
|
691
|
-
|
|
692
614
|
validate_column_name(field)
|
|
693
|
-
|
|
694
615
|
direction = "DESC" if desc else "ASC"
|
|
695
|
-
|
|
696
616
|
order_parts.append(f"{field} {direction}")
|
|
697
|
-
|
|
698
617
|
sql += f" ORDER BY {', '.join(order_parts)}"
|
|
699
618
|
|
|
700
|
-
# ------------------------------------------------
|
|
701
|
-
# LIMIT
|
|
702
|
-
# ------------------------------------------------
|
|
703
|
-
|
|
704
619
|
if builder.take_count is not None:
|
|
705
620
|
sql += " LIMIT %s"
|
|
706
621
|
params.append(builder.take_count)
|
|
707
622
|
|
|
708
|
-
# ------------------------------------------------
|
|
709
|
-
# OFFSET
|
|
710
|
-
# ------------------------------------------------
|
|
711
|
-
|
|
712
623
|
if builder.skip_count:
|
|
713
624
|
sql += " OFFSET %s"
|
|
714
625
|
params.append(builder.skip_count)
|
|
715
626
|
|
|
716
|
-
|
|
717
|
-
|
|
718
|
-
|
|
719
|
-
|
|
720
|
-
cursor.execute(sql, self._serialize_params(params))
|
|
627
|
+
self._timed_execute(
|
|
628
|
+
cursor, sql, self._serialize_params(params), operation="query_linq", table=table
|
|
629
|
+
)
|
|
721
630
|
|
|
722
631
|
if builder.count_only:
|
|
723
632
|
result = cursor.fetchone()[0]
|
|
724
|
-
|
|
725
633
|
else:
|
|
726
634
|
columns = [desc[0] for desc in cursor.description]
|
|
727
|
-
|
|
728
635
|
result = [
|
|
729
636
|
self._deserialize_row(dict(zip(columns, row))) for row in cursor.fetchall()
|
|
730
637
|
]
|
|
731
638
|
|
|
732
639
|
cursor.close()
|
|
733
|
-
|
|
734
640
|
if own_conn:
|
|
735
641
|
conn.commit()
|
|
736
|
-
|
|
737
642
|
return result
|
|
738
|
-
|
|
739
643
|
except Exception as e:
|
|
740
|
-
|
|
741
644
|
if own_conn:
|
|
742
645
|
conn.rollback()
|
|
743
|
-
|
|
744
646
|
raise DatabaseError(f"LINQ query failed: {str(e)}")
|
|
745
|
-
|
|
746
647
|
finally:
|
|
747
|
-
|
|
748
648
|
if own_conn and conn:
|
|
749
649
|
self._return_connection(conn)
|
|
750
650
|
|
|
@@ -771,10 +671,10 @@ class PostgreSQLAdapter:
|
|
|
771
671
|
cursor = None
|
|
772
672
|
try:
|
|
773
673
|
cursor = conn.cursor()
|
|
774
|
-
|
|
775
|
-
|
|
674
|
+
# Log operation shape only, never parameter values
|
|
675
|
+
self.logger.debug("Executing raw SQL (%d params)", len(params or []))
|
|
776
676
|
exec_params = self._serialize_params(params or [])
|
|
777
|
-
|
|
677
|
+
self._timed_execute(cursor, sql, exec_params, operation="execute", table="")
|
|
778
678
|
|
|
779
679
|
if fetch_one:
|
|
780
680
|
row = cursor.fetchone()
|
|
@@ -797,7 +697,6 @@ class PostgreSQLAdapter:
|
|
|
797
697
|
if own_conn:
|
|
798
698
|
conn.commit()
|
|
799
699
|
return None
|
|
800
|
-
|
|
801
700
|
except Exception as e:
|
|
802
701
|
if own_conn:
|
|
803
702
|
try:
|
|
@@ -805,7 +704,6 @@ class PostgreSQLAdapter:
|
|
|
805
704
|
except Exception:
|
|
806
705
|
pass
|
|
807
706
|
raise DatabaseError(f"Execute failed: {str(e)}")
|
|
808
|
-
|
|
809
707
|
finally:
|
|
810
708
|
if cursor:
|
|
811
709
|
try:
|
|
@@ -815,17 +713,234 @@ class PostgreSQLAdapter:
|
|
|
815
713
|
if own_conn and conn:
|
|
816
714
|
self._return_connection(conn)
|
|
817
715
|
|
|
716
|
+
# ---------------------------------------------------------------------
|
|
717
|
+
# ATOMIC OPERATIONS
|
|
718
|
+
# ---------------------------------------------------------------------
|
|
719
|
+
|
|
720
|
+
def atomic_decrement_if_sufficient(
|
|
721
|
+
self,
|
|
722
|
+
table: str,
|
|
723
|
+
balance_field: str,
|
|
724
|
+
amount: Any,
|
|
725
|
+
where_field: str,
|
|
726
|
+
where_value: Any,
|
|
727
|
+
*,
|
|
728
|
+
tx: Optional[Any] = None,
|
|
729
|
+
) -> JsonDict:
|
|
730
|
+
table = validate_table_name(table)
|
|
731
|
+
validate_column_name(balance_field)
|
|
732
|
+
validate_column_name(where_field)
|
|
733
|
+
|
|
734
|
+
conn = tx
|
|
735
|
+
own_conn = False
|
|
736
|
+
if not conn:
|
|
737
|
+
conn = self._get_connection()
|
|
738
|
+
own_conn = True
|
|
739
|
+
|
|
740
|
+
try:
|
|
741
|
+
cursor = conn.cursor()
|
|
742
|
+
sql = (
|
|
743
|
+
f"UPDATE {table} "
|
|
744
|
+
f"SET {balance_field} = {balance_field} - %s "
|
|
745
|
+
f"WHERE {where_field} = %s AND {balance_field} >= %s "
|
|
746
|
+
f"RETURNING *"
|
|
747
|
+
)
|
|
748
|
+
cursor.execute(sql, [amount, where_value, amount])
|
|
749
|
+
result_row = cursor.fetchone()
|
|
750
|
+
|
|
751
|
+
if not result_row:
|
|
752
|
+
if own_conn:
|
|
753
|
+
conn.rollback()
|
|
754
|
+
raise InsufficientBalanceError(
|
|
755
|
+
f"Insufficient balance in {table}.{balance_field} for {where_field}={where_value!r}"
|
|
756
|
+
)
|
|
757
|
+
|
|
758
|
+
columns = [desc[0] for desc in cursor.description]
|
|
759
|
+
result = dict(zip(columns, result_row))
|
|
760
|
+
if own_conn:
|
|
761
|
+
conn.commit()
|
|
762
|
+
cursor.close()
|
|
763
|
+
return self._deserialize_row(result)
|
|
764
|
+
except InsufficientBalanceError:
|
|
765
|
+
raise
|
|
766
|
+
except Exception as e:
|
|
767
|
+
if own_conn:
|
|
768
|
+
try:
|
|
769
|
+
conn.rollback()
|
|
770
|
+
except Exception:
|
|
771
|
+
pass
|
|
772
|
+
raise DatabaseError(f"atomic_decrement_if_sufficient failed: {str(e)}")
|
|
773
|
+
finally:
|
|
774
|
+
if own_conn and conn:
|
|
775
|
+
self._return_connection(conn)
|
|
776
|
+
|
|
777
|
+
def atomic_set_if(
|
|
778
|
+
self,
|
|
779
|
+
table: str,
|
|
780
|
+
id_field: str,
|
|
781
|
+
id_value: Any,
|
|
782
|
+
field: str,
|
|
783
|
+
value: Any,
|
|
784
|
+
expected: Any,
|
|
785
|
+
*,
|
|
786
|
+
tx: Optional[Any] = None,
|
|
787
|
+
) -> Optional[JsonDict]:
|
|
788
|
+
table = validate_table_name(table)
|
|
789
|
+
validate_column_name(field)
|
|
790
|
+
validate_column_name(id_field)
|
|
791
|
+
|
|
792
|
+
conn = tx
|
|
793
|
+
own_conn = not conn
|
|
794
|
+
if own_conn:
|
|
795
|
+
conn = self._get_connection()
|
|
796
|
+
try:
|
|
797
|
+
cursor = conn.cursor()
|
|
798
|
+
sql = (
|
|
799
|
+
f"UPDATE {table} SET {field} = %s "
|
|
800
|
+
f"WHERE {id_field} = %s AND {field} = %s RETURNING *"
|
|
801
|
+
)
|
|
802
|
+
cursor.execute(sql, [value, id_value, expected])
|
|
803
|
+
row = cursor.fetchone()
|
|
804
|
+
if own_conn:
|
|
805
|
+
conn.commit()
|
|
806
|
+
cursor.close()
|
|
807
|
+
if row is None:
|
|
808
|
+
return None
|
|
809
|
+
cols = [d[0] for d in cursor.description]
|
|
810
|
+
return self._deserialize_row(dict(zip(cols, row)))
|
|
811
|
+
except Exception as e:
|
|
812
|
+
if own_conn:
|
|
813
|
+
try:
|
|
814
|
+
conn.rollback()
|
|
815
|
+
except Exception:
|
|
816
|
+
pass
|
|
817
|
+
raise DatabaseError(f"atomic_set_if failed: {e}") from e
|
|
818
|
+
finally:
|
|
819
|
+
if own_conn and conn:
|
|
820
|
+
self._return_connection(conn)
|
|
821
|
+
|
|
822
|
+
def atomic_max(
|
|
823
|
+
self,
|
|
824
|
+
table: str,
|
|
825
|
+
id_field: str,
|
|
826
|
+
id_value: Any,
|
|
827
|
+
field: str,
|
|
828
|
+
value: Any,
|
|
829
|
+
*,
|
|
830
|
+
tx: Optional[Any] = None,
|
|
831
|
+
) -> JsonDict:
|
|
832
|
+
table = validate_table_name(table)
|
|
833
|
+
validate_column_name(field)
|
|
834
|
+
validate_column_name(id_field)
|
|
835
|
+
|
|
836
|
+
conn = tx
|
|
837
|
+
own_conn = not conn
|
|
838
|
+
if own_conn:
|
|
839
|
+
conn = self._get_connection()
|
|
840
|
+
try:
|
|
841
|
+
cursor = conn.cursor()
|
|
842
|
+
cursor.execute(
|
|
843
|
+
f"UPDATE {table} SET {field} = GREATEST({field}, %s) WHERE {id_field} = %s RETURNING *",
|
|
844
|
+
[value, id_value],
|
|
845
|
+
)
|
|
846
|
+
row = cursor.fetchone()
|
|
847
|
+
if own_conn:
|
|
848
|
+
conn.commit()
|
|
849
|
+
cursor.close()
|
|
850
|
+
if row is None:
|
|
851
|
+
raise DatabaseError(f"atomic_max: row not found {id_field}={id_value!r}")
|
|
852
|
+
cols = [d[0] for d in cursor.description]
|
|
853
|
+
return self._deserialize_row(dict(zip(cols, row)))
|
|
854
|
+
except DatabaseError:
|
|
855
|
+
raise
|
|
856
|
+
except Exception as e:
|
|
857
|
+
if own_conn:
|
|
858
|
+
try:
|
|
859
|
+
conn.rollback()
|
|
860
|
+
except Exception:
|
|
861
|
+
pass
|
|
862
|
+
raise DatabaseError(f"atomic_max failed: {e}") from e
|
|
863
|
+
finally:
|
|
864
|
+
if own_conn and conn:
|
|
865
|
+
self._return_connection(conn)
|
|
866
|
+
|
|
867
|
+
def atomic_min(
|
|
868
|
+
self,
|
|
869
|
+
table: str,
|
|
870
|
+
id_field: str,
|
|
871
|
+
id_value: Any,
|
|
872
|
+
field: str,
|
|
873
|
+
value: Any,
|
|
874
|
+
*,
|
|
875
|
+
tx: Optional[Any] = None,
|
|
876
|
+
) -> JsonDict:
|
|
877
|
+
table = validate_table_name(table)
|
|
878
|
+
validate_column_name(field)
|
|
879
|
+
validate_column_name(id_field)
|
|
880
|
+
|
|
881
|
+
conn = tx
|
|
882
|
+
own_conn = not conn
|
|
883
|
+
if own_conn:
|
|
884
|
+
conn = self._get_connection()
|
|
885
|
+
try:
|
|
886
|
+
cursor = conn.cursor()
|
|
887
|
+
cursor.execute(
|
|
888
|
+
f"UPDATE {table} SET {field} = LEAST({field}, %s) WHERE {id_field} = %s RETURNING *",
|
|
889
|
+
[value, id_value],
|
|
890
|
+
)
|
|
891
|
+
row = cursor.fetchone()
|
|
892
|
+
if own_conn:
|
|
893
|
+
conn.commit()
|
|
894
|
+
cursor.close()
|
|
895
|
+
if row is None:
|
|
896
|
+
raise DatabaseError(f"atomic_min: row not found {id_field}={id_value!r}")
|
|
897
|
+
cols = [d[0] for d in cursor.description]
|
|
898
|
+
return self._deserialize_row(dict(zip(cols, row)))
|
|
899
|
+
except DatabaseError:
|
|
900
|
+
raise
|
|
901
|
+
except Exception as e:
|
|
902
|
+
if own_conn:
|
|
903
|
+
try:
|
|
904
|
+
conn.rollback()
|
|
905
|
+
except Exception:
|
|
906
|
+
pass
|
|
907
|
+
raise DatabaseError(f"atomic_min failed: {e}") from e
|
|
908
|
+
finally:
|
|
909
|
+
if own_conn and conn:
|
|
910
|
+
self._return_connection(conn)
|
|
911
|
+
|
|
912
|
+
# ---------------------------------------------------------------------
|
|
913
|
+
# SAVEPOINTS
|
|
914
|
+
# Security: use pg_sql.Identifier to prevent SQL injection via savepoint names.
|
|
915
|
+
# ---------------------------------------------------------------------
|
|
916
|
+
|
|
917
|
+
def begin_savepoint(self, name: str, tx: Any) -> None:
|
|
918
|
+
try:
|
|
919
|
+
with tx.cursor() as cur:
|
|
920
|
+
cur.execute(pg_sql.SQL("SAVEPOINT {}").format(pg_sql.Identifier(name)))
|
|
921
|
+
except Exception as e:
|
|
922
|
+
raise DatabaseError(f"begin_savepoint({name!r}) failed: {str(e)}")
|
|
923
|
+
|
|
924
|
+
def rollback_to_savepoint(self, name: str, tx: Any) -> None:
|
|
925
|
+
try:
|
|
926
|
+
with tx.cursor() as cur:
|
|
927
|
+
cur.execute(pg_sql.SQL("ROLLBACK TO SAVEPOINT {}").format(pg_sql.Identifier(name)))
|
|
928
|
+
except Exception as e:
|
|
929
|
+
raise DatabaseError(f"rollback_to_savepoint({name!r}) failed: {str(e)}")
|
|
930
|
+
|
|
931
|
+
def release_savepoint(self, name: str, tx: Any) -> None:
|
|
932
|
+
try:
|
|
933
|
+
with tx.cursor() as cur:
|
|
934
|
+
cur.execute(pg_sql.SQL("RELEASE SAVEPOINT {}").format(pg_sql.Identifier(name)))
|
|
935
|
+
except Exception as e:
|
|
936
|
+
raise DatabaseError(f"release_savepoint({name!r}) failed: {str(e)}")
|
|
937
|
+
|
|
818
938
|
# ---------------------------------------------------------------------
|
|
819
939
|
# DISTRIBUTED LOCK
|
|
820
940
|
# ---------------------------------------------------------------------
|
|
821
941
|
|
|
822
942
|
@contextmanager
|
|
823
943
|
def distributed_lock(self, lock_name: str) -> Iterator[None]:
|
|
824
|
-
"""
|
|
825
|
-
PostgreSQL advisory lock (session scoped).
|
|
826
|
-
- Always unlock before returning the pooled connection.
|
|
827
|
-
- Never wrap exceptions raised by the user block.
|
|
828
|
-
"""
|
|
829
944
|
conn = None
|
|
830
945
|
cursor = None
|
|
831
946
|
lock_id = int(hashlib.sha256(lock_name.encode()).hexdigest(), 16) % (2**63)
|
|
@@ -833,8 +948,6 @@ class PostgreSQLAdapter:
|
|
|
833
948
|
try:
|
|
834
949
|
conn = self._get_connection()
|
|
835
950
|
cursor = conn.cursor()
|
|
836
|
-
|
|
837
|
-
# Acquire lock (DB op) — if this fails, it's a DatabaseError
|
|
838
951
|
try:
|
|
839
952
|
cursor.execute("SELECT pg_advisory_lock(%s);", (lock_id,))
|
|
840
953
|
self.logger.debug("Acquired distributed lock: %s", lock_name)
|
|
@@ -842,21 +955,15 @@ class PostgreSQLAdapter:
|
|
|
842
955
|
raise DatabaseError(f"Distributed lock acquire failed: {e}") from e
|
|
843
956
|
|
|
844
957
|
try:
|
|
845
|
-
# User code runs here. If it raises, it MUST propagate unchanged.
|
|
846
958
|
yield
|
|
847
959
|
finally:
|
|
848
|
-
# Release lock (DB op). Always attempted.
|
|
849
960
|
try:
|
|
850
961
|
cursor.execute("SELECT pg_advisory_unlock(%s);", (lock_id,))
|
|
851
962
|
self.logger.debug("Released distributed lock: %s", lock_name)
|
|
852
963
|
except Exception as e:
|
|
853
|
-
# IMPORTANT: don't mask user exceptions.
|
|
854
|
-
# If unlock fails, raise DatabaseError only if user block didn't already fail.
|
|
855
|
-
# Easiest safe behavior: just log and continue.
|
|
856
964
|
self.logger.exception(
|
|
857
965
|
"Distributed lock release failed for %s: %s", lock_name, e
|
|
858
966
|
)
|
|
859
|
-
|
|
860
967
|
finally:
|
|
861
968
|
if cursor:
|
|
862
969
|
try:
|