music-assistant-server

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
50
51@asynccontextmanager
52async def debug_query(
53    sql_query: str, query_params: dict[str, Any] | None = None
54) -> AsyncGenerator[None]:
55    """Time the processing time of an sql query."""
56    if not ENABLE_DEBUG:
57        yield
58        return
59    time_start = time.time()
60    try:
61        yield
62    except OperationalError as err:
63        LOGGER.error(f"{err}\n{sql_query}")
64        raise
65    finally:
66        process_time = time.time() - time_start
67        if process_time > 0.5:
68            # log slow queries
69            for key, value in (query_params or {}).items():
70                sql_query = sql_query.replace(f":{key}", repr(value))
71            LOGGER.warning("SQL Query took %s seconds! (\n%s\n", process_time, sql_query)
72
73
74def query_params(query: str, params: dict[str, Any] | None) -> tuple[str, dict[str, Any]]:
75    """Extend query parameters support."""
76    if params is None:
77        return (query, {})
78    count = 0
79    result_query = query
80    result_params = {}
81    for key, value in params.items():
82        # add support for a list within the query params
83        # recreates the params as (:_param_0, :_param_1) etc
84        if isinstance(value, list | tuple):
85            subparams = []
86            for subval in value:
87                subparam_name = f"_param_{count}"
88                result_params[subparam_name] = subval
89                subparams.append(subparam_name)
90                count += 1
91            params_str = ",".join(f":{x}" for x in subparams)
92            # replace the placeholder with the expanded (:_param_x, ...) list;
93            # consume optional parens already around the placeholder and use a
94            # word boundary so placeholders sharing the same prefix are untouched
95            result_query = re.sub(
96                rf"\(\s*:{re.escape(key)}\b\s*\)|:{re.escape(key)}\b",
97                f"({params_str})",
98                result_query,
99            )
100        else:
101            result_params[key] = value
102    return (result_query, result_params)
103
104
105def get_sqlite_memory_settings() -> tuple[int, int]:
106    """
107    Return (cache_size_kib, mmap_size_bytes) scaled to available system memory.
108
109    The page cache is a per-connection ceiling that is filled lazily, so a small database
110    never consumes a large ceiling. Hosts with ample RAM keep the previous generous values
111    (and a much bigger cache on very large hosts) so performance is unaffected — only memory-
112    constrained devices are scaled down. Returns the generous defaults when memory is
113    unknown (e.g. Windows), so those hosts fail open to full performance.
114    """
115    # imported lazily to keep this low-level helper free of the heavier util import chain
116    from music_assistant.helpers.util import get_total_system_memory  # noqa: PLC0415
117
118    # SQLite caps mmap_size at its build-time SQLITE_MAX_MMAP_SIZE (~2GiB), so the previous
119    # 30GB request was already effectively ~2GiB. We keep that ceiling on capable hosts and
120    # only request less on memory-constrained devices (where a large DB would otherwise map
121    # most of the file into reclaimable RSS). The page cache is the only fast path for the
122    # part of a database beyond that ~2GiB window, so hosts with lots of RAM get a much larger
123    # cache to keep a very large library hot.
124    gib = 1024**3
125    total_ram_gb = get_total_system_memory()
126    if total_ram_gb >= 16.0:
127        # very large host: cache enough to keep a multi-GB library hot in memory
128        return 1024000, 2 * gib
129    if total_ram_gb >= 12.0:
130        return 512000, 2 * gib
131    if total_ram_gb >= 8.0:
132        # plenty of RAM: favour performance with a larger page cache
133        return 128000, 2 * gib
134    if total_ram_gb == 0.0 or total_ram_gb >= 4.0:
135        # unknown (fail open) or capable: keep the previous 64MB cache and ~2GiB mmap ceiling
136        return 64000, 2 * gib
137    if total_ram_gb >= 2.0:
138        return 32000, gib
139    # memory-constrained device: keep each connection's footprint small
140    return 16000, 256 * 1024 * 1024
141
142
143class DatabaseConnection:
144    """Class that holds the (connection to the) database with some convenience helper functions."""
145
146    _db: aiosqlite.Connection
147
148    def __init__(self, db_path: str) -> None:
149        """Initialize class."""
150        self.db_path = db_path
151        # per-instance ContextVar (instead of module level) so multiple database
152        # connections (library/cache/auth) track their deferred_commit scopes
153        # independently; only a handful of long-lived instances exist per process
154        self._deferred_commit_depth: ContextVar[int] = ContextVar(
155            "deferred_commit_depth", default=0
156        )
157
158    async def setup(
159        self,
160        cache_size_kib: int | None = None,
161        mmap_size_bytes: int | None = None,
162    ) -> None:
163        """
164        Perform async initialization.
165
166        :param cache_size_kib: SQLite page-cache ceiling for this connection, in KiB.
167            Defaults to a value scaled to the host's available memory.
168        :param mmap_size_bytes: SQLite memory-map ceiling for this connection, in bytes.
169            Defaults to a value scaled to the host's available memory.
170        """
171        default_cache_kib, default_mmap_bytes = get_sqlite_memory_settings()
172        # coerce + clamp to non-negative ints so the values are always safe to interpolate
173        cache_size_kib = max(
174            0, int(default_cache_kib if cache_size_kib is None else cache_size_kib)
175        )
176        mmap_size_bytes = max(
177            0, int(default_mmap_bytes if mmap_size_bytes is None else mmap_size_bytes)
178        )
179        self._db = await aiosqlite.connect(self.db_path)
180        self._db.row_factory = aiosqlite.Row
181        # setup some default settings for more performance
182        await self.execute("PRAGMA analysis_limit=10000;")
183        await self.execute("PRAGMA locking_mode=exclusive;")
184        await self.execute("PRAGMA journal_mode=WAL;")
185        await self.execute("PRAGMA journal_size_limit = 6144000;")
186        await self.execute("PRAGMA synchronous=normal;")
187        await self.execute("PRAGMA temp_store=memory;")
188        await self.execute(f"PRAGMA mmap_size = {mmap_size_bytes};")
189        await self.execute(f"PRAGMA cache_size = -{cache_size_kib};")
190        await self.commit()
191
192    async def close(self) -> None:
193        """Close db connection on exit."""
194        await self.execute("PRAGMA optimize;")
195        await self.commit()
196        await self._db.close()
197
198    async def get_rows(
199        self,
200        table: str,
201        match: dict[str, Any] | None = None,
202        order_by: str | None = None,
203        limit: int = 500,
204        offset: int = 0,
205    ) -> list[Mapping[str, Any]]:
206        """Get all rows for given table."""
207        sql_query = f"SELECT * FROM {table}"
208        if match is not None:
209            sql_query += " WHERE " + " AND ".join(f"{x} = :{x}" for x in match)
210        if order_by is not None:
211            sql_query += f" ORDER BY {order_by}"
212        if limit:
213            sql_query += f" LIMIT {limit} OFFSET {offset}"
214        async with debug_query(sql_query):
215            return cast(
216                "list[Mapping[str, Any]]", await self._db.execute_fetchall(sql_query, match)
217            )
218
219    async def get_rows_from_query(
220        self,
221        query: str,
222        params: dict[str, Any] | None = None,
223        limit: int = 500,
224        offset: int = 0,
225    ) -> list[Mapping[str, Any]]:
226        """Get all rows for given custom query."""
227        if limit:
228            query += f" LIMIT {limit} OFFSET {offset}"
229        _query, _params = query_params(query, params)
230        async with debug_query(_query, _params):
231            return cast("list[Mapping[str, Any]]", await self._db.execute_fetchall(_query, _params))
232
233    async def iter_rows_from_query(
234        self,
235        query: str,
236        params: dict[str, Any] | None = None,
237    ) -> AsyncGenerator[Mapping[str, Any]]:
238        """Stream rows for a given custom query without materializing the full result."""
239        _query, _params = query_params(query, params)
240        async with debug_query(_query, _params), self._db.execute(_query, _params) as cursor:
241            async for row in cursor:
242                yield cast("Mapping[str, Any]", row)
243
244    async def get_count_from_query(
245        self,
246        query: str,
247        params: dict[str, Any] | None = None,
248    ) -> int:
249        """Get row count for given custom query."""
250        query = f"SELECT count() FROM ({query})"
251        _query, _params = query_params(query, params)
252        async with debug_query(_query):
253            async with self._db.execute(_query, _params) as cursor:
254                if result := await cursor.fetchone():
255                    assert isinstance(result[0], int)  # for type checking
256                    return result[0]
257            return 0
258
259    async def get_count(
260        self,
261        table: str,
262    ) -> int:
263        """Get row count for given table."""
264        query = f"SELECT count(*) FROM {table}"
265        async with debug_query(query):
266            async with self._db.execute(query) as cursor:
267                if result := await cursor.fetchone():
268                    assert isinstance(result[0], int)  # for type checking
269                    return result[0]
270            return 0
271
272    async def search(
273        self, table: str, search: str, column: str = "name"
274    ) -> list[Mapping[str, Any]]:
275        """Search table by column."""
276        sql_query = f"SELECT * FROM {table} WHERE {table}.{column} LIKE :search"
277        params = {"search": f"%{search}%"}
278        async with debug_query(sql_query, params):
279            return cast(
280                "list[Mapping[str, Any]]", await self._db.execute_fetchall(sql_query, params)
281            )
282
283    async def get_row(self, table: str, match: dict[str, Any]) -> Mapping[str, Any] | None:
284        """Get single row for given table where column matches keys/values."""
285        sql_query = f"SELECT * FROM {table} WHERE "
286        sql_query += " AND ".join(f"{table}.{x} = :{x}" for x in match)
287        async with debug_query(sql_query, match), self._db.execute(sql_query, match) as cursor:
288            return cast("Mapping[str, Any] | None", await cursor.fetchone())
289
290    async def insert(
291        self,
292        table: str,
293        values: dict[str, Any],
294        allow_replace: bool = False,
295    ) -> int:
296        """Insert data in given table."""
297        # Filter out UNSET values so database defaults are used
298        values = {k: v for k, v in values.items() if v is not UNSET}
299        keys = tuple(values.keys())
300        if allow_replace:
301            sql_query = f"INSERT OR REPLACE INTO {table}({','.join(keys)})"
302        else:
303            sql_query = f"INSERT INTO {table}({','.join(keys)})"
304        sql_query += f" VALUES ({','.join(f':{x}' for x in keys)})"
305        row_id = await self._db.execute_insert(sql_query, values)
306        await self._maybe_commit()
307        assert row_id is not None  # for type checking
308        assert isinstance(row_id[0], int)  # for type checking
309        return row_id[0]
310
311    async def insert_or_replace(self, table: str, values: dict[str, Any]) -> int:
312        """Insert or replace data in given table."""
313        return await self.insert(table=table, values=values, allow_replace=True)
314
315    async def upsert(self, table: str, values: dict[str, Any]) -> None:
316        """Upsert data in given table."""
317        # Filter out UNSET values so database defaults are used
318        values = {k: v for k, v in values.items() if v is not UNSET}
319        keys = tuple(values.keys())
320        sql_query = (
321            f"INSERT INTO {table}({','.join(keys)}) VALUES ({','.join(f':{x}' for x in keys)})"
322        )
323        sql_query += f" ON CONFLICT DO UPDATE SET {','.join(f'{x}=:{x}' for x in keys)}"
324        await self._db.execute(sql_query, values)
325        await self._maybe_commit()
326
327    async def upsert_many(self, table: str, values: Sequence[dict[str, Any]]) -> None:
328        """
329        Upsert multiple rows in the given table with a single commit.
330
331        :param table: The table to upsert the rows into.
332        :param values: The rows to upsert, each given as a column->value dict.
333            Rows do not need to share the same set of columns.
334        """
335        if not values:
336            return
337        # rows are grouped by their column set so each group can be executed as a
338        # single (prepared) statement, while omitted columns keep their existing
339        # value on conflict - identical to calling upsert() per row
340        rows_per_column_set: dict[tuple[str, ...], list[dict[str, Any]]] = {}
341        for row in values:
342            # Filter out UNSET values so database defaults are used
343            filtered_row = {k: v for k, v in row.items() if v is not UNSET}
344            rows_per_column_set.setdefault(tuple(sorted(filtered_row)), []).append(filtered_row)
345        for keys, rows in rows_per_column_set.items():
346            sql_query = (
347                f"INSERT INTO {table}({','.join(keys)}) VALUES ({','.join(f':{x}' for x in keys)})"
348            )
349            sql_query += f" ON CONFLICT DO UPDATE SET {','.join(f'{x}=:{x}' for x in keys)}"
350            await self._db.executemany(sql_query, rows)
351        await self._maybe_commit()
352
353    async def update(
354        self,
355        table: str,
356        match: dict[str, Any],
357        values: dict[str, Any],
358    ) -> None:
359        """Update record."""
360        # Filter out UNSET values so those fields are not updated
361        values = {k: v for k, v in values.items() if v is not UNSET}
362        keys = tuple(values.keys())
363        sql_query = f"UPDATE {table} SET {','.join(f'{x}=:{x}' for x in keys)} WHERE "
364        sql_query += " AND ".join(f"{x} = :{x}" for x in match)
365        await self.execute(sql_query, {**match, **values})
366        await self._maybe_commit()
367
368    async def delete(
369        self, table: str, match: dict[str, Any] | None = None, query: str | None = None
370    ) -> None:
371        """Delete data in given table."""
372        assert not (match and query), "Cannot use both match and query"
373        sql_query = f"DELETE FROM {table} "
374        if match:
375            sql_query += " WHERE " + " AND ".join(f"{x} = :{x}" for x in match)
376        elif query and "where" not in query.lower():
377            sql_query += "WHERE " + query
378        elif query:
379            sql_query += query
380        await self.execute(sql_query, match)
381        await self._maybe_commit()
382
383    async def delete_where_query(self, table: str, query: str | None = None) -> None:
384        """Delete data in given table using given where clausule."""
385        sql_query = f"DELETE FROM {table} WHERE {query}"
386        await self.execute(sql_query)
387        await self._maybe_commit()
388
389    async def execute(self, query: str, values: dict[str, Any] | None = None) -> Any:
390        """Execute command on the database."""
391        return await self._db.execute(query, values)
392
393    async def execute_write(self, query: str, values: dict[str, Any] | None = None) -> None:
394        """
395        Execute a hand-written write statement and commit it.
396
397        Use instead of `execute` for anything that modifies data, so the write is durable
398        even if nothing else happens to commit the shared connection afterwards. Honors
399        `deferred_commit`, so a batch still commits once at the end of its scope.
400
401        :param query: The statement to execute.
402        :param values: The values to bind to the statement's named parameters.
403        """
404        await self._db.execute(query, values)
405        await self._maybe_commit()
406
407    async def commit(self) -> None:
408        """Commit the current transaction."""
409        return await self._db.commit()
410
411    @asynccontextmanager
412    async def deferred_commit(self) -> AsyncGenerator[None]:
413        """
414        Batch all writes of the current task into a single commit when the scope exits.
415
416        Within the scope, the per-statement commit of the insert/upsert/update/delete
417        helpers is skipped for the current task and a single commit is issued when the
418        outermost scope exits (scopes may be nested). This greatly reduces the commit
419        overhead of multi-statement operations such as adding a media item with all
420        its relations to the library.
421
422        Note: this is not an atomic transaction. The scope always commits on exit -
423        also on error or cancellation - and never rolls back. Writes from other tasks
424        are unaffected and still commit immediately.
425        """
426        depth = self._deferred_commit_depth.get()
427        token = self._deferred_commit_depth.set(depth + 1)
428        try:
429            yield
430        finally:
431            self._deferred_commit_depth.reset(token)
432            # always commit on exit, never rollback: the connection is shared by all
433            # tasks, so statements from concurrent writers may interleave with this
434            # scope's statements in the same underlying SQLite transaction and a
435            # rollback would revert their (already acknowledged) writes as well
436            if depth == 0:
437                await self._db.commit()
438
439    async def iter_items(
440        self,
441        table: str,
442        match: dict[str, Any] | None = None,
443    ) -> AsyncGenerator[Mapping[str, Any]]:
444        """Iterate all items within a table."""
445        limit: int = 500
446        offset: int = 0
447        while True:
448            next_items = await self.get_rows(
449                table=table,
450                match=match,
451                offset=offset,
452                limit=limit,
453            )
454            for item in next_items:
455                yield item
456            if len(next_items) < limit:
457                break
458            await asyncio.sleep(0)  # yield to eventloop
459            offset += limit
460
461    async def get_reclaimable_ratio(self) -> float:
462        """
463        Return the fraction (0..1) of the database file that a VACUUM would reclaim.
464
465        This is the share of pages on the free list and is a cheap way to decide
466        whether a (potentially expensive) VACUUM is actually worthwhile.
467        """
468        page_count = await self._get_pragma_int("page_count")
469        if page_count <= 0:
470            return 0.0
471        freelist_count = await self._get_pragma_int("freelist_count")
472        return freelist_count / page_count
473
474    async def vacuum(self) -> None:
475        """Run vacuum command on database."""
476        # VACUUM rebuilds the whole database in temp storage; with temp_store=memory that
477        # copy lives entirely in RAM and OOMs memory constrained devices on large databases,
478        # so spill it to a temp file (located at SQLITE_TMPDIR) for the duration.
479        await self._db.execute("PRAGMA temp_store=FILE;")
480        try:
481            await self._db.execute("VACUUM")
482            await self._db.commit()
483        finally:
484            await self._db.execute("PRAGMA temp_store=memory;")
485
486    async def _get_pragma_int(self, pragma: str) -> int:
487        """Return the integer value of a single-value sqlite PRAGMA."""
488        async with self._db.execute(f"PRAGMA {pragma}") as cursor:
489            row = await cursor.fetchone()
490            return int(row[0]) if row else 0
491
492    async def _maybe_commit(self) -> None:
493        """Commit now, unless the current task is inside a deferred_commit scope."""
494        if self._deferred_commit_depth.get() == 0:
495            await self._db.commit()
496