Repository navigation
Expand file tree
/
Copy pathruntime.py
More file actions
249 lines (220 loc) · 9.14 KB
/
Copy pathruntime.py
File metadata and controls
249 lines (220 loc) · 9.14 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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
"""Instance-owned DuckDB execution, metadata and shutdown resources."""
from __future__ import annotations
import asyncio
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
import logging
import os
from pathlib import Path
import threading
import uuid
import duckdb
from .cache import QueryCache
from .metadata import MetadataStorage
logger = logging.getLogger(__name__)
@dataclass
class Operation:
"""A cursor belongs to its worker until that worker closes it."""
lock: threading.Lock = field(default_factory=threading.Lock)
cursor: object = None
cancelled: bool = False
future: object = None
def interrupt(self):
with self.lock:
self.cancelled = True
if self.cursor is not None:
self.cursor.interrupt()
class DuckDBRuntime(MetadataStorage):
"""One writable database and a bounded executor per workspace.
Query IDs are supplied by the transport as connection-scoped operation IDs.
Writes use DuckDB MVCC and the query runner's bounded conflict retries; a
cursor is never shared between operations. Disconnect interrupts outstanding
work but cannot roll back statements that already committed.
"""
def __init__(
self,
database_path: str,
storage_root: Path,
*,
extensions=None,
meta_namespace="__sqlrooms",
meta_db_path=None,
sync_storage=True,
max_workers=4,
max_operations=64,
):
self.database_path = database_path
self.storage_root = Path(storage_root)
self.extensions = ["httpfs"] if extensions is None else extensions
self.sync_storage = sync_storage
self.meta_namespace = meta_namespace
self.meta_db_path = meta_db_path
self.max_operations = max_operations
self.executor = ThreadPoolExecutor(
max_workers=max_workers, thread_name_prefix="sqlrooms-db"
)
self.connection = None
self.cache = QueryCache()
self.operations: dict[str, Operation] = {}
self.closing = False
self.closed = False
self._started = False
self._startup_future = None
self._shutdown_task = None
self._persisted = False
self._connection_lock = threading.Lock()
@property
def ready(self):
"""Whether new operations may be admitted."""
return self._started and self.connection is not None and not self.closing
async def start(self):
"""Open the database and metadata before publishing readiness."""
def initialize():
self.storage_root.mkdir(parents=True, exist_ok=True)
self.connection = duckdb.connect(self.database_path)
try:
for spec in self.extensions:
name, repo = (spec.split("@", 1) + [None])[:2]
try:
if repo:
self.connection.install_extension(name, repository=repo)
else:
self.connection.install_extension(name)
self.connection.load_extension(name)
except Exception as exc:
logger.warning("Failed to load extension %s: %s", spec, exc)
self.connection.execute(f"SET threads TO {os.cpu_count() or 4}")
self.init_meta_storage(
self.meta_namespace,
self.meta_db_path,
sync_storage=self.sync_storage,
)
except BaseException:
self.connection.close()
self.connection = None
raise
if self._startup_future is not None or self.closing or self.closed:
raise RuntimeError("Runtime cannot be started twice")
self._startup_future = asyncio.get_running_loop().run_in_executor(
self.executor, initialize
)
# A cancelled caller cannot cancel work already running in the executor.
# Keep its future for close()/abort() to settle before touching the DB.
await asyncio.shield(self._startup_future)
if not self.closing:
self._started = True
async def _settle_startup(self):
if self._startup_future is not None:
try:
await asyncio.shield(self._startup_future)
except Exception:
# start() reports the initialization error; shutdown still owns
# the executor and any connection it managed to open.
pass
async def run_db_task(self, execute_with_cursor, *, query_id=None):
"""Run blocking work with bounded admission and isolated cancellation."""
if not self.ready:
raise RuntimeError("Database is not accepting work")
if len(self.operations) >= self.max_operations:
raise RuntimeError("Database busy: too many pending operations")
query_id = query_id or str(uuid.uuid4())
if query_id in self.operations:
raise ValueError("Duplicate active queryId")
operation = Operation()
self.operations[query_id] = operation
def execute():
with operation.lock:
if operation.cancelled:
raise duckdb.InterruptException("Query was cancelled")
with self._connection_lock:
cursor = self.connection.cursor()
operation.cursor = cursor
try:
return execute_with_cursor(cursor)
finally:
with operation.lock:
cursor.close()
operation.cursor = None
future = asyncio.get_running_loop().run_in_executor(self.executor, execute)
operation.future = future
future.add_done_callback(lambda _done: self.operations.pop(query_id, None))
try:
return await asyncio.shield(future)
except asyncio.CancelledError:
operation.interrupt()
# Keep the operation registered until its worker actually finishes.
# Cancelling an asyncio wrapper never proves a write was rolled back.
await asyncio.gather(asyncio.shield(future), return_exceptions=True)
raise
def cancel_query(self, query_id):
"""Interrupt only the cursor belonging to this operation."""
operation = self.operations.get(query_id)
if operation is None:
return False
operation.interrupt()
return True
async def close(self):
"""Stop admission, settle workers, checkpoint and close; propagate failures.
A failed checkpoint leaves the connection available for a close retry,
with admission stopped. It is never reported as a successful save.
"""
if self.closed:
if not self._persisted:
raise RuntimeError("Runtime was aborted without confirmed persistence")
return
task = self._shutdown_task
if task is None or task.done() and task.exception() is not None:
task = self._begin_shutdown(checkpoint=True)
await asyncio.shield(task)
if not self._persisted:
raise RuntimeError("Runtime was aborted without confirmed persistence")
async def abort(self):
"""Release resources after a failed final save, without claiming persistence."""
while not self.closed:
task = self._shutdown_task
if task is not None and not task.done():
try:
await asyncio.shield(task)
except Exception:
pass # A failed checkpoint leaves the connection for abort.
continue
if task is not None and task.done():
try:
task.result()
except Exception:
pass
if not self.closed:
await asyncio.shield(self._begin_shutdown(checkpoint=False))
def _begin_shutdown(self, *, checkpoint):
self.closing = True
task = asyncio.create_task(self._shutdown(checkpoint=checkpoint))
# Retain both the task and its failure if the initiating caller leaves.
task.add_done_callback(
lambda done: done.exception() if not done.cancelled() else None
)
self._shutdown_task = task
return task
async def _shutdown(self, *, checkpoint):
await self._settle_startup()
for operation in list(self.operations.values()):
operation.interrupt()
await asyncio.gather(
*(
asyncio.shield(op.future)
for op in list(self.operations.values())
if op.future
),
return_exceptions=True,
)
def finish():
if self.connection is not None:
if checkpoint:
self.connection.execute("CHECKPOINT")
self.connection.close()
await asyncio.shield(
asyncio.get_running_loop().run_in_executor(self.executor, finish)
)
self.connection = None
self.executor.shutdown(wait=True)
self.closed = True
self._persisted = checkpoint