fix: resolve the AssetDB data source by name

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

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`)

stacklet/mcp/assetdb/redash.py

@@ -32,18 +32,22 @@ 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.
 
         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)
 

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,

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

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(

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)