diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index 35c7f31..7e5821d 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -546,6 +546,7 @@ def _extract_connection_page(payload: object) -> tuple[list[dict], str | None]: def list_databricks_connections(workspace: str) -> list[dict]: env = build_databricks_cli_env(workspace) + profile_name = find_profile_name_for_host(workspace) connections: list[dict] = [] page_token: str | None = None seen_page_tokens: set[str] = set() @@ -561,6 +562,8 @@ def list_databricks_connections(workspace: str) -> list[dict]: "--output", "json", ] + if profile_name: + cmd.extend(["--profile", profile_name]) if page_token: cmd.extend(["--page-token", page_token]) @@ -608,6 +611,7 @@ def _extract_genie_spaces_page(payload: object) -> tuple[list[dict], str | None] def list_genie_spaces(workspace: str) -> list[dict]: env = build_databricks_cli_env(workspace) + profile_name = find_profile_name_for_host(workspace) spaces: list[dict] = [] page_token: str | None = None seen_page_tokens: set[str] = set() @@ -623,6 +627,8 @@ def list_genie_spaces(workspace: str) -> list[dict]: "--output", "json", ] + if profile_name: + cmd.extend(["--profile", profile_name]) if page_token: cmd.extend(["--page-token", page_token]) @@ -667,17 +673,21 @@ def _extract_apps_payload(payload: object) -> list[dict]: def list_databricks_apps(workspace: str) -> list[dict]: env = build_databricks_cli_env(workspace) + profile_name = find_profile_name_for_host(workspace) + cmd = [ + "databricks", + "apps", + "list", + "--limit", + "1000", + "--output", + "json", + ] + if profile_name: + cmd.extend(["--profile", profile_name]) try: result = run( - [ - "databricks", - "apps", - "list", - "--limit", - "1000", - "--output", - "json", - ], + cmd, capture_output=True, text=True, env=env, diff --git a/tests/test_databricks.py b/tests/test_databricks.py index dfcf17c..afc7389 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -328,6 +328,7 @@ def fake_run(args, **kwargs): return subprocess.CompletedProcess(args, 0, stdout=json.dumps(payload)) monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr(db_mod, "find_profile_name_for_host", lambda ws: "test-profile") assert list_databricks_connections(WS) == [ {"name": "confluence-mcp", "connection_type": "HTTP"}, @@ -341,6 +342,8 @@ def fake_run(args, **kwargs): "0", "--output", "json", + "--profile", + "test-profile", ] assert calls[0]["kwargs"]["env"]["DATABRICKS_HOST"] == WS assert calls[1]["args"][-2:] == ["--page-token", "next-page"] @@ -350,6 +353,7 @@ def fake_run(args, **kwargs): return subprocess.CompletedProcess(args, 0, stdout="not-json") monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr(db_mod, "find_profile_name_for_host", lambda ws: None) with pytest.raises(RuntimeError, match="invalid JSON"): list_databricks_connections(WS) @@ -371,6 +375,7 @@ def fake_run(args, **kwargs): return subprocess.CompletedProcess(args, 0, stdout=json.dumps(payload)) monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr(db_mod, "find_profile_name_for_host", lambda ws: "test-profile") assert list_genie_spaces(WS) == [ {"space_id": "space-1", "title": "First"}, @@ -384,6 +389,8 @@ def fake_run(args, **kwargs): "100", "--output", "json", + "--profile", + "test-profile", ] assert calls[0]["kwargs"]["env"]["DATABRICKS_HOST"] == WS assert calls[1]["args"][-2:] == ["--page-token", "next-page"] @@ -393,6 +400,7 @@ def fake_run(args, **kwargs): return subprocess.CompletedProcess(args, 0, stdout="not-json") monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr(db_mod, "find_profile_name_for_host", lambda ws: None) with pytest.raises(RuntimeError, match="invalid JSON"): list_genie_spaces(WS) @@ -413,6 +421,7 @@ def fake_run(args, **kwargs): return subprocess.CompletedProcess(args, 0, stdout=json.dumps(payload)) monkeypatch.setattr(db_mod, "run", fake_run) + monkeypatch.setattr(db_mod, "find_profile_name_for_host", lambda ws: "test-profile") assert list_databricks_apps(WS) == [ { @@ -428,6 +437,8 @@ def fake_run(args, **kwargs): "1000", "--output", "json", + "--profile", + "test-profile", ] assert calls[0]["kwargs"]["env"]["DATABRICKS_HOST"] == WS