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.
@@ -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=int(os.getenv("POSTGRES_MAX_CONNECTIONS", "30")),
99
+ maxconn=_maxconn,
113
100
  dsn=dsn,
114
101
  )
115
- self.logger.info("PostgreSQL pool initialized with Azure optimizations")
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 # Get fresh connection
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 # native -> psycopg2 expands for IN / ANY
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
- cursor.execute(query, values)
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
- cursor.execute(sql, self._serialize_params(params))
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
- cursor.execute(query, params)
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
- cursor.execute(query, values)
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
- cursor.execute(query, params)
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
- # EXECUTE
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
- self.logger.debug("Executing raw SQL: %s", sql)
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
- cursor.execute(sql, exec_params)
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: