Repository navigation
Expand file tree
/
Copy pathquery.py
More file actions
171 lines (144 loc) · 5.39 KB
/
Copy pathquery.py
File metadata and controls
171 lines (144 loc) · 5.39 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
import logging
import random
from hashlib import sha256
from functools import partial
from typing import Optional
import pyarrow as pa
import time
logger = logging.getLogger(__name__)
# Transaction conflict retry configuration
MAX_CONFLICT_RETRIES = 5
BASE_BACKOFF_MS = 10
MAX_BACKOFF_MS = 500
JITTER_FACTOR = 0.3
def _calculate_backoff(attempt: int) -> float:
"""
Calculate exponential backoff with jitter in seconds.
Args:
attempt: Zero-based attempt number (0 = first retry)
Returns:
Backoff duration in seconds
"""
base_delay = min(BASE_BACKOFF_MS * (2**attempt), MAX_BACKOFF_MS)
jitter = base_delay * JITTER_FACTOR * random.random()
return (base_delay + jitter) / 1000.0
def get_key(sql, command):
return f"{sha256(sql.encode('utf-8')).hexdigest()}.{command}"
def retrieve(cache, query, get):
sql = query.get("sql")
command = query.get("type")
key = get_key(sql, command)
if cache is None:
return get(sql)
# Prevent concurrent computes for the same key (avoids DDL races)
lock_name = f"lock:{key}"
with cache.lock(lock_name):
result = cache.get(key)
if result is not None:
logger.debug("Cache hit")
return result
value = get(sql)
if query.get("persist", False):
cache[key] = value
return value
def get_arrow(con, sql):
# Use explicit transaction to keep it active during .to_arrow_table().
# Without this, DuckDB's auto-commit closes the transaction after con.query(),
# causing "ActiveTransaction called without active transaction" when Arrow export
# needs to look up CRS metadata for geometry columns.
started_transaction = False
try:
con.execute("BEGIN TRANSACTION")
started_transaction = True
except Exception:
# Already inside a transaction — proceed without starting a new one
pass
try:
result = con.query(sql)
if result is None:
if started_transaction:
con.execute("COMMIT")
return None
arrow_result = result.to_arrow_table()
if started_transaction:
con.execute("COMMIT")
return arrow_result
except Exception as e:
if started_transaction:
try:
con.execute("ROLLBACK")
except Exception:
pass
logger.error(f"Failed to convert result to Arrow: {str(e)}")
raise
def arrow_to_bytes(table):
if table is None:
return None
sink = pa.BufferOutputStream()
with pa.ipc.new_stream(sink, table.schema) as writer:
# Prefer write_table if available for efficiency
if hasattr(writer, "write_table"):
writer.write_table(table)
else:
for batch in table.to_batches():
writer.write_batch(batch)
return sink.getvalue().to_pybytes()
def get_arrow_bytes(con, sql):
return arrow_to_bytes(get_arrow(con, sql))
def get_json(con, sql):
result = con.query(sql).df()
return result.to_json(orient="records")
async def run_duckdb(runtime, cache, query, query_id: Optional[str] = None):
"""
Run a DuckDB command asynchronously via runtime.run_db_task, returning a structured result.
The actual DB work runs in a thread, using a per-task cursor. Cancellation is handled by the runtime.
"""
logger.debug(
f"Executing DuckDB query:\n{query['sql'][:256]}{'...' if len(query['sql']) > 256 else ''}"
)
def _is_conflict_error(exc: Exception) -> bool:
msg = str(exc).lower()
# Match various DuckDB MVCC conflict patterns:
# - "Transaction conflict" - general conflict
# - "Conflict on update" - update conflicts
# - "write-write conflict" - catalog/DDL conflicts
return (
"transaction conflict" in msg
or "conflict on" in msg
or "write-write conflict" in msg
)
def _execute_once(con):
command = query["type"]
if command == "arrow":
buffer = retrieve(cache, query, partial(get_arrow_bytes, con))
return {"type": "arrow", "data": buffer}
elif command == "json":
data = retrieve(cache, query, partial(get_json, con))
return {"type": "json", "data": data}
elif command == "exec":
sql = query.get("sql")
con.execute(sql)
return {"type": "ok"}
else:
raise ValueError(f"Unknown command {command}")
def _execute_with_cursor(con):
attempts = 0
while True:
try:
return _execute_once(con)
except Exception as e:
if attempts < MAX_CONFLICT_RETRIES and _is_conflict_error(e):
backoff = _calculate_backoff(attempts)
logger.warning(
f"Transaction conflict detected (attempt {attempts + 1}/{MAX_CONFLICT_RETRIES}); "
f"retrying in {backoff * 1000:.1f}ms. Error: {e}"
)
attempts += 1
time.sleep(backoff)
continue
if _is_conflict_error(e):
logger.error(
f"Transaction conflict persisted after {MAX_CONFLICT_RETRIES} retries. Error: {e}"
)
raise
return await runtime.run_db_task(_execute_with_cursor, query_id=query_id)