Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
39 changes: 39 additions & 0 deletions taosmd/_db.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
"""Shared SQLite connection helper.

Every persistent store in the package opens its own SQLite database. By
default ``sqlite3.connect`` uses the rollback-journal mode, which takes an
exclusive lock for the duration of a write and lets only a single writer
touch the file at a time. When several agents (or several processes on the
same machine) share a memory store, that serialisation surfaces as
``SQLITE_BUSY`` errors.

``connect`` centralises the fix: it enables write-ahead logging (WAL), which
lets readers proceed concurrently with a writer, and sets a busy timeout so a
contended writer waits and retries rather than failing immediately. Both are
plain PRAGMAs with no extra dependencies; standalone behaviour is unchanged
apart from the journal mode.
"""

from __future__ import annotations

import sqlite3
from pathlib import Path
from typing import Union

# Allow a contended connection to block-and-retry for this many milliseconds
# before raising ``sqlite3.OperationalError: database is locked``.
BUSY_TIMEOUT_MS = 5000


def connect(db_path: Union[str, Path]) -> sqlite3.Connection:
"""Open a SQLite connection in WAL mode with a busy timeout.

Drop-in replacement for ``sqlite3.connect(db_path)``. Callers that need a
``row_factory`` or other connection attributes should set them on the

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SUGGESTION: Hardcoded busy timeout (5000ms) should be configurable

Consider making this a parameter with a default value so callers can adjust based on their contention expectations:

Suggested change
``row_factory`` or other connection attributes should set them on the
BUSY_TIMEOUT_MS = 5000
def connect(db_path: Union[str, Path], busy_timeout_ms: int = BUSY_TIMEOUT_MS) -> sqlite3.Connection:

Then update the PRAGMA line to use the parameter.

returned connection as before.
"""
conn = sqlite3.connect(db_path)
# WAL is harmless (silently stays "memory") on ``:memory:`` databases.
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(f"PRAGMA busy_timeout={BUSY_TIMEOUT_MS}")

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: WAL mode PRAGMA result not checked

PRAGMA journal_mode=WAL returns the actual journal mode. On some filesystems (network drives, some FUSE mounts) WAL cannot be enabled and it silently falls back to DELETE mode. Check the result to ensure WAL was actually enabled:

result = conn.execute("PRAGMA journal_mode=WAL").fetchone()
if result[0] != "wal":
    logger.warning(f"WAL mode not enabled, got: {result[0]}")

return conn

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SUGGESTION: Use parameter binding for PRAGMA busy_timeout

While the value is a constant here, using parameter binding is safer practice and avoids any injection concerns if this becomes configurable:

Suggested change
return conn
conn.execute("PRAGMA busy_timeout=?", (busy_timeout_ms,))

4 changes: 3 additions & 1 deletion taosmd/access_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
import time
from pathlib import Path

from . import _db

SCHEMA = """
CREATE TABLE IF NOT EXISTS access_log (
memory_key TEXT NOT NULL,
Expand Down Expand Up @@ -49,7 +51,7 @@ def __init__(self, db_path: str | Path = "data/access-tracker.db"):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/archive.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
from pathlib import Path
from typing import Any

from . import _db

logger = logging.getLogger(__name__)

# Event types
Expand Down Expand Up @@ -95,7 +97,7 @@ def __init__(
async def init(self) -> None:
self._archive_dir.mkdir(parents=True, exist_ok=True)
Path(self._index_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._index_path)
self._conn = _db.connect(self._index_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(INDEX_SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/browsing_history.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
import time
from pathlib import Path

from . import _db

SCHEMA = """
CREATE TABLE IF NOT EXISTS browsing_history (
url TEXT PRIMARY KEY,
Expand All @@ -34,7 +36,7 @@ def __init__(self, db_path: str | Path = "data/browsing-history.db"):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/crystallize.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
from pathlib import Path
from typing import TYPE_CHECKING

from . import _db

if TYPE_CHECKING:
from .knowledge_graph import TemporalKnowledgeGraph

Expand Down Expand Up @@ -67,7 +69,7 @@ def __init__(self, db_path: str | Path = "data/crystals.db"):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/knowledge_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
import time
from pathlib import Path

from . import _db

SCHEMA = """
CREATE TABLE IF NOT EXISTS kg_entities (
id TEXT PRIMARY KEY,
Expand Down Expand Up @@ -80,7 +82,7 @@ def __init__(self, db_path: str | Path = "data/knowledge-graph.db"):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/pending_decisions.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,8 @@
from pathlib import Path
from typing import Any

from . import _db


SCHEMA = """
CREATE TABLE IF NOT EXISTS kg_pending_decisions (
Expand Down Expand Up @@ -87,7 +89,7 @@ def __init__(self, db_path: str | Path):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 3 additions & 1 deletion taosmd/reflect.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
from pathlib import Path
from typing import TYPE_CHECKING

from . import _db

if TYPE_CHECKING:
from .knowledge_graph import TemporalKnowledgeGraph

Expand Down Expand Up @@ -134,7 +136,7 @@ def __init__(self, db_path: str | Path = "data/insights.db"):

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
4 changes: 2 additions & 2 deletions taosmd/session_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

import httpx

from taosmd import prompts
from taosmd import _db, prompts

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -143,7 +143,7 @@ def __init__(

async def init(self) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down
6 changes: 3 additions & 3 deletions taosmd/taosmd_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from pathlib import Path
from typing import Any

from . import __version__
from . import __version__, _db
from .backend import MemoryBackend

SETTINGS_SCHEMA = """
Expand Down Expand Up @@ -82,7 +82,7 @@ def __init__(
async def init(self) -> None:
"""Initialise settings database."""
Path(self._settings_db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._settings_db_path)
self._conn = _db.connect(self._settings_db_path)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WARNING: _ensure_conn creates a new connection without closing the old one

If _ensure_conn is called multiple times (e.g. after a connection error), the previous connection is leaked. Consider closing the existing connection first:

if self._conn is not None:
    self._conn.close()

This is a pre-existing issue but worth fixing since you're touching this code.

self._conn.row_factory = sqlite3.Row
self._conn.executescript(SETTINGS_SCHEMA + AGENT_CONFIG_SCHEMA)
self._conn.commit()
Expand Down Expand Up @@ -260,7 +260,7 @@ def _ensure_conn(self) -> None:
"""Ensure the settings DB connection is open (auto-init if needed)."""
if self._conn is None:
Path(self._settings_db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._settings_db_path)
self._conn = _db.connect(self._settings_db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SETTINGS_SCHEMA + AGENT_CONFIG_SCHEMA)
self._conn.commit()
4 changes: 3 additions & 1 deletion taosmd/vector_memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
import time
from pathlib import Path

from . import _db

logger = logging.getLogger(__name__)

SCHEMA = """
Expand Down Expand Up @@ -90,7 +92,7 @@ def __init__(

async def init(self, http_client=None) -> None:
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
self._conn = sqlite3.connect(self._db_path)
self._conn = _db.connect(self._db_path)
self._conn.row_factory = sqlite3.Row
self._conn.executescript(SCHEMA)
self._conn.commit()
Expand Down