diff --git a/CLAUDE.md b/CLAUDE.md index 79e37b1..70fb9b1 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -136,7 +136,8 @@ The server requires Stacklet credentials configured through one of: Additional configuration via environment variables with `STACKLET_MCP_` prefix: - `STACKLET_MCP_DOWNLOADS_PATH` (default: system temp directory) - Directory for storing query result files - `STACKLET_MCP_DOWNLOADS_ENABLED` (default: true) - Write query result files; set false for hosted deployments -- `STACKLET_MCP_ASSETDB_DATASOURCE` (default: 1) - AssetDB data source ID +- `STACKLET_MCP_ASSETDB_DATASOURCE_NAME` (default: AssetDB) - Name the data source id is looked up from on first use +- `STACKLET_MCP_ASSETDB_DATASOURCE` (default: unset) - AssetDB data source ID, when a deployment needs to name one explicitly; unset means look it up by name - `STACKLET_MCP_ASSETDB_ALLOW_SAVE` (default: false) - Enable query save/update functionality - `STACKLET_MCP_ASSETDB_ALLOW_ARCHIVE` (default: false) - Enable query archiving functionality - `STACKLET_MCP_PLATFORM_ALLOW_MUTATIONS` (default: false) - Enable calling mutations in the Platform GraphQL API diff --git a/README.md b/README.md index b96a506..62528a8 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,8 @@ When the MCP is run from an agent, those can be set in the `"env"` section of th The following variables are available: - `STACKLET_MCP_DOWNLOADS_PATH`: directory for storing query result files (default: system temp directory) -- `STACKLET_MCP_ASSETDB_DATASOURCE`: the datasource ID for AssetDB in Redash (default: `1`) +- `STACKLET_MCP_ASSETDB_DATASOURCE_NAME`: the name of the AssetDB data source in Redash, which the id is looked up from on first use (default: `AssetDB`) +- `STACKLET_MCP_ASSETDB_DATASOURCE`: the datasource ID for AssetDB in Redash, for a deployment that needs to name one explicitly. Unset by default, meaning look it up by name — Redash assigns the id at creation, so it differs between deployments and cannot be assumed - `STACKLET_MCP_ASSETDB_ALLOW_SAVE`: whether to enable write operations in AssetDB (default: `false`) - `STACKLET_MCP_ASSETDB_ALLOW_ARCHIVE`: whether to enable query archiving functionality in AssetDB (default: `false`) - `STACKLET_MCP_PLATFORM_ALLOW_MUTATIONS`: whether to enable executing mutations in Platform API (default: `false`) diff --git a/stacklet/mcp/assetdb/redash.py b/stacklet/mcp/assetdb/redash.py index 65c7c9e..9ec3411 100644 --- a/stacklet/mcp/assetdb/redash.py +++ b/stacklet/mcp/assetdb/redash.py @@ -32,7 +32,8 @@ def __init__( self, credentials: StackletCredentials, server_state: ServerStateProtocol, - data_source_id: int = 1, + data_source_id: int | None = None, + data_source_name: str = "AssetDB", ) -> None: """ Initialize AssetDB client with Stacklet credentials. @@ -40,10 +41,13 @@ def __init__( Args: credentials: StackletCredentials object containing endpoint and id_token server_state: Server state for shared resource caching - data_source_id: ID of the Redash data source (default 1 for main AssetDB) + data_source_id: Redash data source id, when one is named explicitly + data_source_name: name to resolve an id from when data_source_id is None """ self.credentials = credentials - self.data_source_id = data_source_id + self.configured_data_source_id = data_source_id + self.data_source_name = data_source_name + self.server_state = server_state self.redash_url = self.credentials.service_endpoint("redash") transport = server_state.ensure_cached("HTTP_TRANSPORT", httpx.AsyncHTTPTransport) @@ -57,7 +61,61 @@ def __init__( @classmethod def get(cls, ctx: Context) -> Self: state = ctx.request_context.lifespan_context # type: ignore[union-attr] - return cls(StackletCredentials.get(ctx), state, SETTINGS.assetdb_datasource) + return cls( + StackletCredentials.get(ctx), + state, + SETTINGS.assetdb_datasource, + SETTINGS.assetdb_datasource_name, + ) + + async def get_data_source_id(self) -> int: + """The data source id to query, looked up by name unless one is configured. + + Redash assigns the id when the data source is created, so it depends on + what else that deployment happened to create first: the same "AssetDB" + source is 1 on one and 6 on another. Querying a non-existent id is not a + clean 404 either -- Redash's handler raises out of a bare `get_by_id`, so + the caller sees a 500 -- which is worth avoiding rather than documenting. + + Cached for the process once resolved, since the mapping only changes if + someone recreates the data source. The lookup runs as the calling user, + because Redash has no other credential here, but the answer is a property + of the org rather than of that user, so sharing it is sound. A failed + lookup is not cached, so a user who cannot see the source does not poison + it for everyone. + + Keyed by Redash URL as well as name. Every caller in a process shares one + endpoint today -- the hosted server reads it once at startup, and a local + one from its own config -- so this cannot currently collide. It is in the + key because the id means nothing without the deployment it came from, and + that is worth stating here rather than in a comment somewhere warning not + to serve two endpoints from one process. + """ + if self.configured_data_source_id is not None: + return self.configured_data_source_id + return await self.server_state.ensure_cached_async( + f"ASSETDB_DATASOURCE_ID:{self.redash_url}:{self.data_source_name}", + self._lookup_data_source_id, + ) + + async def _lookup_data_source_id(self) -> int: + sources = await self._make_request("GET", "api/data_sources") + for source in sources: + if source.get("name") == self.data_source_name: + return int(source["id"]) + raise AnnotatedError( + problem=f"No Redash data source named {self.data_source_name!r}", + likely_cause=( + "this deployment provisions AssetDB under a different name, or the " + "signed-in user cannot see it" + ), + next_steps=( + "check the data sources this deployment has, then set " + "STACKLET_MCP_ASSETDB_DATASOURCE_NAME to the right name, or " + "STACKLET_MCP_ASSETDB_DATASOURCE to its id" + ), + original_error=f"available: {sorted(str(s.get('name')) for s in sources)}", + ) async def _make_request(self, method: str, endpoint: str, **kwargs: Any) -> Any: """ @@ -163,7 +221,7 @@ async def execute_adhoc_query(self, query: str, max_age: int, timeout: int) -> Q """ payload = { "query": query, - "data_source_id": self.data_source_id, + "data_source_id": await self.get_data_source_id(), "max_age": max_age, "parameters": {}, "apply_auto_limit": True, @@ -268,7 +326,7 @@ async def create_query(self, upsert: QueryUpsert) -> Query: Returns: Complete query object with ID, timestamps, and metadata """ - payload = upsert.payload(data_source_id=self.data_source_id) + payload = upsert.payload(data_source_id=await self.get_data_source_id()) result = await self._make_request("POST", "api/queries", json=payload) return Query(**result) diff --git a/stacklet/mcp/settings.py b/stacklet/mcp/settings.py index 5c3d650..0d7d162 100644 --- a/stacklet/mcp/settings.py +++ b/stacklet/mcp/settings.py @@ -30,9 +30,23 @@ class Settings(BaseSettings): ), ) - assetdb_datasource: int = Field( - default=1, - description="AssetDB datasource", + assetdb_datasource: int | None = Field( + default=None, + description=( + "Redash data source id to query, for a deployment that needs to name one " + "explicitly. Left unset, the id is looked up by name -- see " + "assetdb_datasource_name -- which is what every Stacklet deployment wants: " + "Redash assigns the id at creation, so it differs between deployments and " + "cannot be assumed." + ), + ) + assetdb_datasource_name: str = Field( + default="AssetDB", + description=( + "Name of the Redash data source to resolve an id from, used when " + "assetdb_datasource is unset. Matches the name Stacklet provisions, which " + "is stable across deployments in the way the id is not." + ), ) assetdb_allow_save: bool = Field( default=False, diff --git a/tests/test_settings.py b/tests/test_settings.py index 30a804d..5ab45aa 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -21,7 +21,13 @@ def test_default_settings(self) -> None: """Test that default settings are properly initialized.""" assert SETTINGS.downloads_path.name.startswith("downloads") assert "pytest-of-" in str(SETTINGS.downloads_path) + # 1 is what the test fixture names, not the shipped default -- assert that + # separately, or this only tests the fixture. Off the field declarations + # rather than a fresh Settings(), which would read the ambient environment + # and fail for anyone who has STACKLET_MCP_ASSETDB_DATASOURCE exported. assert SETTINGS.assetdb_datasource == 1 + assert Settings.model_fields["assetdb_datasource"].default is None + assert Settings.model_fields["assetdb_datasource_name"].default == "AssetDB" assert SETTINGS.assetdb_allow_save is False assert SETTINGS.assetdb_allow_archive is False assert SETTINGS.platform_allow_mutations is False @@ -30,6 +36,7 @@ def test_env_prefix_configuration(self, monkeypatch: pytest.MonkeyPatch) -> None """Test that environment variables with STACKLET_MCP_ prefix are loaded.""" monkeypatch.setenv("STACKLET_MCP_DOWNLOADS_PATH", "/custom/path") monkeypatch.setenv("STACKLET_MCP_ASSETDB_DATASOURCE", "2") + monkeypatch.setenv("STACKLET_MCP_ASSETDB_DATASOURCE_NAME", "Warehouse") monkeypatch.setenv("STACKLET_MCP_ASSETDB_ALLOW_SAVE", "true") monkeypatch.setenv("STACKLET_MCP_ASSETDB_ALLOW_ARCHIVE", "true") monkeypatch.setenv("STACKLET_MCP_PLATFORM_ALLOW_MUTATIONS", "true") @@ -38,6 +45,7 @@ def test_env_prefix_configuration(self, monkeypatch: pytest.MonkeyPatch) -> None assert settings.downloads_path == Path("/custom/path") assert settings.assetdb_datasource == 2 + assert settings.assetdb_datasource_name == "Warehouse" assert settings.assetdb_allow_save is True assert settings.assetdb_allow_archive is True assert settings.platform_allow_mutations is True diff --git a/tests/test_tools_assetdb.py b/tests/test_tools_assetdb.py index 315390b..f303609 100644 --- a/tests/test_tools_assetdb.py +++ b/tests/test_tools_assetdb.py @@ -719,6 +719,97 @@ async def test_query(self): self.expect_post(self.post_data(query="SELECT 'whatever'"), self.result_response()), ) + def expect_data_sources(self, sources): + return ExpectRequest( + "https://redash.example.com/api/data_sources", + response=sources, + ) + + async def test_data_source_resolved_by_name(self, override_setting): + """With no id configured, the id comes from the named data source. + + Redash assigns ids in creation order, so the same "AssetDB" source is 1 on + one deployment and 6 on another. Querying an id that does not exist there + is not a clean 404 -- Redash raises out of its lookup and answers 500 -- + so the id has to be discovered rather than assumed. + """ + override_setting("assetdb_datasource", None) + + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_data_sources([{"id": 3, "name": "Other"}, {"id": 6, "name": "AssetDB"}]), + self.expect_post(self.post_data(data_source_id=6), self.result_response()), + ) + + async def test_data_source_lookup_is_cached(self, override_setting): + """Resolved once per process, not once per query.""" + override_setting("assetdb_datasource", None) + + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_data_sources([{"id": 6, "name": "AssetDB"}]), + self.expect_post(self.post_data(data_source_id=6), self.result_response()), + ) + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_post(self.post_data(data_source_id=6), self.result_response()), + ) + + async def test_a_configured_id_skips_the_lookup(self, override_setting): + """An explicit id is still honoured, and costs no request.""" + override_setting("assetdb_datasource", 123) + + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_post(self.post_data(data_source_id=123), self.result_response()), + ) + + async def test_a_missing_data_source_says_so(self, override_setting): + """The failure names what is there, rather than surfacing Redash's 500.""" + override_setting("assetdb_datasource", None) + + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_data_sources([{"id": 3, "name": "Other"}]), + expect_error=( + "No Redash data source named 'AssetDB'. This likely means this " + "deployment provisions AssetDB under a different name, or the signed-in " + "user cannot see it. Next steps: check the data sources this deployment " + "has, then set STACKLET_MCP_ASSETDB_DATASOURCE_NAME to the right name, " + "or STACKLET_MCP_ASSETDB_DATASOURCE to its id. " + "Original error: available: ['Other']" + ), + ) + + async def test_a_failed_lookup_is_not_cached(self, override_setting): + """A user who cannot see the source must not poison it for everyone. + + The cache is process-wide, so remembering a failure would turn one + under-permissioned caller -- or one bad moment -- into an outage for every + other caller until the task recycled. Only success is cached, so the next + query looks again. + """ + override_setting("assetdb_datasource", None) + + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_data_sources([{"id": 3, "name": "Other"}]), + expect_error=( + "No Redash data source named 'AssetDB'. This likely means this " + "deployment provisions AssetDB under a different name, or the signed-in " + "user cannot see it. Next steps: check the data sources this deployment " + "has, then set STACKLET_MCP_ASSETDB_DATASOURCE_NAME to the right name, " + "or STACKLET_MCP_ASSETDB_DATASOURCE to its id. " + "Original error: available: ['Other']" + ), + ) + # Looked up again rather than served from a remembered failure. + await self.assert_tool_call( + {"query": "SELECT 1"}, + self.expect_data_sources([{"id": 6, "name": "AssetDB"}]), + self.expect_post(self.post_data(data_source_id=6), self.result_response()), + ) + @json_guard_parametrize([-1, 0, 3600, 3600 * 24 * 365]) async def test_max_age(self, mangle, value): await self.assert_tool_call( diff --git a/tests/testing/settings.py b/tests/testing/settings.py index 952a8b1..86db9de 100644 --- a/tests/testing/settings.py +++ b/tests/testing/settings.py @@ -17,6 +17,10 @@ def default_settings(monkeypatch: pytest.MonkeyPatch, tmp_path_factory) -> Itera defaults = {name: field.default for name, field in Settings.model_fields.items()} # Use pytest temp directory for downloads defaults["downloads_path"] = tmp_path_factory.mktemp("downloads") + # Name a data source id, so a test that exercises a query is not also + # exercising the lookup that resolves one. The lookup has its own tests, + # which opt back in by setting this to None. + defaults["assetdb_datasource"] = 1 for attr, value in defaults.items(): monkeypatch.setattr(SETTINGS, attr, value)