altcodepro-polydb-python 2.5.24__py3-none-any.whl → 2.6.1__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,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: altcodepro-polydb-python
3
- Version: 2.5.24
3
+ Version: 2.6.1
4
4
  Summary: Production-ready multi-cloud database abstraction layer with connection pooling, retry logic, and thread safety
5
5
  Author: AltCodePro
6
6
  Project-URL: Homepage, https://github.com/altcodepro/polydb-python
@@ -23,7 +23,7 @@ Classifier: Operating System :: OS Independent
23
23
  Requires-Python: >=3.11
24
24
  Description-Content-Type: text/markdown
25
25
  License-File: LICENSE
26
- Requires-Dist: psycopg2-binary>=2.9.11
26
+ Requires-Dist: psycopg[binary,pool]>=3.2
27
27
  Requires-Dist: tenacity>=9.1.4
28
28
  Requires-Dist: redis>=6.4.0
29
29
  Requires-Dist: python-dotenv>=1.1.1
@@ -1,4 +1,4 @@
1
- altcodepro_polydb_python-2.5.24.dist-info/licenses/LICENSE,sha256=9X8GLocsBwy-5aR5JGOt2SAMDDPs9Qv-YnqmHBHOXrw,1067
1
+ altcodepro_polydb_python-2.6.1.dist-info/licenses/LICENSE,sha256=9X8GLocsBwy-5aR5JGOt2SAMDDPs9Qv-YnqmHBHOXrw,1067
2
2
  polydb/PolyDB.py,sha256=Lp0QtG7D89Z6P6nt9fVoT8rVHAy8xlE2kFRSS70r-mk,23672
3
3
  polydb/__init__.py,sha256=XuPMkOUcluxI8TIMNZCPyI1MZg77DC-10FKANEVobME,1876
4
4
  polydb/advanced_query.py,sha256=M9sKmdF2PsLZWSc-zKpXRcwaDumbYYTsibmL_IAjKP4,6931
@@ -6,7 +6,7 @@ polydb/aio.py,sha256=0EQ6hXrf6N4If6pcht6L9B7qz0PU2yGahyGzEss2eQg,8340
6
6
  polydb/batch.py,sha256=ulvSdtVeofsOiVNVAQwrE9tXIIaP8fjiW1Sbpz3dcv0,4093
7
7
  polydb/cache.py,sha256=ArI0zRbuhAYd6kn4TPVMBHBptdRmuZyy05_hwIATUsQ,11784
8
8
  polydb/cloudDatabaseFactory.py,sha256=eJiNKycdaC9nO2BKWPjkig3GaNGkNjoJci6yVJC7A_c,25301
9
- polydb/databaseFactory.py,sha256=l0ODaR3HbTN4KNpC_oEPqN1ussorXGfKnRRq2r5xt6Q,60202
9
+ polydb/databaseFactory.py,sha256=MWQhfsumX0_7G3N6XSotCHRqqtzBikT1Ogy4lqJI_hw,61404
10
10
  polydb/errors.py,sha256=AtXCLdzkHE2ncUaE1P0wH5HmIyW3NXx9KHL-fCejjA4,1832
11
11
  polydb/json_safe.py,sha256=jj3ZA6gImGwgTO0dCZVcckjeUHtewa2HP6OA1-u6VZ0,137
12
12
  polydb/models.py,sha256=Pq_63FXu1IiTxl-RshZ2vmoe-5m5k6qft-tXpHLn-70,12600
@@ -15,6 +15,7 @@ polydb/multitenancy.py,sha256=pmbrfdtycQiETnBvOqpvXzkGy2tw_ZVDLo7JOaJ9Z6M,7813
15
15
  polydb/overflow_gc.py,sha256=hgKx6ZZPVNGX6ejQ5hB7SjxI1XT9_r60kMERQ8tl9KA,6499
16
16
  polydb/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
17
17
  polydb/query.py,sha256=17V-e70df9NaiaUYPN5_MQz7O0vRdBaCR4xlZU3F2YE,8511
18
+ polydb/read_batch.py,sha256=kDol5XnIsXKfac5xy6R3FJQbcK2zHj7ce1ndtQJvlXA,6004
18
19
  polydb/registry.py,sha256=RD_elvFXcmhTdCyZDm2f3ej0elxqhArnSJ2aO9k5VCU,2352
19
20
  polydb/retry.py,sha256=0rILN3-auuzomOXPveIAvgnhBbtXXGNKvxg6qcRgN4g,4926
20
21
  polydb/schema.py,sha256=v3lnD9jUoBippMj2eLNwvmykPygO0nFx806X7cmeb8Q,9840
@@ -22,12 +23,13 @@ polydb/security.py,sha256=IF8Oh6Y6YU9xSjz_pi6lotfJuhIUJgP7cCmpGc8bz6w,20158
22
23
  polydb/types.py,sha256=UbO5tQ2oPfviszWAvZvhuYv6PMp_WhgPFgUa5tTVgKI,3026
23
24
  polydb/utils.py,sha256=YS-BhBpQfQsOf-z_i1zw_kvMYwJ05Zn1S2Xt6K4_2ak,3171
24
25
  polydb/validation.py,sha256=DK9-AFndTr0cLodpc3KX0wnABDpw2vPB4H6hp67BJKE,4570
26
+ polydb/write_behind.py,sha256=I8WEtAHZvuswX-ItJAVWhsz1oDcPpJmfp4E8YzWGGmk,5548
25
27
  polydb/adapters/AWSSecretsManagerAdapter.py,sha256=W4B3xhQFcBAlNlwQ4uTgcDmRrOmZr6v-suC4qihTpJw,2747
26
28
  polydb/adapters/AzureBlobStorageAdapter.py,sha256=ZTovSC_CVduGnIoIQUUeNxDQlf-IbFDBRzjGDk8U9Io,7084
27
29
  polydb/adapters/AzureFileStorageAdapter.py,sha256=VZNprqlBXCuWUgtqClNT-NrQmRf-XFYEiRA2BLbf-Sc,7046
28
30
  polydb/adapters/AzureKeyVaultAdapter.py,sha256=EtsrxI3X75nORnO0KbvKTVXJfliHd-gt3y9Mo8Npugs,3365
29
31
  polydb/adapters/AzureQueueAdapter.py,sha256=1e2HiFvkH1BI_JKS9fgays2J41kOfH6pUPbZDe1wmWA,13658
30
- polydb/adapters/AzureTableStorageAdapter.py,sha256=7nNSL1yX3EPDkLptDcUBExOf7sNroNXK4oTwSPDgUsg,40785
32
+ polydb/adapters/AzureTableStorageAdapter.py,sha256=iQSzPqABcT-l_O7s-mVOiioMzJXRdjfZt2Dh65GD-xA,41623
31
33
  polydb/adapters/BlockchainBlobAdapter.py,sha256=85H3PxiL9rGDtK-sTsW26jOuUQ_SGipLSPu3s-juZRA,3298
