From 09b0e6197b934fa9c3987c56f8ae467545d4501d Mon Sep 17 00:00:00 2001 From: Gang Tao Date: Sun, 16 Aug 2026 20:52:19 -0700 Subject: [PATCH] feat(db): put all tpk streams under a dedicated `tpk` database (#58) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every tpk stream (kg_nodes/kg_edges/kg_repos/kg_ingest_log/kg_users/ kg_roles/kg_sessions/chat_audit_log) is now created and queried under a dedicated database — `tpk` by default — instead of the server's `default` database. This keeps tpk's objects out of `default`, lets them be granted or dropped as a unit, and avoids collisions when tpk shares a Timeplus instance (e.g. the k8s app-only mode against an existing Timeplus Enterprise). - config: add `database()` (env TIMEPLUS_DATABASE > [db].database > "tpk"), validated as an identifier. - db: add `qualified(name, prefix)` -> `.`; ensure_schema runs `CREATE DATABASE IF NOT EXISTS` first; all DDL and drop_schema qualify names. get_client stays unscoped (connects to the always-present default DB) to avoid a bootstrap chicken-and-egg — connecting with database= fails if it doesn't exist yet. - Route every stream identifier through db.qualified(): tools, corpus, auth, audit, ingest, api, cli (status), transfer. TPK_STREAM_PREFIX still namespaces within the database. - tests: qualified()/database() precedence + updated backend DDL assertions. Verified end-to-end on live proton (create/upsert/read/delete/drop). - docs: README Configuration matrix + a `tpk` database section with the migration note (existing `default` data isn't moved — re-ingest or export/import, or set TIMEPLUS_DATABASE=default); repos.toml/.env.example. Migration: existing deployments' streams stay in `default`; after upgrade tpk reads from `tpk` (empty) until re-ingested or export/import'd. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01XPzYZXxyTdj5G25KoujpHb --- .env.example | 3 ++- README.md | 17 ++++++++++++++++ deploy/docker/repos.container.toml | 1 + repos.toml | 3 ++- src/tpk/api.py | 4 ++-- src/tpk/audit.py | 11 +++++++---- src/tpk/auth.py | 30 ++++++++++++++--------------- src/tpk/cli.py | 2 +- src/tpk/config.py | 14 ++++++++++++++ src/tpk/corpus.py | 14 +++++++------- src/tpk/db.py | 31 ++++++++++++++++++++++-------- src/tpk/ingest.py | 10 +++++----- src/tpk/tools.py | 8 ++++---- src/tpk/transfer.py | 10 +++++----- tests/test_backend.py | 15 +++++++++++++-- tests/test_config.py | 16 +++++++++++++++ 16 files changed, 134 insertions(+), 55 deletions(-) diff --git a/.env.example b/.env.example index 8db2951..027f114 100644 --- a/.env.example +++ b/.env.example @@ -62,8 +62,9 @@ OPENAI_API_KEY= # Every setting below ALSO has a config-file equivalent in repos.toml's # [db]/[agent]/[server] sections; the env var wins when both are set. See the # README "Configuration" matrix. Uncomment only to override the default shown. +# TIMEPLUS_DATABASE=tpk # [db].database database all tpk streams live under # TPK_DB_BACKEND=timeplusd # [db].backend timeplusd | proton -# TPK_STREAM_PREFIX= # [db].stream_prefix namespace all streams +# TPK_STREAM_PREFIX= # [db].stream_prefix namespace streams within the database # TPK_DB_WAIT_SECONDS=60 # [db].wait_seconds serve: DB connect retry budget # TPK_SESSION_TTL=86400 # [server].session_ttl login session lifetime (s) # TPK_CHAT_AUDIT=1 # [server].chat_audit false/0 disables chat auditing diff --git a/README.md b/README.md index 7b1db0e..a04f332 100644 --- a/README.md +++ b/README.md @@ -66,6 +66,7 @@ export the matching env var — whichever suits your deployment. Secrets are |---|---|---|---| | `TIMEPLUS_HOST` | `[db].host` | `localhost` | DB host | | `TIMEPLUS_USER` | `[db].user` | `default` | DB user | +| `TIMEPLUS_DATABASE` | `[db].database` | `tpk` | Database all tpk streams live under | | `TPK_DB_BACKEND` | `[db].backend` | `timeplusd` | Stream-semantics mode (`timeplusd`\|`proton`) | | `TPK_STREAM_PREFIX` | `[db].stream_prefix` | `` | Namespace prefix for all streams | | `TPK_DB_WAIT_SECONDS` | `[db].wait_seconds` | `60` | `serve`: DB connect retry budget | @@ -83,6 +84,22 @@ Secrets — **env-only**, never in the file: `TIMEPLUS_PASSWORD`, `ANTHROPIC_BASE_URL` / `OPENAI_BASE_URL` / `*_MODEL` endpoint variables the LLM clients read directly). +### The `tpk` database + +All tpk streams (`kg_nodes`, `kg_users`, `chat_audit_log`, …) live under a +dedicated database — **`tpk`** by default, set via `TIMEPLUS_DATABASE` / +`[db].database`. tpk creates it on startup (`CREATE DATABASE IF NOT EXISTS`) and +qualifies every stream as `tpk.`, so its objects don't clutter the server's +`default` database and can be granted or dropped as a unit. This matters most +when pointing tpk at a **shared** Timeplus Enterprise (the k8s app-only mode): +the connecting user needs `CREATE DATABASE` (first run) plus read/write on `tpk`. + +**Upgrading an existing deployment:** streams created before this change live in +`default` and are **not** moved automatically — after upgrading, tpk looks in +`tpk` and finds an empty graph. Either re-ingest (`tpk ingest`), or carry data +over with `tpk export` (old version) → `tpk import` (new version). Point +`TIMEPLUS_DATABASE=default` to keep using the old location instead. + ## Docker compose (quickest start) Two deployment modes are provided: diff --git a/deploy/docker/repos.container.toml b/deploy/docker/repos.container.toml index 3227728..319208a 100644 --- a/deploy/docker/repos.container.toml +++ b/deploy/docker/repos.container.toml @@ -31,6 +31,7 @@ token_budget = 16000 [db] # host = "localhost" # env: TIMEPLUS_HOST (DB+App image sets this to "db") # user = "default" # env: TIMEPLUS_USER (images set this to "tpk") +# database = "tpk" # env: TIMEPLUS_DATABASE (all tpk streams live here) # backend = "timeplusd" # env: TPK_DB_BACKEND (all-in-one image bakes "proton") # stream_prefix = "" # env: TPK_STREAM_PREFIX # wait_seconds = 60 # env: TPK_DB_WAIT_SECONDS diff --git a/repos.toml b/repos.toml index 0140337..edb1743 100644 --- a/repos.toml +++ b/repos.toml @@ -37,8 +37,9 @@ token_budget = 16000 [db] # host = "localhost" # env: TIMEPLUS_HOST # user = "default" # env: TIMEPLUS_USER +# database = "tpk" # env: TIMEPLUS_DATABASE (all tpk streams live here) # backend = "timeplusd" # env: TPK_DB_BACKEND ("timeplusd" | "proton") -# stream_prefix = "" # env: TPK_STREAM_PREFIX (namespaces all streams) +# stream_prefix = "" # env: TPK_STREAM_PREFIX (namespaces streams within the database) # wait_seconds = 60 # env: TPK_DB_WAIT_SECONDS (serve: DB connect retry budget) [agent] diff --git a/src/tpk/api.py b/src/tpk/api.py index f90683a..4f76624 100644 --- a/src/tpk/api.py +++ b/src/tpk/api.py @@ -263,13 +263,13 @@ def list_repos(actor: User = Depends(auth.require_cap(auth_mod.CAP_CORPUS_VIEW)) entries = corpus.list_entries(client, prefix=prefix) counts = dict( client.query( - f"SELECT repo, count() FROM {db.latest(f'{prefix}kg_nodes')} GROUP BY repo" + f"SELECT repo, count() FROM {db.latest(db.qualified('kg_nodes', prefix))} GROUP BY repo" ).result_rows ) last: dict[str, tuple] = {} for repo, sha, status, t in client.query( f"SELECT repo, arg_max(git_sha, _tp_time), arg_max(status, _tp_time)," - f" max(_tp_time) FROM table({prefix}kg_ingest_log) GROUP BY repo" + f" max(_tp_time) FROM table({db.qualified('kg_ingest_log', prefix)}) GROUP BY repo" ).result_rows: last[repo] = (sha, status, t) out = [] diff --git a/src/tpk/audit.py b/src/tpk/audit.py index 595b7c5..02f3491 100644 --- a/src/tpk/audit.py +++ b/src/tpk/audit.py @@ -1,5 +1,6 @@ """Chat support-history audit: one row per chat turn written to the -append-only ``{prefix}chat_audit_log`` stream (see db.ensure_schema). +append-only ``chat_audit_log`` stream (qualified as ``. +chat_audit_log`` via db.qualified; see db.ensure_schema). The sink is deliberately best-effort: auditing must never break or slow the chat response, so every write is wrapped and any failure is logged and @@ -12,6 +13,8 @@ from dataclasses import dataclass, field from datetime import datetime +from tpk import db + logger = logging.getLogger(__name__) # Insert column order for {prefix}chat_audit_log. Kept next to AuditRecord.to_row @@ -81,8 +84,8 @@ def to_row(self) -> list: def make_db_sink(prefix: str, client_factory): - """Return a ``sink(record: AuditRecord) -> None`` that inserts into - ``{prefix}chat_audit_log``. + """Return a ``sink(record: AuditRecord) -> None`` that inserts into the + ``chat_audit_log`` stream (qualified via db.qualified). ``client_factory`` is called per write to obtain a Timeplus client -- a fresh, single-use session avoids timeplus_connect's concurrent-query @@ -96,7 +99,7 @@ def sink(record: AuditRecord) -> None: try: client = client_factory() client.insert( - f"{prefix}chat_audit_log", + db.qualified("chat_audit_log", prefix), [record.to_row()], column_names=CHAT_AUDIT_COLUMNS, ) diff --git a/src/tpk/auth.py b/src/tpk/auth.py index 4d26b7b..d3000df 100644 --- a/src/tpk/auth.py +++ b/src/tpk/auth.py @@ -135,7 +135,7 @@ def _now(): def upsert_user(client, user: User, prefix: str = "") -> None: client.insert( - f"{prefix}kg_users", + db.qualified("kg_users", prefix), [[user.username, user.password_hash, user.role, user.must_change_password, user.disabled, _now(), _now()]], column_names=_USER_COLUMNS, @@ -145,7 +145,7 @@ def upsert_user(client, user: User, prefix: str = "") -> None: def get_user(client, username: str, prefix: str = "") -> User | None: rows = client.query( f"SELECT username, password_hash, role, must_change_password, disabled" - f" FROM {db.latest(f'{prefix}kg_users')} WHERE username = %(u)s", + f" FROM {db.latest(db.qualified('kg_users', prefix))} WHERE username = %(u)s", parameters={"u": username}, ).result_rows if not rows: @@ -157,19 +157,19 @@ def get_user(client, username: str, prefix: str = "") -> User | None: def list_users(client, prefix: str = "") -> list[User]: rows = client.query( f"SELECT username, password_hash, role, must_change_password, disabled" - f" FROM {db.latest(f'{prefix}kg_users')} ORDER BY username" + f" FROM {db.latest(db.qualified('kg_users', prefix))} ORDER BY username" ).result_rows return [User(u, h, r, bool(mc), bool(dis)) for u, h, r, mc, dis in rows] def delete_user(client, username: str, prefix: str = "") -> None: - db.delete(client, f"{prefix}kg_users", "username = %(u)s", + db.delete(client, db.qualified("kg_users", prefix), "username = %(u)s", {"u": username}, ("username",)) def admin_count(client, prefix: str = "") -> int: rows = client.query( - f"SELECT count() FROM {db.latest(f'{prefix}kg_users')}" + f"SELECT count() FROM {db.latest(db.qualified('kg_users', prefix))}" f" WHERE role = %(r)s AND NOT disabled", parameters={"r": ROLE_ADMIN}, ).result_rows @@ -192,7 +192,7 @@ def seed_admin(client, prefix: str = "") -> bool: def upsert_role(client, role: Role, prefix: str = "") -> None: client.insert( - f"{prefix}kg_roles", + db.qualified("kg_roles", prefix), [[role.name, json.dumps(role.entry_keys), json.dumps(role.capabilities), role.description, _now()]], column_names=_ROLE_COLUMNS, @@ -201,7 +201,7 @@ def upsert_role(client, role: Role, prefix: str = "") -> None: def get_role(client, name: str, prefix: str = "") -> Role | None: rows = client.query( - f"SELECT name, entry_keys, capabilities, description FROM {db.latest(f'{prefix}kg_roles')}" + f"SELECT name, entry_keys, capabilities, description FROM {db.latest(db.qualified('kg_roles', prefix))}" f" WHERE name = %(n)s", parameters={"n": name}, ).result_rows @@ -213,7 +213,7 @@ def get_role(client, name: str, prefix: str = "") -> Role | None: def list_roles(client, prefix: str = "") -> list[Role]: rows = client.query( - f"SELECT name, entry_keys, capabilities, description FROM {db.latest(f'{prefix}kg_roles')}" + f"SELECT name, entry_keys, capabilities, description FROM {db.latest(db.qualified('kg_roles', prefix))}" f" ORDER BY name" ).result_rows return [Role(n, json.loads(k) if k else [], d, _parse_capabilities(c)) @@ -221,12 +221,12 @@ def list_roles(client, prefix: str = "") -> list[Role]: def delete_role(client, name: str, prefix: str = "") -> None: - db.delete(client, f"{prefix}kg_roles", "name = %(n)s", {"n": name}, ("name",)) + db.delete(client, db.qualified("kg_roles", prefix), "name = %(n)s", {"n": name}, ("name",)) def usernames_with_role(client, name: str, prefix: str = "") -> list[str]: rows = client.query( - f"SELECT username FROM {db.latest(f'{prefix}kg_users')} WHERE role = %(r)s" + f"SELECT username FROM {db.latest(db.qualified('kg_users', prefix))} WHERE role = %(r)s" f" ORDER BY username", parameters={"r": name}, ).result_rows @@ -242,7 +242,7 @@ def _token_hash(token: str) -> str: def create_session(client, username: str, ttl_seconds: int, prefix: str = "") -> str: token = secrets.token_urlsafe(32) client.insert( - f"{prefix}kg_sessions", + db.qualified("kg_sessions", prefix), [[_token_hash(token), username, _now() + timedelta(seconds=ttl_seconds), _now()]], column_names=_SESSION_COLUMNS, @@ -253,7 +253,7 @@ def create_session(client, username: str, ttl_seconds: int, prefix: str = "") -> def get_session(client, token: str, prefix: str = "") -> str | None: h = _token_hash(token) rows = client.query( - f"SELECT username, expires_at FROM {db.latest(f'{prefix}kg_sessions')}" + f"SELECT username, expires_at FROM {db.latest(db.qualified('kg_sessions', prefix))}" f" WHERE token_hash = %(h)s", parameters={"h": h}, ).result_rows @@ -261,14 +261,14 @@ def get_session(client, token: str, prefix: str = "") -> str | None: return None username, expires_at = rows[0] if expires_at.replace(tzinfo=timezone.utc) < _now(): - db.delete(client, f"{prefix}kg_sessions", "token_hash = %(h)s", + db.delete(client, db.qualified("kg_sessions", prefix), "token_hash = %(h)s", {"h": h}, ("token_hash",)) return None return username def delete_session(client, token: str, prefix: str = "") -> None: - db.delete(client, f"{prefix}kg_sessions", "token_hash = %(h)s", + db.delete(client, db.qualified("kg_sessions", prefix), "token_hash = %(h)s", {"h": _token_hash(token)}, ("token_hash",)) @@ -279,7 +279,7 @@ def delete_user_sessions(client, username: str, prefix: str = "", if keep_token is not None: where += " AND token_hash != %(k)s" params["k"] = _token_hash(keep_token) - db.delete(client, f"{prefix}kg_sessions", where, params, ("token_hash",)) + db.delete(client, db.qualified("kg_sessions", prefix), where, params, ("token_hash",)) # -- HTTP layer ------------------------------------------------------------ diff --git a/src/tpk/cli.py b/src/tpk/cli.py index cbb0c71..8caea69 100644 --- a/src/tpk/cli.py +++ b/src/tpk/cli.py @@ -59,7 +59,7 @@ def status(): rows = client.query( "SELECT repo, max(_tp_time) AS last_run, arg_max(status, _tp_time) AS status," " arg_max(nodes, _tp_time) AS nodes, arg_max(edges, _tp_time) AS edges" - " FROM table(kg_ingest_log) GROUP BY repo ORDER BY repo" + f" FROM table({db.qualified('kg_ingest_log')}) GROUP BY repo ORDER BY repo" ).result_rows for repo, last_run, status_, nodes, edges in rows: typer.echo(f"{repo:35s} {status_:7s} {nodes:>8} nodes {edges:>8} edges {last_run}") diff --git a/src/tpk/config.py b/src/tpk/config.py index 13008dd..4051050 100644 --- a/src/tpk/config.py +++ b/src/tpk/config.py @@ -96,6 +96,20 @@ def db_backend() -> str: return backend +def database() -> str: + """The database all tpk streams live under (env > [db].database > 'tpk'). + tpk qualifies every stream as `.` (see #58) so its + objects don't clutter the server's `default` database and can be granted / + dropped as a unit. Read module-level (like db_backend) so db.py helpers can + qualify names without threading it through every signature.""" + name = setting("TIMEPLUS_DATABASE", "db", "database", "tpk") + if not name.replace("_", "").isalnum(): + raise ValueError( + f"TIMEPLUS_DATABASE / [db].database must be alphanumeric/underscore, got {name!r}" + ) + return name + + @dataclass(frozen=True) class RepoConfig: name: str diff --git a/src/tpk/corpus.py b/src/tpk/corpus.py index 1944618..0b0aba8 100644 --- a/src/tpk/corpus.py +++ b/src/tpk/corpus.py @@ -32,12 +32,12 @@ def _to_cfg(row) -> RepoConfig: def upsert_entry(client, cfg: RepoConfig, prefix: str = "") -> None: - client.insert(f"{prefix}kg_repos", [_row(cfg)], column_names=_COLUMNS) + client.insert(db.qualified("kg_repos", prefix), [_row(cfg)], column_names=_COLUMNS) def list_entries(client, prefix: str = "") -> list[RepoConfig]: rows = client.query( - f"SELECT {', '.join(_COLUMNS)} FROM {db.latest(f'{prefix}kg_repos')}" + f"SELECT {', '.join(_COLUMNS)} FROM {db.latest(db.qualified('kg_repos', prefix))}" " ORDER BY name, ref" ).result_rows return [_to_cfg(r) for r in rows] @@ -45,7 +45,7 @@ def list_entries(client, prefix: str = "") -> list[RepoConfig]: def find_entry(client, name: str, ref: str, prefix: str = "") -> RepoConfig | None: rows = client.query( - f"SELECT {', '.join(_COLUMNS)} FROM {db.latest(f'{prefix}kg_repos')}" + f"SELECT {', '.join(_COLUMNS)} FROM {db.latest(db.qualified('kg_repos', prefix))}" " WHERE name = %(n)s AND ref = %(r)s", parameters={"n": name, "r": ref}, ).result_rows @@ -66,19 +66,19 @@ def delete_entry( cfg = find_entry(client, name, ref, prefix=prefix) if cfg is None: return False - db.delete(client, f"{prefix}kg_repos", "name = %(n)s AND ref = %(r)s", + db.delete(client, db.qualified("kg_repos", prefix), "name = %(n)s AND ref = %(r)s", {"n": name, "r": ref}, ("name", "ref")) if purge: key = entry_key(cfg) - db.delete(client, f"{prefix}kg_nodes", "repo = %(k)s", {"k": key}, ("id",)) - db.delete(client, f"{prefix}kg_edges", "repo = %(k)s", {"k": key}, + db.delete(client, db.qualified("kg_nodes", prefix), "repo = %(k)s", {"k": key}, ("id",)) + db.delete(client, db.qualified("kg_edges", prefix), "repo = %(k)s", {"k": key}, ("src", "dst", "rel")) return True def enabled_keys(client, prefix: str = "") -> list[str]: rows = client.query( - f"SELECT name, ref FROM {db.latest(f'{prefix}kg_repos')} WHERE enabled" + f"SELECT name, ref FROM {db.latest(db.qualified('kg_repos', prefix))} WHERE enabled" " ORDER BY name, ref" ).result_rows return [f"{n}@{r}" if r else n for n, r in rows] diff --git a/src/tpk/db.py b/src/tpk/db.py index 60d5569..9a84634 100644 --- a/src/tpk/db.py +++ b/src/tpk/db.py @@ -23,23 +23,33 @@ import timeplus_connect -from tpk.config import Settings, db_backend +from tpk.config import Settings, database, db_backend logger = logging.getLogger(__name__) +def qualified(name: str, prefix: str = "") -> str: + """Fully-qualified identifier for a tpk stream: `.`. + + All tpk streams live under the configured database (default `tpk`, #58), so + every DDL/DML/read site builds its stream name through here. `prefix` + (TPK_STREAM_PREFIX) still namespaces within the database, e.g. for tests.""" + return f"{database()}.{prefix}{name}" + + def _keyed_stream(prefix: str, name: str, columns: list[str], pk: str) -> str: """DDL for a keyed (upsertable, latest-state) stream, per backend.""" cols = ",\n ".join(columns) + stream = qualified(name, prefix) if db_backend() == "proton": return ( - f"CREATE STREAM IF NOT EXISTS {prefix}{name} (\n" + f"CREATE STREAM IF NOT EXISTS {stream} (\n" f" {cols},\n" f" deleted uint8 DEFAULT 0\n" f" ) PRIMARY KEY ({pk}) SETTINGS mode='versioned_kv'" ) return ( - f"CREATE MUTABLE STREAM IF NOT EXISTS {prefix}{name} (\n" + f"CREATE MUTABLE STREAM IF NOT EXISTS {stream} (\n" f" {cols}\n" f" ) PRIMARY KEY ({pk})" ) @@ -124,6 +134,11 @@ def _add_column_if_missing(client, stream: str, column: str, col_type: str) -> N def ensure_schema(client, prefix: str = "") -> None: + # All tpk streams live under a dedicated database (default `tpk`, #58); + # create it first so the qualified DDL below resolves. CREATE DATABASE is a + # global statement (works from the default-database session get_client uses) + # and is a no-op if it already exists, on both timeplusd and proton. + client.command(f"CREATE DATABASE IF NOT EXISTS {database()}") # Keyed (upsertable, latest-state) streams -- MUTABLE on timeplusd, # versioned_kv + `deleted` tombstone on proton (see _keyed_stream). client.command(_keyed_stream(prefix, "kg_nodes", [ @@ -137,7 +152,7 @@ def ensure_schema(client, prefix: str = "") -> None: "repo string", "updated_at datetime64(3, 'UTC')", ], pk="src, dst, rel")) client.command(f""" - CREATE STREAM IF NOT EXISTS {prefix}kg_ingest_log ( + CREATE STREAM IF NOT EXISTS {qualified('kg_ingest_log', prefix)} ( repo string, run_id string, nodes uint64, @@ -164,7 +179,7 @@ def ensure_schema(client, prefix: str = "") -> None: # column (CREATE ... IF NOT EXISTS won't add it to an existing stream). # Existing role rows keep an empty cell, which auth._parse_capabilities # reads as the chat-only migration default. - _add_column_if_missing(client, f"{prefix}kg_roles", "capabilities", "string") + _add_column_if_missing(client, qualified("kg_roles", prefix), "capabilities", "string") client.command(_keyed_stream(prefix, "kg_sessions", [ "token_hash string", "username string", "expires_at datetime64(3, 'UTC')", "created_at datetime64(3, 'UTC')", @@ -173,7 +188,7 @@ def ensure_schema(client, prefix: str = "") -> None: # answer, and the tool calls made). Not MUTABLE -- like kg_ingest_log, we # want the full event history, not a keyed latest-state view. client.command(f""" - CREATE STREAM IF NOT EXISTS {prefix}chat_audit_log ( + CREATE STREAM IF NOT EXISTS {qualified('chat_audit_log', prefix)} ( ts datetime64(3, 'UTC'), turn_id string, conversation_id string, @@ -200,6 +215,6 @@ def drop_schema(client, prefix: str) -> None: raise ValueError("refusing to drop unprefixed (production) streams") for name in ( "kg_nodes", "kg_edges", "kg_ingest_log", "kg_repos", - "kg_users", "kg_roles", "kg_sessions", + "kg_users", "kg_roles", "kg_sessions", "chat_audit_log", ): - client.command(f"DROP STREAM IF EXISTS {prefix}{name}") + client.command(f"DROP STREAM IF EXISTS {qualified(name, prefix)}") diff --git a/src/tpk/ingest.py b/src/tpk/ingest.py index 0bf4a59..6e26fde 100644 --- a/src/tpk/ingest.py +++ b/src/tpk/ingest.py @@ -31,9 +31,9 @@ def upsert_graph( repos: set[str] | None = None, ) -> None: if nodes: - client.insert(f"{prefix}kg_nodes", node_rows(nodes, run_started_at), column_names=NODE_COLUMNS) + client.insert(db.qualified("kg_nodes", prefix), node_rows(nodes, run_started_at), column_names=NODE_COLUMNS) if edges: - client.insert(f"{prefix}kg_edges", edge_rows(edges, run_started_at), column_names=EDGE_COLUMNS) + client.insert(db.qualified("kg_edges", prefix), edge_rows(edges, run_started_at), column_names=EDGE_COLUMNS) if repos is None: # Derived from the batch: callers seeding multi-repo test data rely on # this. Note this means a repo whose parse yielded zero nodes/edges @@ -44,8 +44,8 @@ def upsert_graph( for repo in repos: params = {"r": repo, "t": run_started_at} where = "repo = %(r)s AND updated_at < %(t)s" - db.delete(client, f"{prefix}kg_nodes", where, params, ("id",)) - db.delete(client, f"{prefix}kg_edges", where, params, ("src", "dst", "rel")) + db.delete(client, db.qualified("kg_nodes", prefix), where, params, ("id",)) + db.delete(client, db.qualified("kg_edges", prefix), where, params, ("src", "dst", "rel")) def _git_sha(repo_path: Path) -> str: @@ -60,7 +60,7 @@ def _git_sha(repo_path: Path) -> str: def _log(client, prefix: str, result: IngestResult, git_sha: str) -> None: client.insert( - f"{prefix}kg_ingest_log", + db.qualified("kg_ingest_log", prefix), [[result.repo, result.run_id, result.nodes, result.edges, git_sha, result.status]], column_names=["repo", "run_id", "nodes", "edges", "git_sha", "status"], ) diff --git a/src/tpk/tools.py b/src/tpk/tools.py index a4e9efb..7f5b523 100644 --- a/src/tpk/tools.py +++ b/src/tpk/tools.py @@ -141,7 +141,7 @@ def _nodes_by_ids(self, ids: list[str]) -> list[dict]: clauses.append("repo IN %(active_repos)s") params["active_repos"] = active or ["__none__"] rows = self._query_rows( - f"SELECT {', '.join(NODE_FIELDS)} FROM {db.latest(f'{self.prefix}kg_nodes')}" + f"SELECT {', '.join(NODE_FIELDS)} FROM {db.latest(db.qualified('kg_nodes', self.prefix))}" f" WHERE {' AND '.join(clauses)}", parameters=params, ) @@ -166,7 +166,7 @@ def _edges_touching(self, ids: list[str], rels, direction: str, confidence) -> l clauses.append("repo IN %(active_repos)s") params["active_repos"] = active or ["__none__"] rows = self._query_rows( - f"SELECT {', '.join(EDGE_FIELDS)} FROM {db.latest(f'{self.prefix}kg_edges')}" + f"SELECT {', '.join(EDGE_FIELDS)} FROM {db.latest(db.qualified('kg_edges', self.prefix))}" f" WHERE {' AND '.join(clauses)}", parameters=params, ) @@ -206,7 +206,7 @@ def search_entities(self, query, kinds=None, repos=None, limit=20): def _run(where_clauses, order): return self._query_rows( - f"SELECT {', '.join(NODE_FIELDS)} FROM {db.latest(f'{self.prefix}kg_nodes')}" + f"SELECT {', '.join(NODE_FIELDS)} FROM {db.latest(db.qualified('kg_nodes', self.prefix))}" f" WHERE {' AND '.join(where_clauses)}" f" ORDER BY {order} LIMIT %(limit)s", parameters=params, @@ -329,7 +329,7 @@ def list_communities(self, repo=None): clause = f" WHERE {' AND '.join(clauses)}" if clauses else "" rows = self._query_rows( f"SELECT repo, community, count() AS node_count" - f" FROM {db.latest(f'{self.prefix}kg_nodes')}{clause}" + f" FROM {db.latest(db.qualified('kg_nodes', self.prefix))}{clause}" " GROUP BY repo, community ORDER BY node_count DESC", parameters=params, ) diff --git a/src/tpk/transfer.py b/src/tpk/transfer.py index 231b234..1c83b7e 100644 --- a/src/tpk/transfer.py +++ b/src/tpk/transfer.py @@ -28,7 +28,7 @@ def _stream_columns(client, stream: str, prefix: str = "") -> list[str]: """Real, insertable columns of a stream — excludes proton's internal `_tp_*` pseudo-columns and any ALIAS/MATERIALIZED columns.""" - rows = client.query(f"DESCRIBE {prefix}{stream}").result_rows + rows = client.query(f"DESCRIBE {db.qualified(stream, prefix)}").result_rows cols = [] for row in rows: name = row[0] @@ -40,7 +40,7 @@ def _stream_columns(client, stream: str, prefix: str = "") -> list[str]: def _row_count(client, stream: str, prefix: str = "") -> int: - r = client.query(f"SELECT count() FROM table({prefix}{stream})").result_rows + r = client.query(f"SELECT count() FROM table({db.qualified(stream, prefix)})").result_rows return int(r[0][0]) if r else 0 @@ -65,7 +65,7 @@ def export_bundle(client, out_dir, streams=None, prefix: str = "", # object; copy it out in chunks so a 300k-row stream never lands # wholly in memory. src = client.raw_stream( - f"SELECT {', '.join(cols)} FROM table({prefix}{stream})", fmt=fmt) + f"SELECT {', '.join(cols)} FROM table({db.qualified(stream, prefix)})", fmt=fmt) with open(out / fname, "wb") as fh: shutil.copyfileobj(src, fh) manifest["streams"][stream] = { @@ -120,10 +120,10 @@ def import_bundle(client, in_dir, prefix: str = "", replace: bool = False) -> di # follows. Done per stream, immediately before its load, so a # mid-run failure only affects the stream in flight rather than # leaving every earlier-dropped stream empty. - client.command(f"DROP STREAM IF EXISTS {prefix}{stream}") + client.command(f"DROP STREAM IF EXISTS {db.qualified(stream, prefix)}") db.ensure_schema(client, prefix) with open(src / info["file"], "rb") as fh: - client.raw_insert(table=f"{prefix}{stream}", + client.raw_insert(table=db.qualified(stream, prefix), column_names=info["columns"], insert_block=fh, fmt=fmt) return manifest diff --git a/tests/test_backend.py b/tests/test_backend.py index f68f69e..6102888 100644 --- a/tests/test_backend.py +++ b/tests/test_backend.py @@ -14,8 +14,10 @@ def test_keyed_stream_ddl_timeplusd(monkeypatch): monkeypatch.setenv("TPK_DB_BACKEND", "timeplusd") + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) ddl = db._keyed_stream("p_", "kg_users", ["username string", "role string"], pk="username") - assert "CREATE MUTABLE STREAM IF NOT EXISTS p_kg_users" in ddl + # streams live under the `tpk` database, qualified as . (#58) + assert "CREATE MUTABLE STREAM IF NOT EXISTS tpk.p_kg_users" in ddl assert "PRIMARY KEY (username)" in ddl assert "versioned_kv" not in ddl assert "deleted" not in ddl @@ -23,13 +25,22 @@ def test_keyed_stream_ddl_timeplusd(monkeypatch): def test_keyed_stream_ddl_proton(monkeypatch): monkeypatch.setenv("TPK_DB_BACKEND", "proton") + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) ddl = db._keyed_stream("p_", "kg_users", ["username string", "role string"], pk="username") - assert "CREATE STREAM IF NOT EXISTS p_kg_users" in ddl + assert "CREATE STREAM IF NOT EXISTS tpk.p_kg_users" in ddl assert "MUTABLE" not in ddl assert "deleted uint8 DEFAULT 0" in ddl assert "SETTINGS mode='versioned_kv'" in ddl +def test_qualified_name(monkeypatch): + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) + assert db.qualified("kg_nodes") == "tpk.kg_nodes" + assert db.qualified("kg_nodes", "p_") == "tpk.p_kg_nodes" + monkeypatch.setenv("TIMEPLUS_DATABASE", "kb") + assert db.qualified("kg_nodes", "p_") == "kb.p_kg_nodes" + + def test_latest_timeplusd(monkeypatch): monkeypatch.setenv("TPK_DB_BACKEND", "timeplusd") assert db.latest("p_kg_users") == "table(p_kg_users)" diff --git a/tests/test_config.py b/tests/test_config.py index 9289790..b55c288 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -7,6 +7,7 @@ Settings, as_bool, checkout_root, + database, db_backend, load_llm, load_repos, @@ -121,6 +122,21 @@ def test_db_backend_rejects_bad_value(config_file, monkeypatch): db_backend() +def test_database_precedence(config_file, monkeypatch): + monkeypatch.delenv("TIMEPLUS_DATABASE", raising=False) + assert database() == "tpk" # built-in default + config_file.write_text('[db]\ndatabase = "kb"\n') + assert database() == "kb" # file over default + monkeypatch.setenv("TIMEPLUS_DATABASE", "graph") + assert database() == "graph" # env over file + + +def test_database_rejects_bad_value(config_file, monkeypatch): + monkeypatch.setenv("TIMEPLUS_DATABASE", "bad-name;drop") + with pytest.raises(ValueError): + database() + + def test_settings_reads_db_section_from_file(config_file, monkeypatch): for var in ("TIMEPLUS_HOST", "TIMEPLUS_USER", "TPK_STREAM_PREFIX", "TPK_DB_BACKEND"): monkeypatch.delenv(var, raising=False)