8000 fix(duckdb_pool): apply MINILAKE_DUCKDB_MEMORY_LIMIT to every DuckDB connection by avison9 · Pull Request #68 · dmux/minilake · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ Everything is set through environment variables. Nothing needs a config file.
| `MINILAKE_DATA_DIR` | `./data` | Where catalogs, workspace files, DBFS and Delta data live |
| `MINILAKE_VERBOSE` | unset | `1` for full `INFO` logs (per-request access logs, SQL and job traces). Off by default — only the startup banner and warnings appear |
| `MINILAKE_SERVICES` | (all) | Comma-separated allowlist, e.g. `unity_catalog,sql_statements,sql_warehouses` |
| `MINILAKE_DUCKDB_MEMORY_LIMIT` | `4GB` | Declared but not yet applied to any connection — currently a no-op |
| `MINILAKE_DUCKDB_MEMORY_LIMIT` | `4GB` | DuckDB `memory_limit`, applied to every connection. Each database (the shared Unity Catalog one and each warehouse) gets its own limit. An invalid value stops startup |
| `MINILAKE_STATEMENT_CACHE_SIZE` | `500` | How many executed statements to keep. Backs both `GET /statements/{id}` and query history; the oldest lose their rows first, then their metadata |
| `MINILAKE_DEV_CORS` | unset | `1` allows cross-origin API access, for running the web UI's Next.js dev server on another port. Leave off otherwise — minilake has no authentication ([Web UI](ui.md#development)) |
| `MINILAKE_DUCKDB_EXTENSION_DIR` | unset (`/opt/duckdb-extensions` in the image) | Directory of pre-installed DuckDB extensions. When set, the server only `LOAD`s the `delta` extension and never downloads it, which is what makes the image work with no internet access |
Expand Down
31 changes: 22 additions & 9 deletions src/minilake/duckdb_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,23 @@ class DuckDBPool:
- All connections are serialized with asyncio.Lock to respect DuckDB's single-writer model
"""

def __init__(self, data_dir: Path, extension_dir: Optional[Path] = None):
def __init__(
self,
data_dir: Path,
extension_dir: Optional[Path] = None,
memory_limit: Optional[str] = None,
):
self.data_dir = Path(data_dir)
self.data_dir.mkdir(parents=True, exist_ok=True)

# Where DuckDB looks for extensions. Set by the Docker image to a directory populated
# at build time; None for a plain pip install, where extensions are downloaded on demand.
self.extension_dir = extension_dir if extension_dir is not None else settings.duckdb_extension_dir

# DuckDB's `memory_limit`, e.g. "4GB". It is a per-database setting, so every
# connection this pool opens (the shared UC one and each warehouse) gets its own cap.
self.memory_limit = memory_limit if memory_limit is not None else settings.duckdb_memory_limit

# Shared Unity Catalog connection
self._uc_connection: Optional[duckdb.DuckDBPyConnection] = None
self._uc_lock = asyncio.Lock()
Expand All @@ -40,19 +49,23 @@ def __init__(self, data_dir: Path, extension_dir: Optional[Path] = None):
self._warehouse_dict_lock = asyncio.Lock()

def _connect_config(self) -> dict:
"""Connection config pinning DuckDB to the pre-installed extension directory.
"""Config shared by every connection: the memory limit, plus the pre-installed
extension directory when there is one.

Empty when there is none, which leaves DuckDB on its default `$HOME/.duckdb`.
Without an extension directory DuckDB stays on its default `$HOME/.duckdb`.
`autoinstall_known_extensions=False` is the part that guarantees offline operation:
without it DuckDB silently reaches for extensions.duckdb.org the first time a query
touches a function from an extension that isn't loaded yet.

An invalid `memory_limit` makes `duckdb.connect` raise. The UC connection is opened
during startup, so a bad MINILAKE_DUCKDB_MEMORY_LIMIT stops the server instead of
being ignored.
"""
if self.extension_dir is None:
return {}
return {
"extension_directory": str(self.extension_dir),
"autoinstall_known_extensions": False,
}
config: dict = {"memory_limit": self.memory_limit}
if self.extension_dir is not None:
config["extension_directory"] = str(self.extension_dir)
config["autoinstall_known_extensions"] = False
return config

async def _get_uc_connection(self) -> duckdb.DuckDBPyConnection:
"""Get or create the shared Unity Catalog DuckDB connection."""
Expand Down
63 changes: 63 additions & 0 deletions tests/test_duckdb_memory_limit.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
"""MINILAKE_DUCKDB_MEMORY_LIMIT reaches DuckDB.

The setting used to be parsed and then never passed to any connection, so DuckDB fell back
to its own default (80% of RAM). These tests read the limit back from the engine itself,
because a test of the parsed setting alone is exactly what let that slip through.
"""

from pathlib import Path

import duckdb
import pytest
from databricks.sdk import WorkspaceClient

from minilake.config import settings
from minilake.duckdb_pool import DuckDBPool

MEMORY_LIMIT_SQL = "SELECT current_setting('memory_limit')"


def _reported_limit(limit: str) -> str:
"""How DuckDB displays `limit`, e.g. "4GB" -> "3.7 GiB"."""
with duckdb.connect(":memory:", config={"memory_limit": limit}) as conn:
return conn.execute(MEMORY_LIMIT_SQL).fetchone()[0]


@pytest.mark.crud
def test_statement_execution_uses_configured_memory_limit(workspace_client: WorkspaceClient):
"""A query through the Statement Execution API sees the configured limit."""
wh = workspace_client.warehouses.create(name="memory_limit_wh")
result = workspace_client.statement_execution.execute_statement(
warehouse_id=wh.id,
statement=MEMORY_LIMIT_SQL,
)

assert result.result.data_array[0][0] == _reported_limit(settings.duckdb_memory_limit)


@pytest.mark.crud
@pytest.mark.asyncio
async def test_every_pool_connection_gets_the_memory_limit(tmp_path: Path):
"""Both connection paths, the shared UC one and a per-warehouse one, apply the limit."""
pool = DuckDBPool(tmp_path, memory_limit="16MB")
try:
uc = await pool.get_uc_connection()
wh = await pool.get_warehouse_connection("wh1")

expected = _reported_limit("16MB")
assert uc.execute(MEMORY_LIMIT_SQL).fetchone()[0] == expected
assert wh.execute(MEMORY_LIMIT_SQL).fetchone()[0] == expected
finally:
await pool.close_all()


@pytest.mark.error
@pytest.mark.asyncio
async def test_invalid_memory_limit_fails_loudly(tmp_path: Path):
"""A malformed limit raises on connect instead of being silently ignored."""
pool = DuckDBPool(tmp_path, memory_limit="lots")
try:
with pytest.raises(duckdb.Error):
await pool.get_uc_connection()
finally:
await pool.close_all()
Loading
0