32
34
  polydb/adapters/BlockchainFileAdapter.py,sha256=lZbRyC18LghZcVh5QyIRP7tCftNF_Ep1FKG0KRkvMmk,8335
33
35
  polydb/adapters/BlockchainKVAdapter.py,sha256=pvpiPshDE5Z_IRcSBvBwgAXT--86lgEyIsEUeqcmoE8,5422
@@ -41,7 +43,7 @@ polydb/adapters/GCPSecretManagerAdapter.py,sha256=ZFL2FGh7uTaBfpUo-HKXZAxS9s-w0w
41
43
  polydb/adapters/GCPStorageAdapter.py,sha256=ce7vCOW46XbdjiupAgMgkgopy3IV7nM5YvIwEhLIKvQ,7133
42
44
  polydb/adapters/KafkaQueueAdapter.py,sha256=_FrcI6JLIoJHPOtSIL2q1uXHeao9ZkGMyUnT68p36HQ,14329
43
45
  polydb/adapters/MongoDBAdapter.py,sha256=Rs1L9gbGj2hHiJuz62Ms7uRK6sppHc1aIO4MNrtZ3Zo,11750
44
- polydb/adapters/PostgreSQLAdapter.py,sha256=n73HfBIXCF8Em4bHPvRVvipOJKdt9qZuZWdzRNCxS3w,48166
46
+ polydb/adapters/PostgreSQLAdapter.py,sha256=8Q-O-6o37p0inwv_EgWMU7M8KthY0CDgbF9XV-lF2SQ,61367
45
47
  polydb/adapters/RabbitMQAdapter.py,sha256=rMeqWXQmnfRX_pakaZX6q1bHkZ8E1boQ_Pkj3byl4E4,19903
46
48
  polydb/adapters/S3Adapter.py,sha256=ZkMR91u1e4SnrXSUbV3siKFYdw0ykYHJfyP2cUJUVso,7564
47
49
  polydb/adapters/S3CompatibleAdapter.py,sha256=Qj9SUtCnOFNL9c05wNVjICxZacFCnee-Wn_-hh0lwwE,7006
@@ -69,7 +71,7 @@ polydb/observability/logging.py,sha256=FWYky-Mg1IhLpwuaKuLIsTROBZc9bI_-AO-y53nvE
69
71
  polydb/services/__init__.py,sha256=Zs3PZbHutgPdPjJatygvtYvKPZMBt1qKSqESbFRPHJs,241
70
72
  polydb/services/compliance_service.py,sha256=UMjCbQ4z13-w6pT4RLQpMi_u1zbShEpiz8kawA2WsSU,5075
71
73
  polydb/services/security_service.py,sha256=gGo_j9Ca_rgvxS7s6C5JUktugQBj1PV1DXpMnOYAGbI,4992
72
- altcodepro_polydb_python-2.5.24.dist-info/METADATA,sha256=6Osh8k039JDAzMEhKD7RIpHstGqXvw7rBgMVis9o8UM,15599
73
- altcodepro_polydb_python-2.5.24.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
74
- altcodepro_polydb_python-2.5.24.dist-info/top_level.txt,sha256=WgLFWJoYjUhwvyPxJFl6jYLrVFuBJDX3OABf4ocwk_E,7
75
- altcodepro_polydb_python-2.5.24.dist-info/RECORD,,
74
+ altcodepro_polydb_python-2.6.1.dist-info/METADATA,sha256=BCUzR9Bpm7q16CWx1YcTyQxzsLuFme6EMRq7m6Bd8NE,15600
75
+ altcodepro_polydb_python-2.6.1.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
76
+ altcodepro_polydb_python-2.6.1.dist-info/top_level.txt,sha256=WgLFWJoYjUhwvyPxJFl6jYLrVFuBJDX3OABf4ocwk_E,7
77
+ altcodepro_polydb_python-2.6.1.dist-info/RECORD,,
@@ -785,22 +785,34 @@ class AzureTableStorageAdapter(NoSQLKVAdapter):
785
785
  self._sanitize_prop_name(orig_k) if orig_k != _MODEL_FIELD else _MODEL_FIELD
786
786
  )
787
787
 
788
- ev = self._encode_value(orig_v)
789
-
790
- if ev is None:
791
- parts.append(f"{sk} eq null")
792
- elif isinstance(ev, bool):
793
- parts.append(f"{sk} eq {str(ev).lower()}")
794
- elif isinstance(ev, (int, float)):
795
- parts.append(f"{sk} eq {ev}")
796
- elif isinstance(ev, datetime):
797
- iso = ev.isoformat()
798
- if not iso.endswith("Z"):
799
- iso = iso + "Z"
800
- parts.append(f"{sk} eq datetime'{iso}'")
801
- else:
802
- sval = str(ev).replace("'", "''")
803
- parts.append(f"{sk} eq '{sval}'")
788
+ def _eq(prop: str, raw: Any) -> str:
789
+ ev = self._encode_value(raw)
790
+ if ev is None:
791
+ return f"{prop} eq null"
792
+ if isinstance(ev, bool):
793
+ return f"{prop} eq {str(ev).lower()}"
794
+ if isinstance(ev, (int, float)):
795
+ return f"{prop} eq {ev}"
796
+ if isinstance(ev, datetime):
797
+ iso = ev.isoformat()
798
+ if not iso.endswith("Z"):
799
+ iso = iso + "Z"
800
+ return f"{prop} eq datetime'{iso}'"
801
+ return f"{prop} eq '{str(ev).replace(chr(39), chr(39) * 2)}'"
802
+
803
+ if isinstance(orig_v, (list, tuple, set, frozenset)) and all(
804
+ isinstance(x, (str, int, float, bool)) for x in orig_v
805
+ ):
806
+ # A list of plain values means "equals any of them" (the SQL adapters already read it that way). Sent as
807
+ # one JSON-encoded string it matched no row at all, so every `field: [a, b]` lookup silently returned [].
808
+ if not orig_v:
809
+ return [] # IN () matches nothing
810
+ if len(orig_v) > 100:
811
+ raise NoSQLError("Azure Table filter: at most 100 values in a list")
812
+ parts.append("(" + " or ".join(_eq(sk, x) for x in orig_v) + ")")
813
+ continue
814
+
815
+ parts.append(_eq(sk, orig_v))
804
816
 
805
817
  # Empty-filter contract: if the caller passed filter keys but every
806
818
  # condition dropped out (e.g. all values None/unresolvable), the query
@@ -5,17 +5,21 @@ import json
5
5
  import os
6
6
  import threading
7
7
  import time
