diff --git a/CHANGELOG.md b/CHANGELOG.md index c814e42..b3b6789 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- `resolve_database_by_id` — fetches a managed database record by id (`GET /databases/{id}`) + with no by-name fallback, and returns an already-resolved `ManagedDatabase` untouched. + `ManagedDatabase` is re-exported for callers that hold one. + - Full-text search tool backed by the engine's BM25 index. `make_hotdata_tools` grows `search_table`/`search_column`/`search_columns`/`search_k`/`search_tool_name` and appends a `hotdata_search_text` tool when a table and column are given; `make_hotdata_search_tool` @@ -30,6 +34,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- **Breaking: managed databases are addressed by id, never by name.** A Hotdata database name + is a display label and is not unique, so a by-name lookup can resolve to the wrong database + — and then every query, load and drop follows it there. The agent-facing + `hotdata_load_managed_table` made that reachable from an LLM, where a wrong target means a + replacing load overwrites another database's table. + - `make_hotdata_tools`, `make_hotdata_search_tool` and `make_hotdata_describe_tables_tool` + take `database_id=` in place of `database=`. It accepts an id, or a `ManagedDatabase` to + skip the lookup. The id is resolved once when the tools are built, so a bad id fails + there rather than on the agent's first query, and queries no longer pay a repeat lookup. + - The `hotdata_load_managed_table` tool's `database` argument is now `database_id`, and its + description names the two tools that hand out ids. + - `load_managed_table` takes `database_id=`; `execute_sql_json`, `bm25_search_json` and + `describe_tables_json` take a resolved `ManagedDatabase` as `database=` and raise + `TypeError` on a string, which would otherwise reach the framework's by-name fallback. + - Passing a name anywhere raises `KeyError`, naming `hotdata_list_managed_databases` as + where ids come from. + + Mirrors `hotdata-dlt-destination`'s move to id-only addressing. To scope tools to the same + database as before, pass its id: `client.list_managed_databases()` reports one per database. - Tool descriptions now state the engine's actual contract instead of a one-line summary. `hotdata_execute_sql` names the dialect and the supported constructs, points at the search tool for text relevance when one is registered, and warns that an aggregate query must diff --git a/README.md b/README.md index e1b1ff2..625b558 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ from langchain.agents import create_agent import hotdata_langchain as hl client = hl.from_env() -tools = hl.make_hotdata_tools(client, database="sales") +tools = hl.make_hotdata_tools(client, database_id="dbid...") agent = create_agent(model=your_model, tools=tools) result = agent.invoke( @@ -32,8 +32,9 @@ result = agent.invoke( print(result["messages"][-1].content) ``` -Queries run against a database scope, so pass `database=` (a managed database name or id). -`hl.from_env().list_managed_databases()` shows what is available in the workspace. +Queries run against a database scope, so pass `database_id=` (a managed database id). +`hl.from_env().list_managed_databases()` shows what is available in the workspace, with the +id of each. ## Tools @@ -42,9 +43,9 @@ Queries run against a database scope, so pass `database=` (a managed database na | Tool | What it does | |------|-------------| | `hotdata_execute_sql` | Run a SQL query and return rows as JSON | -| `hotdata_list_managed_databases` | List available managed databases | -| `hotdata_create_managed_database` | Create a new managed database | -| `hotdata_load_managed_table` | Load a parquet file into a managed table | +| `hotdata_list_managed_databases` | List available managed databases, with the id of each | +| `hotdata_create_managed_database` | Create a new managed database and return its id | +| `hotdata_load_managed_table` | Load a parquet file into a managed table, addressed by database id | | `hotdata_describe_tables` | List tables, or one table's columns and types | | `hotdata_search_text` | Full-text search an indexed column, ranked by relevance (opt-in — see below) | @@ -58,8 +59,8 @@ table with its column count; called with a table name it returns that table's co types. Without it an agent has to guess column names, and a guess that misses fails the query. ```python -tools = hl.make_hotdata_tools(client, database="sales") # included -tools = hl.make_hotdata_tools(client, database="sales", describe_tables=False) # omitted +tools = hl.make_hotdata_tools(client, database_id="dbid...") # included +tools = hl.make_hotdata_tools(client, database_id="dbid...", describe_tables=False) # omitted ``` It reads `information_schema` in whichever database the tools are scoped to, so it needs no @@ -71,19 +72,21 @@ extra permissions. With it turned off, the SQL tool's description tells the agen You can also invoke tools outside of an agent loop: ```python -tools = {t.name: t for t in hl.make_hotdata_tools(client, database="sales")} +import json + +tools = {t.name: t for t in hl.make_hotdata_tools(client, database_id="dbid...")} result = tools["hotdata_execute_sql"].invoke({"sql": "SELECT * FROM orders LIMIT 10"}) print(result) # JSON rows -tools["hotdata_create_managed_database"].invoke({ - "name": "sales", +created = tools["hotdata_create_managed_database"].invoke({ + "name": "sales", # a display label, not an identifier "schema_name": "public", "tables": "orders,customers", }) tools["hotdata_load_managed_table"].invoke({ - "database": "sales", + "database_id": json.loads(created)["id"], "table": "orders", "file": "/path/to/orders.parquet", }) @@ -96,7 +99,7 @@ Point the agent at a text column carrying a BM25 index and it gets a search tool ```python tools = hl.make_hotdata_tools( client, - database="sf_airbnb", + database_id="dbid...", search_table="default.public.listings", # catalog.schema.table search_column="description", # must have a BM25 index search_columns=["id", "name", "price", "description"], # what each hit returns @@ -114,7 +117,7 @@ the tool surface lets an agent discover which columns are indexed, and the engin outright rather than falling back to a scan when a column has no BM25 index. Inside a managed database the built-in catalog is always `default`, so a managed table reads -as `default..` when `database=` scopes the query to it. +as `default..
` when `database_id=` scopes the query to it. For more than one searchable corpus, build the tools yourself and give each a distinct name and description — the agent then routes on the descriptions: @@ -126,14 +129,14 @@ tools = [ # and the agent goes back to trying to match text in SQL. *hl.make_hotdata_tools( client, - database="sf_airbnb", + database_id="dbid...", search_table="default.public.listings", search_column="description", search_tool_name="search_listings", ), hl.make_hotdata_search_tool( client, table="default.public.reviews", column="comments", - name="search_reviews", database="sf_airbnb", + name="search_reviews", database_id="dbid...", ), ] ``` @@ -144,12 +147,27 @@ index, then an agent that picks between search and SQL. ## Scoping queries to a managed database -`database=` resolves all SQL the agent runs against a specific managed database, by name or -id. The API requires a database scope, so queries fail with `a database is required` without -it: +`database_id=` scopes all SQL the agent runs to one managed database. The API requires a +database scope, so queries fail with `a database is required` without it: + +```python +tools = hl.make_hotdata_tools(client, database_id="dbid...") +``` + +**Databases are addressed by id, never by name.** A database name is a display label and is +not unique, so a name lookup can silently resolve to the wrong database — and the agent's +`hotdata_load_managed_table` overwrites the table it loads into. Passing a name raises +`KeyError`. Ids come from `client.list_managed_databases()`, the +`hotdata_list_managed_databases` tool, or the response of a create. + +The id is resolved once when the tools are built, so a bad id fails there rather than on the +agent's first query, and no query pays a repeat lookup. If you already hold a +`ManagedDatabase` — from `list_managed_databases()` or `create_managed_database()` — pass it +instead of its id to skip the lookup entirely: ```python -tools = hl.make_hotdata_tools(client, database="sales") +db = client.create_managed_database(description="sales", schema="public", tables=["orders"]) +tools = hl.make_hotdata_tools(client, database_id=db) ``` ## Controlling result size diff --git a/demo/README.md b/demo/README.md index 4770417..dc2a2e7 100644 --- a/demo/README.md +++ b/demo/README.md @@ -24,6 +24,10 @@ The script is safe to re-run: it reuses the managed database, skips the load whe table already has rows, and reuses an existing index. ```bash +# bind an existing managed database by id instead of letting the demo find its own +uv run --group demo --env-file .env python demo/bm25_search_demo.py \ + --database-id dbid... + # different search text, more hits uv run --group demo --env-file .env python demo/bm25_search_demo.py \ --query "quiet garden studio near the park" --k 10 @@ -45,6 +49,7 @@ LANGSMITH_TRACING=true LANGSMITH_PROJECT=hotdata-langchain-bm25 \ |---|---|---| | `HOTDATA_API_KEY` | steps 1–5 | everything except the agent run | | `HOTDATA_WORKSPACE` | optional | pins a workspace; first available otherwise | +| `DEMO_DATABASE_ID` | optional | pins the managed database by id (same as `--database-id`) | | your model provider's key | step 6 | skipped without it, or without `--model` | | `LANGSMITH_API_KEY` + `LANGSMITH_TRACING=true` | optional | traces the run to LangSmith | | `LANGSMITH_PROJECT` | optional | project the traces land in (default: `default`) | @@ -60,8 +65,12 @@ and SQL are both inspectable after the fact. ## What each step does -1. **Managed database** — creates `langchain_bm25_demo` with `public.listings` declared - up front, so the load materialises into it directly. +1. **Managed database** — binds the database given by `--database-id`, or else creates one + labelled `langchain_bm25_demo` with `public.listings` declared up front, so the load + materialises into it directly. Everything downstream addresses the resolved record, so + the label never selects the target; the create path prints the new id to pin. Finding a + previous run's database by its label is the one by-label lookup here, and it exists only + so the demo is re-runnable without pinning — pass `--database-id` to skip it. 2. **Listings data** — downloads the fixture parquet (cached in the system temp dir) and loads it into the managed table. 3. **BM25 index** — creates a `bm25` index on `description` through `IndexesApi` and @@ -83,16 +92,17 @@ and SQL are both inspectable after the fact. Set the model with `--model` or `DEMO_MODEL`; any tool-calling model works, and the step is skipped when its provider key is absent. - **What this step demonstrates is the routing.** Across five runs of the current - configuration the agent called the search tool and the schema tool every time, and the - final table matched the whole-neighbourhood figures every time. Two of those runs also - recovered from a bad query after the error was handed back to them. - - Do not read five runs as a guarantee. Answering a compound question correctly is a - property of the model rather than of these tools, and an earlier configuration — which - shared one row budget between the search and SQL tools, silently truncating the - aggregate — got it right in only four of five. The printed tool calls are the reliable - part; treat the prose answer as illustrative. + **What this step demonstrates is the routing, and only the routing.** Across five runs + the agent called the search tool and the schema tool every time, in either order. That + part is what the tools are responsible for, and it holds. + + The prose answer is a different matter: the final table matched the true + whole-neighbourhood figures in three of five runs. The two failures were the model + aggregating over the handful of listings search returned instead of over every listing + in those neighbourhoods, and a join that inflated the counts. Five runs on the merged + `main`, measured the same way against the same ground-truth query, also gave three of + five — so this is the model's ceiling on a compound question, not something these tools + introduced or can fix. Read the printed tool calls, not the table. ## What makes the agent run work diff --git a/demo/bm25_search_demo.py b/demo/bm25_search_demo.py index f12fe8c..59648fe 100644 --- a/demo/bm25_search_demo.py +++ b/demo/bm25_search_demo.py @@ -29,7 +29,9 @@ PARQUET_URL = "https://www.hotdata.dev/data/sf-airbnb-listings.parquet" USER_AGENT = "hotdata-langchain-demo" -DATABASE = "langchain_bm25_demo" +#: Display label for the database this demo creates. Not an identifier — see +#: `ensure_database`, which addresses the database by id once it has one. +DATABASE_LABEL = "langchain_bm25_demo" SCHEMA = "public" TABLE = "listings" SEARCH_COLUMN = "description" @@ -88,33 +90,45 @@ def download_parquet() -> Path: return path -def find_database(client: hl.HotdataClient, name: str) -> Any | None: +def find_database_by_label(client: hl.HotdataClient, label: str) -> Any | None: + """Scan the workspace for a database carrying this display label. + + Bootstrap convenience for a re-runnable demo, and the only by-label lookup here. + Labels are not unique, so this is not how an application should find its database — + pass `--database-id` to bind one by id instead. + """ for db in client.list_managed_databases(): - if db.description == name or db.id == name: + if db.description == label: return db return None -def ensure_database(client: hl.HotdataClient) -> Any: - existing = find_database(client, DATABASE) +def ensure_database(client: hl.HotdataClient, database_id: str | None) -> Any: + """Return the demo's managed database record, bound by id or created.""" + if database_id: + db = hl.resolve_database_by_id(client, database_id) + print(f"Bound managed database {db.id} by id (label={db.description!r})") + return db + + existing = find_database_by_label(client, DATABASE_LABEL) if existing is not None: - print(f"Reusing managed database {DATABASE!r} (id={existing.id})") + print(f"Reusing managed database {existing.id} (label={DATABASE_LABEL!r})") return existing - db = client.create_managed_database(description=DATABASE, schema=SCHEMA, tables=[TABLE]) - print(f"Created managed database {DATABASE!r} (id={db.id}) with {SCHEMA}.{TABLE} declared") + + db = client.create_managed_database(description=DATABASE_LABEL, schema=SCHEMA, tables=[TABLE]) + print(f"Created managed database {db.id} (label={DATABASE_LABEL!r}) with {SCHEMA}.{TABLE}") + print(f" Pin it for later runs with --database-id {db.id} (or DEMO_DATABASE_ID)") return db -def load_listings(client: hl.HotdataClient, parquet: Path) -> None: - loaded = client.load_managed_table(DATABASE, TABLE, schema=SCHEMA, file=str(parquet)) +def load_listings(client: hl.HotdataClient, db: Any, parquet: Path) -> None: + loaded = client.load_managed_table(db, TABLE, schema=SCHEMA, file=str(parquet)) print(f"Loaded {loaded.row_count} rows into {loaded.full_name}") -def table_columns(client: hl.HotdataClient) -> list[str]: +def table_columns(client: hl.HotdataClient, db: Any) -> list[str]: """Return the table's column names, through the same tool the agent gets.""" - described = json.loads( - hl.describe_tables_json(client, table=f"{SCHEMA}.{TABLE}", database=DATABASE) - ) + described = json.loads(hl.describe_tables_json(client, table=f"{SCHEMA}.{TABLE}", database=db)) return [column["name"] for column in described["columns"]] @@ -250,6 +264,12 @@ def main() -> None: help="tool-calling model for the agent step, e.g. ':' " "(or set DEMO_MODEL); the agent step is skipped without one", ) + parser.add_argument( + "--database-id", + default=os.environ.get("DEMO_DATABASE_ID"), + help="bind an existing managed database by id (or set DEMO_DATABASE_ID); " + "without one the demo reuses or creates its own and prints the id to pin", + ) parser.add_argument("--skip-agent", action="store_true", help="stop after direct tool use") parser.add_argument("--reload", action="store_true", help="reload the parquet even if loaded") parser.add_argument( @@ -261,18 +281,22 @@ def main() -> None: print(f"Connected to {client.host} (workspace={client.workspace_id})") if args.cleanup: - existing = find_database(client, DATABASE) - if existing is None: - print(f"No managed database named {DATABASE!r} to delete") + target = ( + hl.resolve_database_by_id(client, args.database_id) + if args.database_id + else find_database_by_label(client, DATABASE_LABEL) + ) + if target is None: + print(f"No managed database labelled {DATABASE_LABEL!r} to delete") else: - client.delete_managed_database(existing.id) - print(f"Deleted managed database {DATABASE!r} (id={existing.id})") + client.delete_managed_database(target) + print(f"Deleted managed database {target.id} (label={target.description!r})") client.close() return try: step("1. Managed database") - db = ensure_database(client) + db = ensure_database(client, args.database_id) step("2. Listings data") already_loaded = False @@ -280,20 +304,20 @@ def main() -> None: # Probing with a row read rather than COUNT(*): the engine rejects a # projection that is aggregates only. try: - probe = client.execute_sql(f"SELECT id FROM {TABLE_REF} LIMIT 1", database=DATABASE) + probe = client.execute_sql(f"SELECT id FROM {TABLE_REF} LIMIT 1", database=db) already_loaded = bool(probe.rows) if already_loaded: print(f"{TABLE_REF} already holds data; skipping load (--reload to force)") except Exception as e: print(f"Table not queryable yet ({type(e).__name__}); loading fixture") if not already_loaded: - load_listings(client, download_parquet()) + load_listings(client, db, download_parquet()) step("3. BM25 index") ensure_bm25_index(client, db.default_connection_id) step("4. Tools") - available = table_columns(client) + available = table_columns(client, db) columns = [c for c in PREFERRED_COLUMNS if c in available] if SEARCH_COLUMN not in columns: columns.append(SEARCH_COLUMN) @@ -301,7 +325,8 @@ def main() -> None: tools = hl.make_hotdata_tools( client, - database=DATABASE, + # The resolved record, so the tool set does not re-look-up what step 1 has. + database_id=db, # Not args.k: max_rows also caps the SQL tool, and the agent's aggregate in # step 6 groups over whole neighbourhoods, which a search-sized budget would # silently truncate. Search hits and SQL rows are different budgets. diff --git a/docs/engine-contract.md b/docs/engine-contract.md index aeb854e..673c557 100644 --- a/docs/engine-contract.md +++ b/docs/engine-contract.md @@ -97,7 +97,9 @@ discover which columns are searchable, and why the search tool pins its corpus. translates "an attachment alias or `default`", so attachment is presumably the supported route — **not verified here**. - **Database names are not unique.** `name` is a display label; `resolve_managed_database` tries - the id first and then scans `list_databases()` matching on name. Ids are the only safe handle. + the id first and then scans `list_databases()` matching on name. Ids are the only safe handle, + so this package never calls that resolver: `resolve_database_by_id` goes straight to + `GET /databases/{id}`, and a resolved `ManagedDatabase` is what scopes every query. - **`from_env()` picks a workspace silently** when `HOTDATA_WORKSPACE` is unset — first active, else first overall, no warning. `HotdataClient(api_key, workspace_id)` takes it explicitly. diff --git a/examples/langchain_basic.py b/examples/langchain_basic.py index 27003d9..d24e179 100644 --- a/examples/langchain_basic.py +++ b/examples/langchain_basic.py @@ -9,14 +9,16 @@ def main() -> None: tools = {tool.name: tool for tool in hl.make_hotdata_tools(client)} print(tools["hotdata_list_managed_databases"].invoke({})) - # Queries need a database scope, so pick one from the workspace. + # Queries need a database scope, addressed by id — a database name is a display + # label and is not unique. Pass the listed record straight through to reuse the + # lookup that produced it. databases = client.list_managed_databases() if not databases: print("No managed databases in this workspace; create one to run a query.") client.close() return - scoped = hl.make_hotdata_tools(client, database=databases[0].id) + scoped = hl.make_hotdata_tools(client, database_id=databases[0]) by_name = {tool.name: tool for tool in scoped} print(by_name["hotdata_execute_sql"].invoke({"sql": "SELECT 1 AS ok"})) diff --git a/examples/langchain_managed_db.py b/examples/langchain_managed_db.py index 2ac8228..44d5ee1 100644 --- a/examples/langchain_managed_db.py +++ b/examples/langchain_managed_db.py @@ -1,5 +1,7 @@ """Managed database tools for LangChain agents.""" +import json + import hotdata_langchain as hl @@ -9,21 +11,23 @@ def main() -> None: by_name = {tool.name: tool for tool in tools} create = by_name["hotdata_create_managed_database"] - print( - create.invoke( - { - "name": "demo_sales", - "schema_name": "public", - "tables": "orders\ncustomers", - } - ) + created = create.invoke( + { + "name": "demo_sales", + "schema_name": "public", + "tables": "orders\ncustomers", + } ) + print(created) + + # 'name' was a display label; the id is what addresses the database from here on. + database_id = json.loads(created)["id"] load = by_name["hotdata_load_managed_table"] print( load.invoke( { - "database": "demo_sales", + "database_id": database_id, "table": "orders", "file": "/path/to/orders.parquet", "schema_name": "public", diff --git a/hotdata_langchain/__init__.py b/hotdata_langchain/__init__.py index 818cbc7..b69869b 100644 --- a/hotdata_langchain/__init__.py +++ b/hotdata_langchain/__init__.py @@ -7,7 +7,7 @@ except PackageNotFoundError: __version__ = "0.0.0+unknown" -from hotdata_framework import HotdataClient, QueryResult, from_env +from hotdata_framework import HotdataClient, ManagedDatabase, QueryResult, from_env from hotdata_langchain.databases import ( create_managed_database, @@ -15,6 +15,7 @@ load_managed_table, load_result_summary, managed_database_summary, + resolve_database_by_id, ) from hotdata_langchain.schema import ( DEFAULT_DESCRIBE_TOOL_NAME, @@ -41,6 +42,7 @@ "DEFAULT_SEARCH_TOOL_NAME", "SCORE_COLUMN", "HotdataClient", + "ManagedDatabase", "QueryResult", "__version__", "bm25_search_json", @@ -56,5 +58,6 @@ "make_hotdata_search_tool", "make_hotdata_tools", "managed_database_summary", + "resolve_database_by_id", "result_rows_for_llm", ] diff --git a/hotdata_langchain/databases.py b/hotdata_langchain/databases.py index c3cc79d..b262499 100644 --- a/hotdata_langchain/databases.py +++ b/hotdata_langchain/databases.py @@ -5,12 +5,62 @@ import json from typing import Any +from hotdata.api.databases_api import DatabasesApi +from hotdata.exceptions import ApiException from hotdata_framework import ( DEFAULT_SCHEMA, HotdataClient, LoadManagedTableResult, ManagedDatabase, ) +from hotdata_framework.databases import api_error_message, managed_database_from_detail + + +def resolve_database_by_id( + client: HotdataClient, + database_id: str | ManagedDatabase, +) -> ManagedDatabase: + """Fetch a managed database record by id (``GET /databases/{id}``). + + Addresses the database by id only. A Hotdata database name is a display label and is + not unique, so there is deliberately no by-name fallback: a name that collides with + another database's label would otherwise resolve to the wrong database, and every + query, load and drop would follow it there. Ids come from + :func:`list_managed_databases_json` or :func:`create_managed_database`. + + An already-resolved ``ManagedDatabase`` is returned as-is, so a caller holding one + pays no lookup. + + Raises ``KeyError`` when the workspace has no database with that id. + """ + if isinstance(database_id, ManagedDatabase): + return database_id + try: + detail = DatabasesApi(client.api).get_database(database_id) + except ApiException as e: + if e.status == 404: + raise KeyError( + f"no managed database with id {database_id!r} in this workspace. " + "Ids are listed by hotdata_list_managed_databases; a database name is " + "not accepted here, because names are not unique." + ) from e + raise RuntimeError(api_error_message(e)) from e + return managed_database_from_detail(detail) + + +def query_scope(database: ManagedDatabase | None) -> ManagedDatabase | None: + """Return ``database`` unchanged, rejecting a scope that was never resolved. + + A string reaching ``HotdataClient`` would go through its name-or-id resolver, whose + by-name fallback matches a non-unique display label. Resolve with + :func:`resolve_database_by_id` first. + """ + if database is None or isinstance(database, ManagedDatabase): + return database + raise TypeError( + f"database must be a resolved ManagedDatabase, got {type(database).__name__}; " + "resolve it with resolve_database_by_id(client, database_id) first" + ) def list_managed_databases_json(client: HotdataClient) -> str: @@ -32,17 +82,29 @@ def create_managed_database( schema: str = DEFAULT_SCHEMA, tables: list[str] | None = None, ) -> ManagedDatabase: + """Create a managed database, labelled ``name``. + + ``name`` is a display label only; address the result by its ``id`` from here on. + """ return client.create_managed_database(description=name, schema=schema, tables=tables) def load_managed_table( client: HotdataClient, *, - database: str, + database_id: str | ManagedDatabase, table: str, file: str, schema: str = DEFAULT_SCHEMA, ) -> LoadManagedTableResult: + """Load a local parquet file into a declared table of the database with that id. + + ``database_id`` is resolved by id (see :func:`resolve_database_by_id`) and the + resolved record is what addresses the load, so a display label never selects the + target. This load replaces the table's contents, which is why addressing it + unambiguously matters. + """ + database = resolve_database_by_id(client, database_id) return client.load_managed_table(database, table, schema=schema, file=file) diff --git a/hotdata_langchain/schema.py b/hotdata_langchain/schema.py index 3302bd6..3ea79e9 100644 --- a/hotdata_langchain/schema.py +++ b/hotdata_langchain/schema.py @@ -4,10 +4,11 @@ import json -from hotdata_framework import HotdataClient +from hotdata_framework import HotdataClient, ManagedDatabase from langchain_core.tools import StructuredTool from hotdata_langchain._sql import validate_identifier +from hotdata_langchain.databases import query_scope, resolve_database_by_id DEFAULT_DESCRIBE_TOOL_NAME = "hotdata_describe_tables" @@ -56,7 +57,7 @@ def describe_tables_json( client: HotdataClient, *, table: str | None = None, - database: str | None = None, + database: ManagedDatabase | None = None, max_columns: int = DEFAULT_MAX_COLUMNS, ) -> str: """Describe the scoped database's tables, or one table's columns, as JSON. @@ -66,12 +67,16 @@ def describe_tables_json( declaration order, capped at ``max_columns`` so a wide table cannot flood the model's context; the payload says so when the cap truncated the list. + ``database`` is a resolved ``ManagedDatabase``, not an id or a name — resolve one + with :func:`hotdata_langchain.databases.resolve_database_by_id`. + Raises ``ValueError`` for a non-positive ``max_columns``. """ if max_columns < 1: raise ValueError(f"max_columns must be >= 1, got {max_columns}") + scope = query_scope(database) if table is None: - result = client.execute_sql(table_overview_sql(), database=database) + result = client.execute_sql(table_overview_sql(), database=scope) tables = [ { "table": f"{row['table_schema']}.{row['table_name']}", @@ -84,7 +89,7 @@ def describe_tables_json( # One row past the cap, so a table with exactly `max_columns` columns is reported as # complete rather than flagged as truncated — telling the model part of the schema is # missing is the one thing likely to send it back to guessing. - result = client.execute_sql(table_columns_sql(table, limit=max_columns + 1), database=database) + result = client.execute_sql(table_columns_sql(table, limit=max_columns + 1), database=scope) records = result.to_records() if not records: return json.dumps( @@ -117,17 +122,22 @@ def default_describe_description() -> str: def make_hotdata_describe_tables_tool( client: HotdataClient, *, - database: str | None = None, + database_id: str | ManagedDatabase | None = None, name: str = DEFAULT_DESCRIBE_TOOL_NAME, description: str | None = None, max_columns: int = DEFAULT_MAX_COLUMNS, ) -> StructuredTool: """Return a LangChain tool that reports the scoped database's tables and columns. + ``database_id`` scopes the introspection to one managed database, by id and never by + name; it is resolved once here. Pass an already-resolved ``ManagedDatabase`` to skip + the lookup. + Fails fast on a non-positive ``max_columns`` rather than at first invocation. """ if max_columns < 1: raise ValueError(f"max_columns must be >= 1, got {max_columns}") + database = resolve_database_by_id(client, database_id) if database_id is not None else None def hotdata_describe_tables(table: str | None = None) -> str: """List the tables in the database, or one table's columns and types.""" diff --git a/hotdata_langchain/search.py b/hotdata_langchain/search.py index af798ed..35eecaa 100644 --- a/hotdata_langchain/search.py +++ b/hotdata_langchain/search.py @@ -6,10 +6,11 @@ import re from collections.abc import Sequence -from hotdata_framework import HotdataClient +from hotdata_framework import HotdataClient, ManagedDatabase from langchain_core.tools import StructuredTool from hotdata_langchain._sql import quote_literal, validate_identifier +from hotdata_langchain.databases import query_scope, resolve_database_by_id #: Column the engine appends to every ``bm25_search`` result, holding the BM25 relevance score. SCORE_COLUMN = "score" @@ -95,16 +96,19 @@ def bm25_search_json( k: int = DEFAULT_SEARCH_LIMIT, columns: Sequence[str] | None = None, max_rows: int = 100, - database: str | None = None, + database: ManagedDatabase | None = None, ) -> str: """Run a BM25 search and return ``{"metadata": ..., "rows": [...]}`` as JSON. Mirrors the envelope :func:`hotdata_langchain.tools.execute_sql_json` returns, so an agent sees one result shape across every Hotdata tool. Rows arrive ranked by ``score`` descending. + + ``database`` is a resolved ``ManagedDatabase``, not an id or a name — resolve one + with :func:`hotdata_langchain.databases.resolve_database_by_id`. """ sql = bm25_search_sql(table=table, column=column, query=query, k=k, columns=columns) - result = client.execute_sql(sql, database=database) + result = client.execute_sql(sql, database=query_scope(database)) payload = { "metadata": result.metadata_dict(), "rows": result.to_records(max_rows=max_rows), @@ -141,7 +145,7 @@ def make_hotdata_search_tool( name: str = DEFAULT_SEARCH_TOOL_NAME, description: str | None = None, max_rows: int = 100, - database: str | None = None, + database_id: str | ManagedDatabase | None = None, ) -> StructuredTool: """Return a LangChain tool that full-text searches one indexed column. @@ -150,6 +154,10 @@ def make_hotdata_search_tool( errors outright when one is missing. The agent supplies only ``query`` and an optional ``k``. + ``database_id`` scopes the search to one managed database, by id and never by name; + it is resolved once here. Pass an already-resolved ``ManagedDatabase`` to skip the + lookup. + Register the factory more than once, with distinct ``name`` and ``description`` values, to expose several searchable corpora; the agent then routes on the descriptions. @@ -167,6 +175,7 @@ def make_hotdata_search_tool( if columns is not None: _projection(column, columns) default_k = k + database = resolve_database_by_id(client, database_id) if database_id is not None else None def hotdata_search_text(query: str, k: int | None = None) -> str: """Search indexed text by relevance and return ranked rows as JSON.""" diff --git a/hotdata_langchain/tools.py b/hotdata_langchain/tools.py index b6c4027..f878ef4 100644 --- a/hotdata_langchain/tools.py +++ b/hotdata_langchain/tools.py @@ -6,7 +6,7 @@ from collections.abc import Sequence from typing import Any -from hotdata_framework import DEFAULT_SCHEMA, HotdataClient, QueryResult +from hotdata_framework import DEFAULT_SCHEMA, HotdataClient, ManagedDatabase, QueryResult from langchain_core.tools import StructuredTool from hotdata_langchain.databases import ( @@ -15,6 +15,8 @@ load_managed_table, load_result_summary, managed_database_summary, + query_scope, + resolve_database_by_id, ) from hotdata_langchain.schema import ( DEFAULT_DESCRIBE_TOOL_NAME, @@ -83,9 +85,14 @@ def execute_sql_json( sql: str, *, max_rows: int = 100, - database: str | None = None, + database: ManagedDatabase | None = None, ) -> str: - result = client.execute_sql(sql, database=database) + """Run SQL scoped to an already-resolved managed database and return JSON. + + ``database`` is a resolved ``ManagedDatabase``, not an id or a name — resolve one + with :func:`hotdata_langchain.databases.resolve_database_by_id`. + """ + result = client.execute_sql(sql, database=query_scope(database)) payload = { "metadata": result.metadata_dict(), "rows": result.to_records(max_rows=max_rows), @@ -97,7 +104,7 @@ def make_hotdata_tools( client: HotdataClient, *, max_rows: int = 100, - database: str | None = None, + database_id: str | ManagedDatabase | None = None, search_table: str | None = None, search_column: str | None = None, search_columns: Sequence[str] | None = None, @@ -107,6 +114,13 @@ def make_hotdata_tools( ) -> list[StructuredTool]: """Return LangChain tools for SQL and managed database workflows. + ``database_id`` scopes every query these tools run to one managed database. It is a + database id, never a name: names are display labels and are not unique. The id is + resolved once here and the resolved record is what each query carries, so a + non-existent id fails at build time rather than on the agent's first query. Pass an + already-resolved ``ManagedDatabase`` to skip the lookup. Ids come from + ``client.list_managed_databases()`` or the ``hotdata_list_managed_databases`` tool. + ``describe_tables`` (on by default) adds a schema-introspection tool, so the agent can look up tables and columns instead of guessing them. It reads ``information_schema`` in whichever database the tools are scoped to. @@ -123,6 +137,8 @@ def make_hotdata_tools( if (search_table is None) != (search_column is None): raise ValueError("search_table and search_column must be provided together") + database = resolve_database_by_id(client, database_id) if database_id is not None else None + def hotdata_execute_sql(sql: str) -> str: """Run SQL against the Hotdata workspace and return JSON rows.""" return execute_sql_json(client, sql, max_rows=max_rows, database=database) @@ -147,7 +163,7 @@ def hotdata_create_managed_database( return json.dumps(managed_database_summary(db), indent=2) def hotdata_load_managed_table( - database: str, + database_id: str, table: str, file: str, schema_name: str = DEFAULT_SCHEMA, @@ -155,7 +171,7 @@ def hotdata_load_managed_table( """Load a local parquet file into a declared managed table.""" loaded = load_managed_table( client, - database=database, + database_id=database_id, table=table, file=file, schema=schema_name or DEFAULT_SCHEMA, @@ -178,7 +194,9 @@ def hotdata_load_managed_table( description=( "List the managed databases in this workspace. Returns each database's " "'id' and its human-readable 'description'. Names are display labels and " - "are not unique — pass the 'id' to other tools, never the description." + "are not unique — pass the 'id' to other tools, never the description. " + "An id cannot be guessed or built from a name; it only comes from here or " + "from creating a database." ), ), StructuredTool.from_function( @@ -186,9 +204,10 @@ def hotdata_load_managed_table( name="hotdata_create_managed_database", description=( "Create a managed database to hold tables you load. 'name' is a display " - "label; the response carries the 'id' to use with the other tools. Declare " - "the tables you intend to load up front as a comma- or newline-separated " - "list, so data loads straight into them." + "label only and is not an identifier; the response carries the 'id', which " + "is what every other tool needs — keep it. Declare the tables you intend " + "to load up front as a comma- or newline-separated list, so data loads " + "straight into them." ), ), StructuredTool.from_function( @@ -197,14 +216,18 @@ def hotdata_load_managed_table( description=( "Load a parquet file from the local filesystem into a table that was " "declared on a managed database, replacing whatever the table held. " - "'database' should be a database id. Only local parquet paths are " - "accepted — not URLs, and not other file formats." + "'database_id' must be a database id returned by " + "hotdata_list_managed_databases or hotdata_create_managed_database — call " + "one of those first if you do not have an id. A database name is rejected: " + "names are not unique, and this load overwrites the table, so the wrong " + "target would destroy data. Only local parquet paths are accepted — not " + "URLs, and not other file formats." ), ), ] if describe_tables: - tools.append(make_hotdata_describe_tables_tool(client, database=database)) + tools.append(make_hotdata_describe_tables_tool(client, database_id=database)) if has_search: assert search_table is not None and search_column is not None @@ -217,7 +240,7 @@ def hotdata_load_managed_table( k=search_k, name=search_tool_name, max_rows=max_rows, - database=database, + database_id=database, ) ) diff --git a/tests/conftest.py b/tests/conftest.py index 1455fbc..999351c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,42 @@ from __future__ import annotations -from unittest.mock import MagicMock +from collections.abc import Iterator +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest -from hotdata_framework import QueryResult +from hotdata_framework import ManagedDatabase, QueryResult + + +@pytest.fixture +def managed_db() -> ManagedDatabase: + """A resolved managed database, as ``resolve_database_by_id`` returns one. + + Query scopes are resolved records rather than id strings, so tests pass this where + an id would previously have been threaded through. + """ + return ManagedDatabase( + id="dbidsf000000000000000000000001", + description="sf_airbnb", + default_connection_id="connsf00000000000000000000001", + ) + + +@pytest.fixture +def databases_api(managed_db: ManagedDatabase) -> Iterator[MagicMock]: + """Patch the raw databases API so an id lookup resolves to ``managed_db``. + + ``GET /databases/{id}`` is the only lookup the package is allowed to make, so tests + stub it here and assert against the calls it received. + """ + detail = SimpleNamespace( + id=managed_db.id, + name=managed_db.description, + default_connection_id=managed_db.default_connection_id, + ) + with patch("hotdata_langchain.databases.DatabasesApi") as api: + api.return_value.get_database.return_value = detail + yield api @pytest.fixture diff --git a/tests/test_database_ids.py b/tests/test_database_ids.py new file mode 100644 index 0000000..80a2a70 --- /dev/null +++ b/tests/test_database_ids.py @@ -0,0 +1,217 @@ +"""Managed databases are addressed by id, never by name. + +A Hotdata database name is a display label and is not unique, so a by-name lookup can +resolve to the wrong database — and then every query, load and drop follows it there. +The framework's ``resolve_managed_database`` still offers that by-name fallback; these +tests pin that this package never reaches it, and that ``GET /databases/{id}`` is the +only lookup it makes. +""" + +from __future__ import annotations + +import json +from unittest.mock import MagicMock + +import pytest +from hotdata.exceptions import ApiException +from hotdata_framework import LoadManagedTableResult, ManagedDatabase + +from hotdata_langchain.databases import query_scope, resolve_database_by_id +from hotdata_langchain.schema import make_hotdata_describe_tables_tool +from hotdata_langchain.search import make_hotdata_search_tool +from hotdata_langchain.tools import execute_sql_json, make_hotdata_tools + +TABLE = "default.public.listings" +COLUMN = "description" + + +# --- resolution -------------------------------------------------------------------- + + +def test_resolve_fetches_the_record_by_id( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + resolved = resolve_database_by_id(mock_client, managed_db.id) + assert resolved == managed_db + databases_api.return_value.get_database.assert_called_once_with(managed_db.id) + + +def test_resolve_never_lists_or_matches_on_a_name( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + """Listing is how a name match would be found, so it must not happen at all.""" + resolve_database_by_id(mock_client, managed_db.id) + databases_api.return_value.list_databases.assert_not_called() + mock_client.list_managed_databases.assert_not_called() + mock_client.resolve_managed_database.assert_not_called() + + +def test_resolve_passes_an_already_resolved_record_through_without_a_lookup( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + assert resolve_database_by_id(mock_client, managed_db) is managed_db + databases_api.return_value.get_database.assert_not_called() + + +def test_resolve_raises_key_error_for_an_unknown_id( + mock_client: MagicMock, databases_api: MagicMock +) -> None: + """A name lands here: it is not an id, so it 404s rather than matching a label.""" + databases_api.return_value.get_database.side_effect = ApiException( + status=404, reason="Not Found" + ) + with pytest.raises(KeyError) as excinfo: + resolve_database_by_id(mock_client, "sf_airbnb") + message = str(excinfo.value) + assert "hotdata_list_managed_databases" in message + assert "not unique" in message + + +def test_resolve_surfaces_other_api_errors( + mock_client: MagicMock, databases_api: MagicMock +) -> None: + databases_api.return_value.get_database.side_effect = ApiException( + status=403, reason="Forbidden", body="workspace does not permit reads" + ) + with pytest.raises(RuntimeError, match="workspace does not permit reads"): + resolve_database_by_id(mock_client, "dbid000000000000000000000000x") + + +# --- query scopes ------------------------------------------------------------------ + + +def test_query_scope_accepts_a_resolved_record_and_none(managed_db: ManagedDatabase) -> None: + assert query_scope(managed_db) is managed_db + assert query_scope(None) is None + + +def test_query_scope_rejects_an_id_or_name_string() -> None: + """Strings reach the framework's name-or-id resolver, which is what we are avoiding.""" + with pytest.raises(TypeError, match="resolve_database_by_id"): + query_scope("dbid000000000000000000000000x") # type: ignore[arg-type] + + +def test_execute_sql_json_refuses_an_unresolved_scope(mock_client: MagicMock) -> None: + with pytest.raises(TypeError, match="resolve_database_by_id"): + execute_sql_json(mock_client, "SELECT 1", database="sf_airbnb") # type: ignore[arg-type] + + +# --- factories resolve once -------------------------------------------------------- + + +def test_tool_set_resolves_the_database_exactly_once( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + """SQL, schema and search tools share one resolved record — not one lookup each.""" + make_hotdata_tools( + mock_client, + database_id=managed_db.id, + search_table=TABLE, + search_column=COLUMN, + ) + databases_api.return_value.get_database.assert_called_once_with(managed_db.id) + + +def test_tool_set_scopes_every_query_to_the_resolved_record( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + tools = {t.name: t for t in make_hotdata_tools(mock_client, database_id=managed_db.id)} + tools["hotdata_execute_sql"].invoke({"sql": "SELECT 1"}) + assert mock_client.execute_sql.call_args.kwargs == {"database": managed_db} + + +def test_tool_set_rejects_a_database_name_at_build_time( + mock_client: MagicMock, databases_api: MagicMock +) -> None: + """Fail while wiring the tools up, not on the agent's first query.""" + databases_api.return_value.get_database.side_effect = ApiException( + status=404, reason="Not Found" + ) + with pytest.raises(KeyError): + make_hotdata_tools(mock_client, database_id="sf_airbnb") + + +def test_an_unscoped_tool_set_makes_no_lookup( + mock_client: MagicMock, databases_api: MagicMock +) -> None: + make_hotdata_tools(mock_client) + databases_api.return_value.get_database.assert_not_called() + + +@pytest.mark.parametrize( + "factory", + [ + lambda client, database_id: make_hotdata_describe_tables_tool( + client, database_id=database_id + ), + lambda client, database_id: make_hotdata_search_tool( + client, table=TABLE, column=COLUMN, database_id=database_id + ), + ], +) +def test_standalone_factories_resolve_by_id( + factory: object, + mock_client: MagicMock, + managed_db: ManagedDatabase, + databases_api: MagicMock, +) -> None: + factory(mock_client, managed_db.id) # type: ignore[operator] + databases_api.return_value.get_database.assert_called_once_with(managed_db.id) + + +# --- the agent-facing load tool ---------------------------------------------------- + + +def load_tool(client: MagicMock) -> object: + return {t.name: t for t in make_hotdata_tools(client)}["hotdata_load_managed_table"] + + +def test_load_tool_takes_a_database_id_argument(mock_client: MagicMock) -> None: + """The argument name is what the model sees in the schema, so it must say 'id'.""" + tool = load_tool(mock_client) + assert set(tool.args) == {"database_id", "table", "file", "schema_name"} # type: ignore[attr-defined] + + +def test_load_tool_resolves_the_agent_supplied_id_by_id( + mock_client: MagicMock, managed_db: ManagedDatabase, databases_api: MagicMock +) -> None: + mock_client.load_managed_table.return_value = LoadManagedTableResult( + connection_id=managed_db.default_connection_id, + schema_name="public", + table_name="orders", + row_count=1, + full_name=f"{managed_db.id}.public.orders", + ) + payload = json.loads( + load_tool(mock_client).invoke( # type: ignore[attr-defined] + { + "database_id": managed_db.id, + "table": "orders", + "file": "/tmp/orders.parquet", + } + ) + ) + databases_api.return_value.get_database.assert_called_once_with(managed_db.id) + # The resolved record addresses the load, so no name can select the overwrite target. + assert mock_client.load_managed_table.call_args.args[0] == managed_db + assert payload["row_count"] == 1 + + +def test_load_tool_rejects_a_database_name( + mock_client: MagicMock, databases_api: MagicMock +) -> None: + """An agent that passes a label gets an error naming the tool that yields ids.""" + databases_api.return_value.get_database.side_effect = ApiException( + status=404, reason="Not Found" + ) + with pytest.raises(KeyError, match="hotdata_list_managed_databases"): + load_tool(mock_client).invoke( # type: ignore[attr-defined] + {"database_id": "sales", "table": "orders", "file": "/tmp/orders.parquet"} + ) + mock_client.load_managed_table.assert_not_called() + + +def test_load_description_tells_the_model_where_ids_come_from(mock_client: MagicMock) -> None: + description = load_tool(mock_client).description or "" # type: ignore[attr-defined] + assert "hotdata_list_managed_databases" in description + assert "hotdata_create_managed_database" in description diff --git a/tests/test_schema.py b/tests/test_schema.py index c312ad8..e9f8242 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock import pytest -from hotdata_framework import QueryResult +from hotdata_framework import ManagedDatabase, QueryResult from hotdata_langchain.schema import ( DEFAULT_MAX_COLUMNS, @@ -158,11 +158,17 @@ def test_describe_reports_an_unknown_table_rather_than_empty_success() -> None: assert "no table named" in payload["error"] -def test_describe_scopes_queries_to_the_database() -> None: +def test_describe_scopes_queries_to_the_database(managed_db: ManagedDatabase) -> None: client = MagicMock() client.execute_sql.return_value = OVERVIEW - describe_tables_json(client, database="sf_airbnb") - assert client.execute_sql.call_args.kwargs == {"database": "sf_airbnb"} + describe_tables_json(client, database=managed_db) + assert client.execute_sql.call_args.kwargs == {"database": managed_db} + + +def test_describe_refuses_an_unresolved_database_scope() -> None: + """A bare string would reach the framework's by-name fallback.""" + with pytest.raises(TypeError, match="resolve_database_by_id"): + describe_tables_json(MagicMock(), database="sf_airbnb") # type: ignore[arg-type] # --- Tool surface ----------------------------------------------------------------- diff --git a/tests/test_search.py b/tests/test_search.py index 1186010..0e34764 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -6,7 +6,7 @@ from unittest.mock import MagicMock import pytest -from hotdata_framework import QueryResult +from hotdata_framework import ManagedDatabase, QueryResult from hotdata_langchain.search import ( DEFAULT_SEARCH_LIMIT, @@ -159,11 +159,23 @@ def test_bm25_search_json_returns_metadata_and_rows( def test_bm25_search_json_scopes_query_to_database( - mock_client: MagicMock, search_result: QueryResult + mock_client: MagicMock, search_result: QueryResult, managed_db: ManagedDatabase ) -> None: mock_client.execute_sql.return_value = search_result - bm25_search_json(mock_client, table=TABLE, column=COLUMN, query=QUERY, database="sf_airbnb") - assert mock_client.execute_sql.call_args.kwargs == {"database": "sf_airbnb"} + bm25_search_json(mock_client, table=TABLE, column=COLUMN, query=QUERY, database=managed_db) + assert mock_client.execute_sql.call_args.kwargs == {"database": managed_db} + + +def test_bm25_search_json_refuses_an_unresolved_database_scope(mock_client: MagicMock) -> None: + """A bare string would reach the framework's by-name fallback.""" + with pytest.raises(TypeError, match="resolve_database_by_id"): + bm25_search_json( + mock_client, + table=TABLE, + column=COLUMN, + query=QUERY, + database="sf_airbnb", # type: ignore[arg-type] + ) def test_bm25_search_json_truncates_rows_to_max_rows( @@ -352,17 +364,17 @@ def test_make_hotdata_tools_requires_both_search_arguments( def test_make_hotdata_tools_shares_database_scope_with_search( - mock_client: MagicMock, search_result: QueryResult + mock_client: MagicMock, search_result: QueryResult, managed_db: ManagedDatabase ) -> None: mock_client.execute_sql.return_value = search_result tools = { tool.name: tool for tool in make_hotdata_tools( - mock_client, database="sf_airbnb", search_table=TABLE, search_column=COLUMN + mock_client, database_id=managed_db, search_table=TABLE, search_column=COLUMN ) } tools["hotdata_search_text"].invoke({"query": QUERY}) - assert mock_client.execute_sql.call_args.kwargs == {"database": "sf_airbnb"} + assert mock_client.execute_sql.call_args.kwargs == {"database": managed_db} def test_make_hotdata_tools_shares_max_rows_with_search( diff --git a/tests/test_tools.py b/tests/test_tools.py index 7d0b400..2b0b6ba 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -24,9 +24,9 @@ def test_execute_sql_json(mock_client, sample_result): mock_client.execute_sql.assert_called_once_with("select 1", database=None) -def test_execute_sql_json_with_database(mock_client, sample_result): - execute_sql_json(mock_client, "select 1", database="my_db") - mock_client.execute_sql.assert_called_once_with("select 1", database="my_db") +def test_execute_sql_json_with_database(mock_client, sample_result, managed_db): + execute_sql_json(mock_client, "select 1", database=managed_db) + mock_client.execute_sql.assert_called_once_with("select 1", database=managed_db) def test_list_managed_databases_json(mock_client): @@ -52,7 +52,8 @@ def test_create_managed_database_delegates(mock_client): assert db.description == "sales" -def test_load_managed_table_delegates(mock_client): +def test_load_managed_table_delegates(mock_client, managed_db): + """The load addresses the resolved record, so no name can select the target.""" mock_client.load_managed_table.return_value = LoadManagedTableResult( connection_id="c1", schema_name="public", @@ -62,12 +63,12 @@ def test_load_managed_table_delegates(mock_client): ) loaded = load_managed_table( mock_client, - database="sales", + database_id=managed_db, table="orders", file="/tmp/orders.parquet", ) mock_client.load_managed_table.assert_called_once_with( - "sales", + managed_db, "orders", schema="public", file="/tmp/orders.parquet", @@ -75,7 +76,7 @@ def test_load_managed_table_delegates(mock_client): assert loaded.row_count == 3 -def test_make_hotdata_tools(mock_client, sample_result): +def test_make_hotdata_tools(mock_client, sample_result, managed_db, databases_api): mock_client.create_managed_database.return_value = ManagedDatabase( id="c1", description="sales", @@ -106,7 +107,7 @@ def test_make_hotdata_tools(mock_client, sample_result): json.loads( by_name["hotdata_load_managed_table"].invoke( { - "database": "sales", + "database_id": managed_db.id, "table": "orders", "file": "/tmp/orders.parquet", }