/
/
/
1"""Database helpers and logic."""
2
3from __future__ import annotations
4
5import asyncio
6import logging
7import os
8import re
9import time
10from collections.abc import Mapping
11from contextlib import asynccontextmanager
12from contextvars import ContextVar
13from sqlite3 import OperationalError
14from typing import TYPE_CHECKING, Any, cast
15
16import aiosqlite
17
18from music_assistant.constants import MASS_LOGGER_NAME
19
20if TYPE_CHECKING:
21 from collections.abc import AsyncGenerator, Sequence
22
23LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.database")
24
25
26class _UnsetType:
27 """Sentinel value to indicate a field should use the database default."""
28
29 _instance: _UnsetType | None = None
30
31 def __new__(cls) -> _UnsetType: # noqa: PYI034 # singleton sentinel always returns the one instance, not Self
32 """Create singleton instance."""
33 if cls._instance is None:
34 cls._instance = super().__new__(cls)
35 return cls._instance
36
37 def __repr__(self) -> str:
38 """Return string representation."""
39 return "UNSET"
40
41 def __bool__(self) -> bool:
42 """Return False for boolean context."""
43 return False
44
45
46UNSET: _UnsetType = _UnsetType()
47
48ENABLE_DEBUG = os.environ.get("PYTHONDEVMODE") == "1"
49
50SLOW_QUERY_THRESHOLD = 0.5
51_STALL_SAMPLE_INTERVAL = 0.05
52
53
54class _LoopStallTracker:
55 """Samples how long the event loop spends unavailable, so query timings can discount it."""
56
57 def __init__(self) -> None:
58 """Initialize class."""
59 self._total = 0.0
60 self._last_tick = 0.0
61 self._task: asyncio.Task[None] | None = None
62 self._users = 0
63
64 @property
65 def stalled(self) -> float:
66 """Return the total time the event loop was unavailable, to compare two readings of."""
67 if self._task is None:
68 return 0.0
69 # the sampler can only record a stall once the loop frees up again, and it hands any
70 # awaiting query its result first, so a stall in progress lives in how overdue the
71 # sampler currently is. Counting it here as well keeps this reading continuous, which
72 # is what lets two readings bracket exactly the stalls that fall between them
73 overdue = asyncio.get_running_loop().time() - self._last_tick - _STALL_SAMPLE_INTERVAL
74 return self._total + max(0.0, overdue)
75
76 def acquire(self) -> bool:
77 """Start sampling on behalf of one more user; False when there is nothing to sample."""
78 if not ENABLE_DEBUG:
79 return False
80 self._users += 1
81 if self._task is None:
82 loop = asyncio.get_running_loop()
83 self._last_tick = loop.time()
84 self._task = loop.create_task(self._sample())
85 return True
86
87 def release(self) -> None:
88 """Stop sampling once the last user has released the tracker."""
89 self._users = max(0, self._users - 1)
90 if self._users == 0 and self._task is not None:
91 self._task.cancel()
92 self._task = None
93
94 async def _sample(self) -> None:
95 """Record how late each wake-up is, which is the time the loop was unavailable."""
96 loop = asyncio.get_running_loop()
97 while True:
98 await asyncio.sleep(_STALL_SAMPLE_INTERVAL)
99 now = loop.time()
100 self._total += max(0.0, now - self._last_tick - _STALL_SAMPLE_INTERVAL)
101 self._last_tick = now
102
103
104_loop_stalls = _LoopStallTracker()
105
106
107@asynccontextmanager
108async def debug_query(
109 sql_query: str, query_params: dict[str, Any] | None = None
110) -> AsyncGenerator[None]:
111 """Time the processing time of an sql query."""
112 if not ENABLE_DEBUG:
113 yield
114 return
115 time_start = time.monotonic()
116 stalled_start = _loop_stalls.stalled
117 try:
118 yield
119 except OperationalError as err:
120 LOGGER.error(f"{err}\n{sql_query}")
121 raise
122 finally:
123 # queries run on aiosqlite's connection thread, so the awaited wall time also covers any
124 # stretch the loop was blocked elsewhere and could not deliver the result. Discounting
125 # that keeps an unrelated blocking callback from reporting as a slow query; a stall that
126 # overlaps a genuinely slow query is discounted too, so this under-reports rather than
127 # points at the wrong culprit.
128 process_time = time.monotonic() - time_start - (_loop_stalls.stalled - stalled_start)
129 if process_time > SLOW_QUERY_THRESHOLD:
130 # log slow queries
131 for key, value in (query_params or {}).items():
132 sql_query = sql_query.replace(f":{key}", repr(value))
133 LOGGER.warning("SQL Query took %s seconds! (\n%s\n", process_time, sql_query)
134
135
136def query_params(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any]]:
137 """Extend query parameters support."""
138 if params is None:
139 return (query, {})
140 count = 0
141 result_query = query
142 result_params = {}
143 for key, value in params.items():
144 # add support for a list within the query params
145 # recreates the params as (:_param_0, :_param_1) etc
146 if isinstance(value, list | tuple):
147 subparams = []
148 for subval in value:
149 subparam_name = f"_param_{count}"
150 result_params[subparam_name] = subval
151 subparams.append(subparam_name)
152 count += 1
153 params_str = ",".join(f":{x}" for x in subparams)
154 # replace the placeholder with the expanded (:_param_x, ...) list;
155 # consume optional parens already around the placeholder and use a
156 # word boundary so placeholders sharing the same prefix are untouched
157 result_query = re.sub(
158 rf"\(\s*:{re.escape(key)}\b\s*\)|:{re.escape(key)}\b",
159 f"({params_str})",
160 result_query,
161 )
162 else:
163 result_params[key] = value
164 return (result_query, result_params)
165
166
167def get_sqlite_memory_settings() -> tuple[int, int]:
168 """
169 Return (cache_size_kib, mmap_size_bytes) scaled to available system memory.
170
171 The page cache is a per-connection ceiling that is filled lazily, so a small database
172 never consumes a large ceiling. Hosts with ample RAM keep the previous generous values
173 (and a much bigger cache on very large hosts) so performance is unaffected â only memory-
174 constrained devices are scaled down. Returns the generous defaults when memory is
175 unknown (e.g. Windows), so those hosts fail open to full performance.
176 """
177 # imported lazily to keep this low-level helper free of the heavier util import chain
178 from music_assistant.helpers.util import get_total_system_memory # noqa: PLC0415
179
180 # SQLite caps mmap_size at its build-time SQLITE_MAX_MMAP_SIZE (~2GiB), so the previous
181 # 30GB request was already effectively ~2GiB. We keep that ceiling on capable hosts and
182 # only request less on memory-constrained devices (where a large DB would otherwise map
183 # most of the file into reclaimable RSS). The page cache is the only fast path for the
184 # part of a database beyond that ~2GiB window, so hosts with lots of RAM get a much larger
185 # cache to keep a very large library hot.
186 gib = 1024**3
187 total_ram_gb = get_total_system_memory()
188 if total_ram_gb >= 16.0:
189 # very large host: cache enough to keep a multi-GB library hot in memory
190 return 1024000, 2 * gib
191 if total_ram_gb >= 12.0:
192 return 512000, 2 * gib
193 if total_ram_gb >= 8.0:
194 # plenty of RAM: favour performance with a larger page cache
195 return 128000, 2 * gib
196 if total_ram_gb == 0.0 or total_ram_gb >= 4.0:
197 # unknown (fail open) or capable: keep the previous 64MB cache and ~2GiB mmap ceiling
198 return 64000, 2 * gib
199 if total_ram_gb >= 2.0:
200 return 32000, gib
201 # memory-constrained device: keep each connection's footprint small
202 return 16000, 256 * 1024 * 1024
203
204
205class DatabaseConnection:
206 """Class that holds the (connection to the) database with some convenience helper functions."""
207
208 _db: aiosqlite.Connection
209
210 def __init__(self, db_path: str) -> None:
211 """Initialize class."""
212 self.db_path = db_path
213 # per-instance ContextVar (instead of module level) so multiple database
214 # connections (library/cache/auth) track their deferred_commit scopes
215 # independently; only a handful of long-lived instances exist per process
216 self._deferred_commit_depth: ContextVar[int] = ContextVar(
217 "deferred_commit_depth", default=0
218 )
219 self._tracking_loop_stalls = False
220
221 async def setup(
222 self,
223 cache_size_kib: int | None = None,
224 mmap_size_bytes: int | None = None,
225 ) -> None:
226 """
227 Perform async initialization.
228
229 :param cache_size_kib: SQLite page-cache ceiling for this connection, in KiB.
230 Defaults to a value scaled to the host's available memory.
231 :param mmap_size_bytes: SQLite memory-map ceiling for this connection, in bytes.
232 Defaults to a value scaled to the host's available memory.
233 """
234 default_cache_kib, default_mmap_bytes = get_sqlite_memory_settings()
235 # coerce + clamp to non-negative ints so the values are always safe to interpolate
236 cache_size_kib = max(
237 0, int(default_cache_kib if cache_size_kib is None else cache_size_kib)
238 )
239 mmap_size_bytes = max(
240 0, int(default_mmap_bytes if mmap_size_bytes is None else mmap_size_bytes)
241 )
242 self._db = await aiosqlite.connect(self.db_path)
243 self._db.row_factory = aiosqlite.Row
244 # setup some default settings for more performance
245 await self.execute("PRAGMA analysis_limit=10000;")
246 await self.execute("PRAGMA locking_mode=exclusive;")
247 await self.execute("PRAGMA journal_mode=WAL;")
248 await self.execute("PRAGMA journal_size_limit = 6144000;")
249 await self.execute("PRAGMA synchronous=normal;")
250 await self.execute("PRAGMA temp_store=memory;")
251 await self.execute(f"PRAGMA mmap_size = {mmap_size_bytes};")
252 await self.execute(f"PRAGMA cache_size = -{cache_size_kib};")
253 await self.commit()
254 self._tracking_loop_stalls = _loop_stalls.acquire()
255
256 async def close(self) -> None:
257 """Close db connection on exit."""
258 await self.execute("PRAGMA optimize;")
259 await self.commit()
260 await self._db.close()
261 # mirror the acquire in setup() exactly, so a connection that failed to set up or is
262 # closed twice cannot release a slot that belongs to one of the other connections
263 if self._tracking_loop_stalls:
264 self._tracking_loop_stalls = False
265 _loop_stalls.release()
266
267 async def get_rows(
268 self,
269 table: str,
270 match: dict[str, Any] | None = None,
271 order_by: str | None = None,
272 limit: int = 500,
273 offset: int = 0,
274 ) -> list[Mapping[str, Any]]:
275 """Get all rows for given table."""
276 sql_query = f"SELECT * FROM {table}"
277 if match is not None:
278 sql_query += " WHERE " + " AND ".join(f"{x} = :{x}" for x in match)
279 if order_by is not None:
280 sql_query += f" ORDER BY {order_by}"
281 if limit:
282 sql_query += f" LIMIT {limit} OFFSET {offset}"
283 async with debug_query(sql_query):
284 return cast(
285 "list[Mapping[str, Any]]", await self._db.execute_fetchall(sql_query, match)
286 )
287
288 async def get_rows_from_query(
289 self,
290 query: str,
291 params: dict[str, Any] | None = None,
292 limit: int = 500,
293 offset: int = 0,
294 ) -> list[Mapping[str, Any]]:
295 """Get all rows for given custom query."""
296 if limit:
297 query += f" LIMIT {limit} OFFSET {offset}"
298 _query, _params = query_params(query, params)
299 async with debug_query(_query, _params):
300 return cast("list[Mapping[str, Any]]", await self._db.execute_fetchall(_query, _params))
301
302 async def iter_rows_from_query(
303 self,
304 query: str,
305 params: dict[str, Any] | None = None,
306 ) -> AsyncGenerator[Mapping[str, Any]]:
307 """Stream rows for a given custom query without materializing the full result."""
308 _query, _params = query_params(query, params)
309 async with debug_query(_query, _params), self._db.execute(_query, _params) as cursor:
310 async for row in cursor:
311 yield cast("Mapping[str, Any]", row)
312
313 async def get_count_from_query(
314 self,
315 query: str,
316 params: dict[str, Any] | None = None,
317 ) -> int:
318 """Get row count for given custom query."""
319 query = f"SELECT count() FROM ({query})"
320 _query, _params = query_params(query, params)
321 async with debug_query(_query):
322 async with self._db.execute(_query, _params) as cursor:
323 if result := await cursor.fetchone():
324 assert isinstance(result[0], int) # for type checking
325 return result[0]
326 return 0
327
328 async def get_count(
329 self,
330 table: str,
331 ) -> int:
332 """Get row count for given table."""
333 query = f"SELECT count(*) FROM {table}"
334 async with debug_query(query):
335 async with self._db.execute(query) as cursor:
336 if result := await cursor.fetchone():
337 assert isinstance(result[0], int) # for type checking
338 return result[0]
339 return 0
340
341 async def search(
342 self, table: str, search: str, column: str = "name"
343 ) -> list[Mapping[str, Any]]:
344 """Search table by column."""
345 sql_query = f"SELECT * FROM {table} WHERE {table}.{column} LIKE :search"
346 params = {"search": f"%{search}%"}
347 async with debug_query(sql_query, params):
348 return cast(
349 "list[Mapping[str, Any]]", await self._db.execute_fetchall(sql_query, params)
350 )
351
352 async def get_row(self, table: str, match: dict[str, Any]) -> Mapping[str, Any] | None:
353 """Get single row for given table where column matches keys/values."""
354 sql_query = f"SELECT * FROM {table} WHERE "
355 sql_query += " AND ".join(f"{table}.{x} = :{x}" for x in match)
356 async with debug_query(sql_query, match), self._db.execute(sql_query, match) as cursor:
357 return cast("Mapping[str, Any] | None", await cursor.fetchone())
358
359 async def insert(
360 self,
361 table: str,
362 values: dict[str, Any],
363 allow_replace: bool = False,
364 ) -> int:
365 """Insert data in given table."""
366 # Filter out UNSET values so database defaults are used
367 values = {k: v for k, v in values.items() if v is not UNSET}
368 keys = tuple(values.keys())
369 if allow_replace:
370 sql_query = f"INSERT OR REPLACE INTO {table}({','.join(keys)})"
371 else:
372 sql_query = f"INSERT INTO {table}({','.join(keys)})"
373 sql_query += f" VALUES ({','.join(f':{x}' for x in keys)})"
374 row_id = await self._db.execute_insert(sql_query, values)
375 await self._maybe_commit()
376 assert row_id is not None # for type checking
377 assert isinstance(row_id[0], int) # for type checking
378 return row_id[0]
379
380 async def insert_or_replace(self, table: str, values: dict[str, Any]) -> int:
381 """Insert or replace data in given table."""
382 return await self.insert(table=table, values=values, allow_replace=True)
383
384 async def upsert(self, table: str, values: dict[str, Any]) -> None:
385 """Upsert data in given table."""
386 # Filter out UNSET values so database defaults are used
387 values = {k: v for k, v in values.items() if v is not UNSET}
388 keys = tuple(values.keys())
389 sql_query = (
390 f"INSERT INTO {table}({','.join(keys)}) VALUES ({','.join(f':{x}' for x in keys)})"
391 )
392 sql_query += f" ON CONFLICT DO UPDATE SET {','.join(f'{x}=:{x}' for x in keys)}"
393 await self._db.execute(sql_query, values)
394 await self._maybe_commit()
395
396 async def upsert_many(self, table: str, values: Sequence[dict[str, Any]]) -> None:
397 """
398 Upsert multiple rows in the given table with a single commit.
399
400 :param table: The table to upsert the rows into.
401 :param values: The rows to upsert, each given as a column->value dict.
402 Rows do not need to share the same set of columns.
403 """
404 if not values:
405 return
406 # rows are grouped by their column set so each group can be executed as a
407 # single (prepared) statement, while omitted columns keep their existing
408 # value on conflict - identical to calling upsert() per row
409 rows_per_column_set: dict[tuple[str, ...], list[dict[str, Any]]] = {}
410 for row in values:
411 # Filter out UNSET values so database defaults are used
412 filtered_row = {k: v for k, v in row.items() if v is not UNSET}
413 rows_per_column_set.setdefault(tuple(sorted(filtered_row)), []).append(filtered_row)
414 for keys, rows in rows_per_column_set.items():
415 sql_query = (
416 f"INSERT INTO {table}({','.join(keys)}) VALUES ({','.join(f':{x}' for x in keys)})"
417 )
418 sql_query += f" ON CONFLICT DO UPDATE SET {','.join(f'{x}=:{x}' for x in keys)}"
419 await self._db.executemany(sql_query, rows)
420 await self._maybe_commit()
421
422 async def update(
423 self,
424 table: str,
425 match: dict[str, Any],
426 values: dict[str, Any],
427 ) -> None:
428 """Update record."""
429 # Filter out UNSET values so those fields are not updated
430 values = {k: v for k, v in values.items() if v is not UNSET}
431 keys = tuple(values.keys())
432 sql_query = f"UPDATE {table} SET {','.join(f'{x}=:{x}' for x in keys)} WHERE "
433 sql_query += " AND ".join(f"{x} = :{x}" for x in match)
434 await self.execute(sql_query, {**match, **values})
435 await self._maybe_commit()
436
437 async def delete(
438 self, table: str, match: dict[str, Any] | None = None, query: str | None = None
439 ) -> None:
440 """Delete data in given table."""
441 assert not (match and query), "Cannot use both match and query"
442 sql_query = f"DELETE FROM {table} "
443 if match:
444 sql_query += " WHERE " + " AND ".join(f"{x} = :{x}" for x in match)
445 elif query and "where" not in query.lower():
446 sql_query += "WHERE " + query
447 elif query:
448 sql_query += query
449 await self.execute(sql_query, match)
450 await self._maybe_commit()
451
452 async def delete_where_query(self, table: str, query: str | None = None) -> None:
453 """Delete data in given table using given where clausule."""
454 sql_query = f"DELETE FROM {table} WHERE {query}"
455 await self.execute(sql_query)
456 await self._maybe_commit()
457
458 async def execute(self, query: str, values: dict[str, Any] | None = None) -> Any:
459 """Execute command on the database."""
460 return await self._db.execute(query, values)
461
462 async def execute_write(self, query: str, values: dict[str, Any] | None = None) -> None:
463 """
464 Execute a hand-written write statement and commit it.
465
466 Use instead of `execute` for anything that modifies data, so the write is durable
467 even if nothing else happens to commit the shared connection afterwards. Honors
468 `deferred_commit`, so a batch still commits once at the end of its scope.
469
470 :param query: The statement to execute.
471 :param values: The values to bind to the statement's named parameters.
472 """
473 await self._db.execute(query, values)
474 await self._maybe_commit()
475
476 async def commit(self) -> None:
477 """Commit the current transaction."""
478 return await self._db.commit()
479
480 @asynccontextmanager
481 async def deferred_commit(self) -> AsyncGenerator[None]:
482 """
483 Batch all writes of the current task into a single commit when the scope exits.
484
485 Within the scope, the per-statement commit of the insert/upsert/update/delete
486 helpers is skipped for the current task and a single commit is issued when the
487 outermost scope exits (scopes may be nested). This greatly reduces the commit
488 overhead of multi-statement operations such as adding a media item with all
489 its relations to the library.
490
491 Note: this is not an atomic transaction. The scope always commits on exit -
492 also on error or cancellation - and never rolls back. Writes from other tasks
493 are unaffected and still commit immediately.
494 """
495 depth = self._deferred_commit_depth.get()
496 token = self._deferred_commit_depth.set(depth + 1)
497 try:
498 yield
499 finally:
500 self._deferred_commit_depth.reset(token)
501 # always commit on exit, never rollback: the connection is shared by all
502 # tasks, so statements from concurrent writers may interleave with this
503 # scope's statements in the same underlying SQLite transaction and a
504 # rollback would revert their (already acknowledged) writes as well
505 if depth == 0:
506 await self._db.commit()
507
508 async def iter_items(
509 self,
510 table: str,
511 match: dict[str, Any] | None = None,
512 ) -> AsyncGenerator[Mapping[str, Any]]:
513 """Iterate all items within a table."""
514 limit: int = 500
515 offset: int = 0
516 while True:
517 next_items = await self.get_rows(
518 table=table,
519 match=match,
520 offset=offset,
521 limit=limit,
522 )
523 for item in next_items:
524 yield item
525 if len(next_items) < limit:
526 break
527 await asyncio.sleep(0) # yield to eventloop
528 offset += limit
529
530 async def get_reclaimable_ratio(self) -> float:
531 """
532 Return the fraction (0..1) of the database file that a VACUUM would reclaim.
533
534 This is the share of pages on the free list and is a cheap way to decide
535 whether a (potentially expensive) VACUUM is actually worthwhile.
536 """
537 page_count = await self._get_pragma_int("page_count")
538 if page_count <= 0:
539 return 0.0
540 freelist_count = await self._get_pragma_int("freelist_count")
541 return freelist_count / page_count
542
543 async def vacuum(self) -> None:
544 """Run vacuum command on database."""
545 # VACUUM rebuilds the whole database in temp storage; with temp_store=memory that
546 # copy lives entirely in RAM and OOMs memory constrained devices on large databases,
547 # so spill it to a temp file (located at SQLITE_TMPDIR) for the duration.
548 await self._db.execute("PRAGMA temp_store=FILE;")
549 try:
550 await self._db.execute("VACUUM")
551 await self._db.commit()
552 finally:
553 await self._db.execute("PRAGMA temp_store=memory;")
554
555 async def _get_pragma_int(self, pragma: str) -> int:
556 """Return the integer value of a single-value sqlite PRAGMA."""
557 async with self._db.execute(f"PRAGMA {pragma}") as cursor:
558 row = await cursor.fetchone()
559 return int(row[0]) if row else 0
560
561 async def _maybe_commit(self) -> None:
562 """Commit now, unless the current task is inside a deferred_commit scope."""
563 if self._deferred_commit_depth.get() == 0:
564 await self._db.commit()
565