8
+ import weakref
8
9
  from contextlib import contextmanager
9
10
  from datetime import date, datetime
10
11
  from decimal import Decimal
11
12
  from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
12
13
 
13
- import psycopg2.extensions
14
- from psycopg2 import sql as pg_sql
15
- from psycopg2.extras import Json
14
+ import psycopg
15
+ import psycopg.pq
16
+ import psycopg_pool
17
+ from psycopg import sql as pg_sql
18
+ from psycopg.types.json import Jsonb as Json
16
19
 
17
20
  from ..errors import ConnectionError, DatabaseError, InsufficientBalanceError
18
21
  from ..query import Operator, QueryBuilder
22
+ from ..read_batch import SelectSpec, current_batch
19
23
  from ..retry import retry
20
24
  from ..types import JsonDict, Lookup
21
25
  from ..utils import validate_column_name, validate_table_name
@@ -35,6 +39,80 @@ from ..utils import validate_column_name, validate_table_name
35
39
  ALLOWED_SESSION_VARS = frozenset({"app.tenant_id"})
36
40
 
37
41
 
42
+
43
+ class PoolError(Exception):
44
+ """Raised when no connection frees up within the acquire timeout."""
45
+
46
+
47
+ def _close_pool_quietly(pool) -> None:
48
+ # A pool that is dropped without closeall() (a test, a process shutting down) must not raise from __del__ -- psycopg_pool's own
49
+ # finalizer does, because its worker threads outlive the object.
50
+ try:
51
+ pool.close(timeout=1.0)
52
+ except Exception:
53
+ pass
54
+
55
+
56
+ class LazyPool:
57
+ """A small, lazy, bounded connection pool on psycopg 3 (`psycopg_pool.ConnectionPool`).
58
+
59
+ Why not the old eager pool: every process used to open `POSTGRES_MIN_CONNECTIONS` connections the moment it started and keep
60
+ them for its whole life. With a handful of processes (api plus one worker per lane) and a database that only allows a few dozen
61
+ clients, the idle reservations alone filled the server ("remaining connection slots are reserved ...") and a new revision or
62
+ worker could not even start.
63
+
64
+ This pool opens connections ONLY when a request needs one, keeps at most `keep` (default 1) idle for reuse, and closes every
65
+ other idle connection after `max_idle` seconds. `max_size` is a hard ceiling per process. A request that finds every connection
66
+ in use waits up to `acquire_timeout` seconds instead of failing at once."""
67
+
68
+ def __init__(self, dsn: str, *, max_size: int, keep: int = 1, acquire_timeout: float = 10.0,
69
+ max_idle: float = 30.0, max_lifetime: float = 1800.0):
70
+ self.maxconn = max_size
71
+ self._timeout = float(acquire_timeout)
72
+ keep = max(0, min(keep, max_size))
73
+ self._pool = psycopg_pool.ConnectionPool(
74
+ conninfo=dsn,
75
+ min_size=keep,
76
+ max_size=max_size,
77
+ timeout=self._timeout,
78
+ max_idle=max_idle,
79
+ max_lifetime=max_lifetime,
80
+ reconnect_timeout=30.0,
81
+ # prepare_threshold=None: no server-side prepared statements (they pin per-connection state and break behind a
82
+ # transaction-mode pooler); autocommit off keeps the engine's explicit commit/rollback contract.
83
+ kwargs={"autocommit": False, "prepare_threshold": None},
84
+ open=False,
85
+ name="polydb-pg",
86
+ )
87
+ self._pool.open()
88
+ weakref.finalize(self, _close_pool_quietly, self._pool)
89
+
90
+ def getconn(self, key=None):
91
+ try:
92
+ return self._pool.getconn(timeout=self._timeout)
93
+ except psycopg_pool.PoolTimeout as e:
94
+ raise PoolError(
95
+ f"connection pool exhausted: waited {self._timeout:g}s for one of {self.maxconn} connections"
96
+ ) from e
97
+
98
+ def putconn(self, conn=None, key=None, close=False):
99
+ if conn is None:
100
+ return
101
+ if close:
102
+ try:
103
+ conn.close()
104
+ except Exception:
105
+ pass
106
+ self._pool.putconn(conn)
107
+
108
+ def closeall(self):
109
+ self._pool.close()
110
+
111
+ def used_count(self) -> int:
112
+ stats = self._pool.get_stats()
113
+ return max(0, int(stats.get("pool_size", 0)) - int(stats.get("pool_available", 0)))
114
+
115
+
38
116
  class PostgreSQLAdapter:
39
117
  """PostgreSQL with full LINQ support, connection pooling, JSON/JSONB support"""
40
118
 
@@ -59,7 +137,7 @@ class PostgreSQLAdapter:
59
137
 
60
138
  def _is_idle(self, conn) -> bool:
61
139
  try:
62
- return conn.info.transaction_status == psycopg2.extensions.TRANSACTION_STATUS_IDLE
140
+ return conn.info.transaction_status == psycopg.pq.TransactionStatus.IDLE
63
141
  except Exception:
64
142
  return False
65
143
 
@@ -91,8 +169,6 @@ class PostgreSQLAdapter:
91
169
  try:
92
170
  from urllib.parse import parse_qs, quote, urlencode, urlparse, urlunparse
93
171
 
94
- import psycopg2.pool
95
-
96
172
  dsn = self.connection_string
97
173
  if "postgresql://" in dsn:
98
174
  parsed = urlparse(dsn)
@@ -132,17 +208,23 @@ class PostgreSQLAdapter:
132
208
  parsed = parsed._replace(query=new_query)
133
209
  dsn = urlunparse(parsed)
134
210
 
135
- _maxconn = int(os.getenv("POSTGRES_MAX_CONNECTIONS", "100"))
211
+ _maxconn = int(os.getenv("POSTGRES_MAX_CONNECTIONS", "10"))
136
212
  with self._lock:
137
213
  if not self._pool:
138
- self._pool = psycopg2.pool.ThreadedConnectionPool(
139
- minconn=int(os.getenv("POSTGRES_MIN_CONNECTIONS", "5")),
140
- maxconn=_maxconn,
141
- dsn=dsn,
214
+ # One probe connection up front so a bad DSN / full server fails here, loudly, exactly as the old eager pool
215
+ # did -- and is closed again at once (the pool itself opens connections only on demand).
216
+ with psycopg.connect(dsn, connect_timeout=int(os.getenv("POSTGRES_CONNECT_TIMEOUT_S", "30"))) as probe:
217
+ probe.execute("SELECT 1")
218
+ self._pool = LazyPool(
219
+ dsn,
220
+ max_size=_maxconn,
221
+ keep=int(os.getenv("POSTGRES_POOL_IDLE_KEEP", "1")),
222
+ acquire_timeout=float(os.getenv("POSTGRES_POOL_ACQUIRE_TIMEOUT_S", "10")),
223
+ max_idle=float(os.getenv("POSTGRES_POOL_MAX_IDLE_S", "30")),
142
224
  )
143
225
  self.logger.info(
144
- "PostgreSQL pool initialized: min=%s max=%s",
145
- os.getenv("POSTGRES_MIN_CONNECTIONS", "5"),
226
+ "PostgreSQL pool initialized (lazy): idle-keep=%s max=%s",
227
+ os.getenv("POSTGRES_POOL_IDLE_KEEP", "1"),
146
228
  _maxconn,
147
229
  )
148
230
  except Exception as e:
@@ -152,8 +234,8 @@ class PostgreSQLAdapter:
152
234
  if not self._pool:
153
235
  return
154
236
  try:
155
- maxconn = int(os.getenv("POSTGRES_MAX_CONNECTIONS", "100"))
156
- used_count = len(getattr(self._pool, "_used", {}))
237
+ maxconn = self._pool.maxconn
238
+ used_count = self._pool.used_count()
157
239
  pct = (used_count / maxconn * 100) if maxconn else 0
158
240
  if pct > 80:
159
241
  self.logger.warning(
@@ -204,7 +286,9 @@ class PostgreSQLAdapter:
204
286
  table: str = "",
205
287
  ) -> float:
206
288
  t0 = time.perf_counter()
207
- cursor.execute(sql, params)
289
+ # psycopg 3 parses `%` as a placeholder whenever a params sequence is given, even an empty one (psycopg2 did not):
290
+ # no parameters means None, so raw SQL / DDL containing a literal `%` or several statements still runs.
291
+ cursor.execute(sql, params if params else None)
208
292
  duration_ms = (time.perf_counter() - t0) * 1000.0
209
293
 
210
294
  self.logger.debug(
@@ -323,7 +407,7 @@ class PostgreSQLAdapter:
323
407
  return Json(self._json_safe(v))
324
408
  if isinstance(v, (list, tuple)):
325
409
  # A bare list/tuple of scalars is a real Postgres array column
326
- # (e.g. TEXT[]) -- psycopg2 adapts a plain Python list to that
410
+ # (e.g. TEXT[]) -- psycopg adapts a plain Python list to that
327
411
  # natively, so it must pass through untouched here, exactly
328
412
  # like _serialize_param already does for query parameters
329
413
  # (below). Only a list/tuple actually containing dicts is
@@ -394,7 +478,7 @@ class PostgreSQLAdapter:
394
478
  both are simply "a string that starts with `{` and ends with
395
479
  `}`." This adapter deliberately does NOT try to solve that
396
480
  ambiguity generically (declaring the column JSONB instead sidesteps
397
- it entirely -- psycopg2 registers global JSON/JSONB typecasters
481
+ it entirely -- psycopg registers JSON/JSONB loaders
398
482
  for OIDs 114/3802 automatically at import time, so a genuine
399
483
  `json`/`jsonb` column never even reaches this heuristic, `isinstance
400
484
  (v, str)` already `False` by the time this method sees it); a
@@ -473,10 +557,91 @@ class PostgreSQLAdapter:
473
557
  if own_conn and conn:
474
558
  self._return_connection(conn)
475
559
 
560
+ # ---------------------------------------------------------------------
561
+ # INSERT MANY (pipelined)
562
+ # ---------------------------------------------------------------------
563
+
564
+ def insert_many(
565
+ self,
566
+ table: str,
567
+ rows: List[JsonDict],
568
+ *,
569
+ session_vars: Optional[Dict[str, str]] = None,
570
+ ) -> int:
571
+ """Insert many rows in ONE transaction and ONE network round trip (psycopg 3 `executemany` runs in pipeline mode).
572
+
573
+ All-or-nothing: if any row fails the whole batch is rolled back and a DatabaseError is raised (the caller -- see
574
+ `write_behind.WriteBehind` -- then retries the rows one by one). Rows may have different column sets; they are grouped by
575
+ column set, and every group is sent inside the same transaction. No RETURNING: this is for rows nobody reads back."""
576
+ table = validate_table_name(table)
577
+ if not rows:
578
+ return 0
579
+ groups: Dict[Tuple[str, ...], List[JsonDict]] = {}
580
+ for row in rows:
581
+ cols = tuple(row.keys())
582
+ for c in cols:
583
+ validate_column_name(c)
584
+ groups.setdefault(cols, []).append(row)
585
+
586
+ conn = self._get_connection()
587
+ try:
588
+ self._apply_session_vars(conn, session_vars)
589
+ t0 = time.perf_counter()
590
+ with conn.cursor() as cursor:
591
+ for cols, group in groups.items():
592
+ placeholders = ", ".join(["%s"] * len(cols))
593
+ sql = f"INSERT INTO {table} ({', '.join(cols)}) VALUES ({placeholders})"
594
+ cursor.executemany(sql, [[self._serialize_value(r[c]) for c in cols] for r in group])
595
+ conn.commit()
596
+ self._log_batch(len(rows), (time.perf_counter() - t0) * 1000.0)
597
+ return len(rows)
598
+ except Exception as e:
599
+ try:
600
+ conn.rollback()
601
+ except Exception:
602
+ pass
603
+ raise DatabaseError(f"Insert many failed: {str(e)}")
604
+ finally:
605
+ self._return_connection(conn)
606
+
476
607
  # ---------------------------------------------------------------------
477
608
  # SELECT
478
609
  # ---------------------------------------------------------------------
479
610
 
611
+ @staticmethod
612
+ def _build_select(
613
+ table: str,
614
+ query: Optional[Lookup],
615
+ limit: Optional[int],
616
+ offset: Optional[int],
617
+ fields: Optional[List[str]],
618
+ ) -> Tuple[str, List[Any]]:
619
+ """The SQL text and bind parameters of a plain `select()` -- a pure function, shared by the single and the pipelined path."""
620
+ sql = f"SELECT {', '.join(fields)} FROM {table}" if fields else f"SELECT * FROM {table}"
621
+ params: List[Any] = []
622
+ if query:
623
+ where_parts: List[str] = []
624
+ for k, v in query.items():
625
+ validate_column_name(k)
626
+ if v is None:
627
+ where_parts.append(f"{k} IS NULL")
628
+ elif isinstance(v, (list, tuple)):
629
+ placeholders = ",".join(["%s"] * len(v))
630
+ where_parts.append(f"{k} IN ({placeholders})")
631
+ params.extend(list(v))
632
+ else:
633
+ where_parts.append(f"{k} = %s")
634
+ params.append(v)
635
+ if where_parts:
636
+ sql += " WHERE " + " AND ".join(where_parts)
637
+ if limit:
638
+ sql += " LIMIT %s"
639
+ params.append(limit)
640
+ if offset:
641
+ sql += " OFFSET %s"
642
+ params.append(offset)
643
+ return sql, params
644
+
480
645
  @retry(max_attempts=3, delay=1.0, exceptions=(DatabaseError,))
481
646
  def select(
482
647
  self,
@@ -493,6 +658,16 @@ class PostgreSQLAdapter:
493
658
  if fields:
494
659
  for f in fields:
495
660
  validate_column_name(f)
661
+
662
+ # Inside a ReadBatch (the engine's coll.parallel) this query is queued and sent together with its siblings in one
663
+ # pipelined round trip on one connection -- see polydb/read_batch.py. A caller-supplied transaction is never batched.
664
+ batch = current_batch() if tx is None else None
665
+ if batch is not None:
666
+ return batch.submit(
667
+ self,
668
+ SelectSpec(table=table, query=query, limit=limit, offset=offset, fields=fields, session_vars=session_vars),
669
+ )
670
+
496
671
  conn = tx
497
672
  own_conn = False
498
673
  if not conn:
@@ -502,31 +677,7 @@ class PostgreSQLAdapter:
502
677
  try:
503
678
  self._apply_session_vars(conn, session_vars)
504
679
  cursor = conn.cursor()
505
- sql = f"SELECT {', '.join(fields)} FROM {table}" if fields else f"SELECT * FROM {table}"
506
- params: List[Any] = []
507
-
508
- if query:
509
- where_parts: List[str] = []
510
- for k, v in query.items():
511
- validate_column_name(k)
512
- if v is None:
513
- where_parts.append(f"{k} IS NULL")
514
- elif isinstance(v, (list, tuple)):
515
- placeholders = ",".join(["%s"] * len(v))
516
- where_parts.append(f"{k} IN ({placeholders})")
517
- params.extend(list(v))
518
- else:
519
- where_parts.append(f"{k} = %s")
520
- params.append(v)
521
- if where_parts:
522
- sql += " WHERE " + " AND ".join(where_parts)
523
-
524
- if limit:
525
- sql += " LIMIT %s"
526
- params.append(limit)
527
- if offset:
528
- sql += " OFFSET %s"
529
- params.append(offset)
680
+ sql, params = self._build_select(table, query, limit, offset, fields)
530
681
 
531
682
  self._timed_execute(
532
683
  cursor, sql, self._serialize_params(params), operation="select", table=table
@@ -548,6 +699,119 @@ class PostgreSQLAdapter:
548
699
  if own_conn and conn:
549
700
  self._return_connection(conn)
550
701
 
702
+ def select_many(self, specs: List[SelectSpec]) -> List[Any]:
703
+ """Run several SELECTs in ONE round trip on ONE connection (psycopg 3 pipeline mode).
704
+
705
+ Returns one entry per spec, in order: the rows (list of dicts), or a `DatabaseError` for a query that failed. A failing
706
+ query never poisons the others: Postgres skips every statement after an error until the next sync point, so the specs
707
+ that were skipped are simply sent again in a second pipeline (one extra round trip per failure, nothing in the common
708
+ case).
709
+
710
+ Every query keeps the guarantees a connection of its own would give it: it runs in its own READ ONLY transaction, with its
711
+ own `session_vars` applied via `set_config(name, value, true)` (transaction-local, so row-level security sees exactly the
712
+ tenant that spec asked for and nothing leaks to the next spec or the next borrower of the connection)."""
713
+ results: List[Any] = [None] * len(specs)
714
+ if not specs:
715
+ return results
716
+
717
+ built: List[Optional[Tuple[str, List[Any]]]] = []
718
+ for i, sp in enumerate(specs):
719
+ try:
720
+ table = validate_table_name(sp.table)
721
+ if sp.fields:
722
+ for f in sp.fields:
723
+ validate_column_name(f)
724
+ if sp.builder is not None:
725
+ sql, params = self._build_linq(table, sp.builder)
726
+ else:
727
+ sql, params = self._build_select(table, sp.query, sp.limit, sp.offset, sp.fields)
728
+ for k in (sp.session_vars or {}):
729
+ if k not in ALLOWED_SESSION_VARS:
730
+ raise DatabaseError(
731
+ f"session_vars: {k!r} is not in the allow-list of Postgres "
732
+ f"session variables this adapter will set"
733
+ )
734
+ built.append((sql, self._serialize_params(params)))
735
+ except Exception as e: # a malformed spec fails alone
736
+ built.append(None)
737
+ results[i] = e if isinstance(e, DatabaseError) else DatabaseError(f"Select failed: {e}")
738
+
739
+ pending = [i for i in range(len(specs)) if built[i] is not None]
740
+ if not pending:
741
+ return results
742
+
743
+ conn = self._get_connection()
744
+ t0 = time.perf_counter()
745
+ broken = False
746
+ try:
747
+ conn.autocommit = True
748
+ while pending:
749
+ cursors: Dict[int, Any] = {}
750
+ pipeline_error: Optional[BaseException] = None
751
+ try:
752
+ with conn.pipeline():
753
+ for i in pending:
754
+ sql, params = built[i] # type: ignore[misc]
755
+ conn.execute("ROLLBACK") # no-op when idle; clears a transaction the previous failure left aborted
756
+ conn.execute("BEGIN READ ONLY")
757
+ for k, v in (specs[i].session_vars or {}).items():
758
+ conn.execute("SELECT set_config(%s, %s, true)", [k, str(v)])
759
+ cur = conn.cursor()
760
+ cur.execute(sql, params)
761
+ conn.execute("COMMIT")
762
+ cursors[i] = cur
763
+ except Exception as e: # the first failing statement; everything after it was skipped by the server
764
+ pipeline_error = e
765
+
766
+ skipped: List[int] = []
767
+ failed_seen = False
768
+ for i in pending:
769
+ if failed_seen:
770
+ skipped.append(i)
771
+ continue
772
+ cur = cursors.get(i)
773
+ try:
774
+ if specs[i].builder is not None and specs[i].builder.count_only:
775
+ results[i] = cur.fetchone()[0] # type: ignore[union-attr]
776
+ else:
777
+ columns = [d[0] for d in cur.description] # type: ignore[union-attr]
778
+ results[i] = [self._deserialize_row(dict(zip(columns, r))) for r in cur.fetchall()] # type: ignore[union-attr]
779
+ except Exception as e:
780
+ failed_seen = True
781
+ cause = pipeline_error if pipeline_error is not None else e
782
+ results[i] = DatabaseError(f"Select failed: {str(cause)}")
783
+ if not failed_seen and pipeline_error is not None:
784
+ # the pipeline failed after every result was already in (e.g. COMMIT of the last spec): nothing is lost
785
+ pass
786
+ pending = skipped
787
+ except Exception:
788
+ broken = True
789
+ raise
790
+ finally:
791
+ try:
792
+ if not conn.closed:
793
+ try:
794
+ conn.rollback()
795
+ except Exception:
796
+ pass
797
+ conn.autocommit = False
798
+ except Exception:
799
+ broken = True
800
+ if broken:
801
+ self._pool.putconn(conn, close=True) # type: ignore[union-attr]
802
+ else:
803
+ self._return_connection(conn)
804
+ self._log_batch(len(specs), (time.perf_counter() - t0) * 1000.0)
805
+ return results
806
+
807
+ def _log_batch(self, n: int, duration_ms: float) -> None:
808
+ self.logger.debug("SQL pipeline executed: statements=%d duration_ms=%.1f", n, duration_ms)
809
+ if duration_ms > self._slow_query_ms:
810
+ self.logger.warning(
811
+ "Slow query detected: operation=select_many statements=%d duration_ms=%.1f threshold=%.1f",
812
+ n, duration_ms, self._slow_query_ms,
813
+ )
814
+
551
815
  # ---------------------------------------------------------------------
552
816
  # SELECT PAGE
553
817
  # ---------------------------------------------------------------------
@@ -767,11 +1031,63 @@ class PostgreSQLAdapter:
767
1031
  # LINQ QUERY
768
1032
  # ---------------------------------------------------------------------
769
1033
 
1034
+ @staticmethod
1035
+ def _build_linq(table: str, builder: QueryBuilder) -> Tuple[str, List[Any]]:
1036
+ """The SQL text and bind parameters of a `query_linq()` -- a pure function, shared by the single and the pipelined path."""
1037
+ if builder.count_only:
1038
+ sql = f"SELECT COUNT(*) FROM {table}"
1039
+ elif builder.selected_fields:
1040
+ for f in builder.selected_fields:
1041
+ validate_column_name(f)
1042
+ fields = ", ".join(builder.selected_fields)
1043
+ if builder.distinct_flag:
1044
+ sql = f"SELECT DISTINCT {fields} FROM {table}"
1045
+ else:
1046
+ sql = f"SELECT {fields} FROM {table}"
1047
+ else:
1048
+ sql = f"SELECT * FROM {table}"
1049
+
1050
+ params: List[Any] = []
1051
+ where_clause, where_params = builder.to_sql_where()
1052
+ if where_clause:
1053
+ sql += f" WHERE {where_clause}"
1054
+ params.extend(where_params)
1055
+
1056
+ if builder.group_by_fields:
1057
+ for f in builder.group_by_fields:
1058
+ validate_column_name(f)
1059
+ sql += f" GROUP BY {', '.join(builder.group_by_fields)}"
1060
+
1061
+ if builder.order_by_fields:
1062
+ order_parts = []
1063
+ for field, desc in builder.order_by_fields:
1064
+ validate_column_name(field)
1065
+ direction = "DESC" if desc else "ASC"
1066
+ order_parts.append(f"{field} {direction}")
1067
+ sql += f" ORDER BY {', '.join(order_parts)}"
1068
+
1069
+ if builder.take_count is not None:
1070
+ sql += " LIMIT %s"
1071
+ params.append(builder.take_count)
1072
+
1073
+ if builder.skip_count:
1074
+ sql += " OFFSET %s"
1075
+ params.append(builder.skip_count)
1076
+ return sql, params
1077
+
770
1078
  @retry(max_attempts=3, delay=1.0, exceptions=(DatabaseError,))
771
1079
  def query_linq(
772
- self, table: str, builder: QueryBuilder, tx: Optional[Any] = None
1080
+ self,
1081
+ table: str,
1082
+ builder: QueryBuilder,
1083
+ tx: Optional[Any] = None,
1084
+ *,
1085
+ session_vars: Optional[Dict[str, str]] = None,
773
1086
  ) -> Union[List[JsonDict], int]:
774
1087
  table = validate_table_name(table)
1088
+ batch = current_batch() if tx is None else None
1089
+ if batch is not None:
1090
+ return batch.submit(self, SelectSpec(table=table, builder=builder, session_vars=session_vars))
775
1091
  conn = tx
776
1092
  own_conn = False
777
1093
  if not conn:
@@ -779,46 +1095,10 @@ class PostgreSQLAdapter:
779
1095
  own_conn = True
780
1096
 
781
1097
  try:
1098
+ if own_conn:
1099
+ self._apply_session_vars(conn, session_vars)
782
1100
  cursor = conn.cursor()
783
- if builder.count_only:
784
- sql = f"SELECT COUNT(*) FROM {table}"
785
- elif builder.selected_fields:
786
- for f in builder.selected_fields:
787
- validate_column_name(f)
788
- fields = ", ".join(builder.selected_fields)
789
- if builder.distinct_flag:
790
- sql = f"SELECT DISTINCT {fields} FROM {table}"
791
- else:
792
- sql = f"SELECT {fields} FROM {table}"
793
- else:
794
- sql = f"SELECT * FROM {table}"
795
-
796
- params: List[Any] = []
797
- where_clause, where_params = builder.to_sql_where()
798
- if where_clause:
799
- sql += f" WHERE {where_clause}"
800
- params.extend(where_params)
801
-
802
- if builder.group_by_fields:
803
- for f in builder.group_by_fields:
804
- validate_column_name(f)
805
- sql += f" GROUP BY {', '.join(builder.group_by_fields)}"
806
-
807
- if builder.order_by_fields:
808
- order_parts = []
809
- for field, desc in builder.order_by_fields:
810
- validate_column_name(field)
811
- direction = "DESC" if desc else "ASC"
812
- order_parts.append(f"{field} {direction}")
813
- sql += f" ORDER BY {', '.join(order_parts)}"
814
-
815
- if builder.take_count is not None:
816
- sql += " LIMIT %s"
817
- params.append(builder.take_count)
818
-
819
- if builder.skip_count:
820
- sql += " OFFSET %s"
821
- params.append(builder.skip_count)
1101
+ sql, params = self._build_linq(table, builder)
822
1102
 
823
1103
  self._timed_execute(
824
1104
  cursor, sql, self._serialize_params(params), operation="query_linq", table=table
polydb/databaseFactory.py CHANGED
@@ -425,6 +425,20 @@ class DatabaseFactory:
425
425
  # CREATE
426
426
  # ──────────────────────────────────────────────────────────
427
427
 
428
+ def _write_behind_queue(self):
429
+ wb = getattr(self, "_wb", None)
430
+ if wb is None:
431
+ from .write_behind import WriteBehind
432
+
433
+ wb = self._wb = WriteBehind()
434
+ return wb
435
+
436
+ def flush_background_writes(self, timeout: float = 10.0) -> None:
437
+ """Block until every `create(..., background=True)` row has been written. A no-op when none were queued."""
438
+ wb = getattr(self, "_wb", None)
439
+ if wb is not None:
440
+ wb.flush(timeout)
441
+
428
442
  def create(
429
443
  self,
430
444
  model: Union[type, str],
@@ -432,10 +446,19 @@ class DatabaseFactory:
432
446
  *,
433
447
  engine_override: Optional[EngineOverride] = None,
434
448
  session_vars: Optional[Dict[str, str]] = None,
449
+ background: bool = False,
435
450
  ) -> JsonDict:
451
+ """`background=True` is for telemetry-style rows nobody reads back in the same request (audit / event / query log): the row is
452
+ queued and written by one background thread in batches -- see write_behind.py. It returns the row as it will be stored (not a
453
+ database round trip), and applies only to a SQL model; anything else is written inline as usual."""
436
454
  meta = _extract_meta(model)
437
455
  name = _model_name(model)
438
456
  data = self._inject_audit_fields(data, is_create=True)
457
+ if background and self._is_sql(meta, engine_override) and not getattr(meta, "encrypted_fields", None):
458
+ self._write_behind_queue().submit(
459
+ self._adapters_for(model, meta, engine_override).sql, meta.table, data, session_vars
460
+ )
461
+ return data
439
462
 
440
463
  encrypted_fields = getattr(meta, "encrypted_fields", [])
441
464
  if self.encryption and encrypted_fields:
polydb/read_batch.py ADDED
@@ -0,0 +1,134 @@
1
+ # src/polydb/read_batch.py
2
+ """Coalesce the SQL reads several threads make at the same moment into ONE pipelined round trip on ONE connection.
3
+
4
+ The problem this solves: a request that needs N independent reads used to pay N network round trips in a row (each query waits
5
+ for the previous answer), or -- when run on N threads -- held N Postgres connections at once and could drain the connection pool
6
+ or the server's connection slots. Postgres runs one statement at a time per connection, but with psycopg 3 *pipeline mode* a
7
+ client can send many statements without waiting for each answer, so N reads cost about one round trip and one connection.
8
+
9
+ How it plugs in: a caller (the engine's `coll.parallel`) creates a `ReadBatch` for the tasks it is about to run and makes it the
10
+ current batch for those threads. `PostgreSQLAdapter.select()` -- the single place every SQL read goes through -- checks for a
11
+ current batch and, if there is one, hands its query to the batch instead of executing it, then waits for its own rows. Everything
12
+ above `select()` (permission checks, tenant scoping, audit, caches, decryption) runs untouched on each task's own thread; only the
13
+ final "send this SELECT to the database" step is shared.
14
+
15
+ A batch flushes as soon as every still-running task is waiting on it, or after `window_s` (so a task busy with non-database work
16
+ cannot hold the others up). Each query still runs in its own read-only transaction with its own `session_vars`
17
+ (`set_config(..., true)`), so row-level security sees exactly the tenant it would have seen on a connection of its own.
18
+ """
19
+ from __future__ import annotations
20
+
21
+ import contextvars
22
+ import os
23
+ import threading
24
+ from contextlib import contextmanager
25
+ from dataclasses import dataclass
26
+ from typing import Any, Dict, Iterator, List, Optional
27
+
28
+
29
+ @dataclass
30
+ class SelectSpec:
31
+ """One SELECT to run, already validated by `PostgreSQLAdapter.select()`."""
32
+
33
+ table: str
34
+ query: Optional[Dict[str, Any]] = None
35
+ limit: Optional[int] = None
36
+ offset: Optional[int] = None
37
+ fields: Optional[List[str]] = None
38
+ session_vars: Optional[Dict[str, str]] = None
39
+ # A `QueryBuilder` (LINQ-style query: projection, ORDER BY, operators, count). When set, `query`/`limit`/`offset`/`fields` are
40
+ # not used; the result is the row list, or the integer for a count-only builder.
41
+ builder: Optional[Any] = None
42
+
43
+
44
+ class _Pending:
45
+ __slots__ = ("adapter", "spec", "event", "rows", "error")
46
+
47
+ def __init__(self, adapter: Any, spec: SelectSpec):
48
+ self.adapter = adapter
49
+ self.spec = spec
50
+ self.event = threading.Event()
51
+ self.rows: Any = None
52
+ self.error: Optional[BaseException] = None
53
+
54
+
55
+ _current: contextvars.ContextVar[Optional["ReadBatch"]] = contextvars.ContextVar("polydb_read_batch", default=None)
56
+
57
+
58
+ def current_batch() -> Optional["ReadBatch"]:
59
+ return _current.get()
60
+
61
+
62
+ @contextmanager
63
+ def use_read_batch(batch: Optional["ReadBatch"]) -> Iterator[None]:
64
+ """Make `batch` the current batch for this context (and for contexts copied from it afterwards)."""
65
+ token = _current.set(batch)
66
+ try:
67
+ yield
68
+ finally:
69
+ _current.reset(token)
70
+
71
+
72
+ class ReadBatch:
73
+ """Collects SELECTs submitted by up to `parties` threads and runs them together.
74
+
75
+ Each party must call `task_done()` exactly once when it has finished (success or failure), so the batch knows nobody else is
76
+ coming and can flush what is waiting."""
77
+
78
+ def __init__(self, parties: int, window_s: Optional[float] = None):
79
+ self._active = max(0, int(parties))
80
+ self._window = (
81
+ float(window_s) if window_s is not None else float(os.getenv("POLYDB_READ_BATCH_WINDOW_MS", "5")) / 1000.0
82
+ )
83
+ self._lock = threading.Lock()
84
+ self._queue: List[_Pending] = []
85
+ self.flushes = 0 # how many pipelined round trips were made (for tests / diagnostics)
86
+ self.statements = 0 # how many SELECTs went through the batch
87
+
88
+ # -- called from a task's thread, inside PostgreSQLAdapter.select() ------------------------------------------------------
89
+ def submit(self, adapter: Any, spec: SelectSpec) -> Any:
90
+ item = _Pending(adapter, spec)
91
+ with self._lock:
92
+ self._queue.append(item)
93
+ everyone_is_waiting = len(self._queue) >= self._active
94
+ if everyone_is_waiting:
95
+ self._flush()
96
+ elif not item.event.wait(self._window):
97
+ self._flush() # the window elapsed: send whatever has queued up
98
+ item.event.wait() # another thread may be executing our query right now
99
+ if item.error is not None:
100
+ raise item.error
101
+ return item.rows if item.rows is not None else []
102
+
103
+ def task_done(self) -> None:
104
+ with self._lock:
105
+ self._active -= 1
106
+ waiting = bool(self._queue) and len(self._queue) >= self._active
107
+ if waiting:
108
+ self._flush()
109
+
110
+ # -- internals ------------------------------------------------------------------------------------------------------------
111
+ def _flush(self) -> None:
112
+ with self._lock:
113
+ batch, self._queue = self._queue, []
114
+ if not batch:
115
+ return
116
+ groups: Dict[int, List[_Pending]] = {}
117
+ for it in batch:
118
+ groups.setdefault(id(it.adapter), []).append(it)
119
+ try:
120
+ for items in groups.values():
121
+ try:
122
+ outcomes = items[0].adapter.select_many([i.spec for i in items])
123
+ except BaseException as exc: # noqa: BLE001 -- every waiter must be released
124
+ outcomes = [exc] * len(items)
125
+ self.flushes += 1
126
+ self.statements += len(items)
127
+ for it, out in zip(items, outcomes):
128
+ if isinstance(out, BaseException):
129
+ it.error = out
130
+ else:
131
+ it.rows = out
132
+ finally:
133
+ for it in batch:
134
+ it.event.set()
polydb/write_behind.py ADDED
@@ -0,0 +1,110 @@
1
+ # src/polydb/write_behind.py
2
+ """Write-behind queue: telemetry-style inserts (audit log, event log, query log) written by ONE background thread, in batches.
3
+
4
+ Why: the engine writes an audit row, an event row and a query-log row for every action it runs. Written inline, each of those is its
5
+ own transaction on a pooled connection -- two network round trips apiece, taken on whatever thread is running the action. A request
6
+ that runs 40 steps paid ~120 extra transactions before it could answer, and a `coll.parallel` group did it from several threads at
7
+ once, each holding its own connection. Here the caller only appends the row to an in-memory queue and moves on; one writer thread
8
+ groups what has queued up per table and sends each group as ONE pipelined transaction on ONE connection
9
+ (`PostgreSQLAdapter.insert_many`).
10
+
11
+ Trade-off, stated plainly: a row still queued when the process is killed hard (SIGKILL / OOM) is lost; a normal exit flushes
12
+ (`atexit`). A full queue falls back to writing inline, so nothing is dropped under load. Only for rows nobody reads back in the
13
+ same request -- never for data a caller depends on.
14
+ """
15
+ from __future__ import annotations
16
+
17
+ import atexit
18
+ import logging
19
+ import os
20
+ import queue
21
+ import threading
22
+ import time
23
+ from typing import Any, Dict, List, Optional, Tuple
24
+
25
+ logger = logging.getLogger(__name__)
26
+
27
+ _QUEUE_MAX = int(os.getenv("POLYDB_WRITE_BEHIND_QUEUE_MAX", "20000"))
28
+ _BATCH_MAX = int(os.getenv("POLYDB_WRITE_BEHIND_BATCH_MAX", "200"))
29
+ _LINGER_S = float(os.getenv("POLYDB_WRITE_BEHIND_LINGER_MS", "25")) / 1000.0
30
+
31
+ _Item = Tuple[Any, str, Dict[str, Any], Optional[Dict[str, str]]] # (sql adapter, table, row, session_vars)
32
+
33
+
34
+ class WriteBehind:
35
+ def __init__(self) -> None:
36
+ self._queue: "queue.Queue[_Item]" = queue.Queue(maxsize=_QUEUE_MAX)
37
+ self._thread: Optional[threading.Thread] = None
38
+ self._start_lock = threading.Lock()
39
+ self.batches = 0 # pipelined transactions written (diagnostics / tests)
40
+ self.rows = 0 # rows written through the queue
41
+ atexit.register(self.flush)
42
+
43
+ # ------------------------------------------------------------------------------------------------------------------
44
+ def submit(self, adapter: Any, table: str, row: Dict[str, Any], session_vars: Optional[Dict[str, str]] = None) -> None:
45
+ self._ensure_thread()
46
+ try:
47
+ self._queue.put_nowait((adapter, table, row, session_vars))
48
+ except queue.Full:
49
+ # Never drop a row: when the queue is full the caller writes it inline, exactly as it would without write-behind.
50
+ adapter.insert(table, row, session_vars=session_vars)
51
+
52
+ def flush(self, timeout: float = 10.0) -> None:
53
+ """Block until every queued row has been written (or `timeout` seconds pass)."""
54
+ if self._thread is None:
55
+ return
56
+ deadline = time.monotonic() + timeout
57
+ while (not self._queue.empty() or self._queue.unfinished_tasks) and time.monotonic() < deadline:
58
+ time.sleep(0.005)
59
+
60
+ # ------------------------------------------------------------------------------------------------------------------
61
+ def _ensure_thread(self) -> None:
62
+ if self._thread is not None and self._thread.is_alive():
63
+ return
64
+ with self._start_lock:
65
+ if self._thread is None or not self._thread.is_alive():
66
+ self._thread = threading.Thread(target=self._run, name="polydb-write-behind", daemon=True)
67
+ self._thread.start()
68
+
69
+ def _run(self) -> None:
70
+ while True:
71
+ first = self._queue.get()
72
+ batch: List[_Item] = [first]
73
+ deadline = time.monotonic() + _LINGER_S
74
+ while len(batch) < _BATCH_MAX:
75
+ remaining = deadline - time.monotonic()
76
+ if remaining <= 0:
77
+ break
78
+ try:
79
+ batch.append(self._queue.get(timeout=remaining))
80
+ except queue.Empty:
81
+ break
82
+ try:
83
+ self._write(batch)
84
+ except Exception: # noqa: BLE001 -- best-effort, the writer thread must never die
85
+ logger.warning("write-behind: batch write failed", exc_info=True)
86
+ finally:
87
+ for _ in batch:
88
+ self._queue.task_done()
89
+
90
+ def _write(self, batch: List[_Item]) -> None:
91
+ groups: Dict[Tuple[int, str, Tuple[Tuple[str, str], ...]], List[_Item]] = {}
92
+ for item in batch:
93
+ adapter, table, _row, session_vars = item
94
+ key = (id(adapter), table, tuple(sorted((session_vars or {}).items())))
95
+ groups.setdefault(key, []).append(item)
96
+ for items in groups.values():
97
+ adapter, table, _row, session_vars = items[0]
98
+ rows = [i[2] for i in items]
99
+ try:
100
+ adapter.insert_many(table, rows, session_vars=session_vars)
101
+ self.batches += 1
102
+ self.rows += len(rows)
103
+ except Exception: # noqa: BLE001 -- one bad row must not cost the rest: retry them one by one
104
+ logger.warning("write-behind: pipelined insert into %s failed; retrying row by row", table, exc_info=True)
105
+ for r in rows:
106
+ try:
107
+ adapter.insert(table, r, session_vars=session_vars)
108
+ self.rows += 1
109
+ except Exception: # noqa: BLE001
110
+ logger.warning("write-behind: dropped one row for %s", table, exc_info=True)