From 5881e9dd4ae586c86dbc1f0d18315d6acd7d6bc0 Mon Sep 17 00:00:00 2001 From: Steven Tan Date: Tue, 10 Mar 2026 10:25:49 +0800 Subject: [PATCH 1/2] Add Related Skills and Resources sections to 18 skills Consolidated from previously individual PRs (#263-268) per maintainer feedback to submit as a single combined PR. Added Related Skills to 4 skills that were missing them: - databricks-config, databricks-dbsql, databricks-metric-views, databricks-spark-structured-streaming Added Resources (external doc links) to 18 skills that were missing them: - databricks-agent-bricks, databricks-aibi-dashboards, databricks-app-python, databricks-config, databricks-dbsql, databricks-docs, databricks-genie, databricks-lakebase-autoscale, databricks-lakebase-provisioned, databricks-metric-views, databricks-mlflow-evaluation, databricks-parsing, databricks-python-sdk, databricks-spark-declarative-pipelines, databricks-spark-structured-streaming, databricks-synthetic-data-gen, databricks-unstructured-pdf-generation, databricks-vector-search --- .../databricks-agent-bricks/SKILL.md | 7 +++++++ .../databricks-aibi-dashboards/SKILL.md | 6 ++++++ .../databricks-app-python/SKILL.md | 7 +++++++ databricks-skills/databricks-config/SKILL.md | 11 +++++++++++ databricks-skills/databricks-dbsql/SKILL.md | 17 +++++++++++++++++ databricks-skills/databricks-docs/SKILL.md | 7 +++++++ databricks-skills/databricks-genie/SKILL.md | 6 ++++++ .../databricks-lakebase-autoscale/SKILL.md | 6 ++++++ .../databricks-lakebase-provisioned/SKILL.md | 6 ++++++ .../databricks-metric-views/SKILL.md | 6 ++++++ .../databricks-mlflow-evaluation/SKILL.md | 7 +++++++ databricks-skills/databricks-parsing/SKILL.md | 6 ++++++ .../databricks-python-sdk/SKILL.md | 7 +++++++ .../SKILL.md | 8 ++++++++ .../SKILL.md | 15 +++++++++++++++ .../databricks-synthetic-data-gen/SKILL.md | 6 ++++++ .../SKILL.md | 6 ++++++ .../databricks-vector-search/SKILL.md | 7 +++++++ 18 files changed, 141 insertions(+) diff --git a/databricks-skills/databricks-agent-bricks/SKILL.md b/databricks-skills/databricks-agent-bricks/SKILL.md index 04be7dad..551a6ded 100644 --- a/databricks-skills/databricks-agent-bricks/SKILL.md +++ b/databricks-skills/databricks-agent-bricks/SKILL.md @@ -204,6 +204,13 @@ manage_mas( - **[databricks-model-serving](../databricks-model-serving/SKILL.md)** - Deploy custom agent endpoints used as MAS agents - **[databricks-vector-search](../databricks-vector-search/SKILL.md)** - Build vector indexes for RAG applications paired with KAs +## Resources + +- [Databricks Agent Bricks](https://docs.databricks.com/generative-ai/agent-bricks/index.html) +- [Knowledge Assistants](https://docs.databricks.com/generative-ai/agent-bricks/knowledge-assistants.html) +- [Supervisor Agents](https://docs.databricks.com/generative-ai/agent-bricks/supervisor-agents.html) +- [Genie Spaces](https://docs.databricks.com/genie/index.html) + ## See Also - `1-knowledge-assistants.md` - Detailed KA patterns and examples diff --git a/databricks-skills/databricks-aibi-dashboards/SKILL.md b/databricks-skills/databricks-aibi-dashboards/SKILL.md index e722cb0d..46231cb6 100644 --- a/databricks-skills/databricks-aibi-dashboards/SKILL.md +++ b/databricks-skills/databricks-aibi-dashboards/SKILL.md @@ -914,6 +914,12 @@ print(result["url"]) - Use **separate text widgets** for title and subtitle at different y positions - Example: title at y=0 with height=1, subtitle at y=1 with height=1 +## Resources + +- [AI/BI Dashboards](https://docs.databricks.com/dashboards/index.html) +- [Lakeview API Reference](https://docs.databricks.com/api/workspace/lakeview) +- [Dashboard Embedding](https://docs.databricks.com/dashboards/embed.html) + ## Related Skills - **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - for querying the underlying data and system tables diff --git a/databricks-skills/databricks-app-python/SKILL.md b/databricks-skills/databricks-app-python/SKILL.md index d5bb4d27..7a3e1125 100644 --- a/databricks-skills/databricks-app-python/SKILL.md +++ b/databricks-skills/databricks-app-python/SKILL.md @@ -199,6 +199,13 @@ class EntityIn(BaseModel): - **[app.yaml Reference](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/app-runtime)** — command and env config - **[System Environment](https://docs.databricks.com/aws/en/dev-tools/databricks-apps/system-env)** — pre-installed packages, runtime details +## Resources + +- [Databricks Apps Overview](https://docs.databricks.com/dev-tools/databricks-apps/index.html) +- [App Configuration (app.yaml)](https://docs.databricks.com/dev-tools/databricks-apps/configuration.html) +- [App Resources and Permissions](https://docs.databricks.com/dev-tools/databricks-apps/app-resources.html) +- [Streamlit on Databricks](https://docs.databricks.com/dev-tools/databricks-apps/streamlit.html) + ## Related Skills - **[databricks-app-apx](../databricks-app-apx/SKILL.md)** - full-stack apps with FastAPI + React diff --git a/databricks-skills/databricks-config/SKILL.md b/databricks-skills/databricks-config/SKILL.md index 88382c33..a74c50b5 100644 --- a/databricks-skills/databricks-config/SKILL.md +++ b/databricks-skills/databricks-config/SKILL.md @@ -20,3 +20,14 @@ Use the `manage_workspace` MCP tool for all workspace operations. Do NOT edit `~ 4. Present the result. For `status`/`switch`/`login`: show host, profile, username. For `list`: formatted table with the active profile marked. > **Note:** The switch is session-scoped — it resets on MCP server restart. For permanent profile setup, use `databricks auth login -p ` and update `~/.databrickscfg` with `cluster_id` or `serverless_compute_id = auto`. + +## Related Skills + +- **[databricks-python-sdk](../databricks-python-sdk/SKILL.md)** - SDK patterns that depend on workspace configuration +- **[databricks-docs](../databricks-docs/SKILL.md)** - Documentation reference for authentication methods + +## Resources + +- [Databricks CLI Authentication](https://docs.databricks.com/dev-tools/cli/authentication.html) +- [Unified Auth Configuration](https://docs.databricks.com/dev-tools/auth/unified-auth.html) +- [Environment Variables](https://docs.databricks.com/dev-tools/auth/env-variables.html) diff --git a/databricks-skills/databricks-dbsql/SKILL.md b/databricks-skills/databricks-dbsql/SKILL.md index 24bf2694..5864ffcc 100644 --- a/databricks-skills/databricks-dbsql/SKILL.md +++ b/databricks-skills/databricks-dbsql/SKILL.md @@ -298,3 +298,20 @@ Load these for detailed syntax, full parameter lists, and advanced patterns: - **Define PK/FK constraints** on dimensional models for query optimization - **Use `COLLATE UTF8_LCASE`** for user-facing string columns that need case-insensitive search - **Use MCP tools** (`execute_sql`, `execute_sql_multi`) to test and validate all SQL before deploying + +## Related Skills + +- **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Governance, grants, and catalog management for SQL objects +- **[databricks-aibi-dashboards](../databricks-aibi-dashboards/SKILL.md)** - Building dashboards that consume DBSQL queries +- **[databricks-metric-views](../databricks-metric-views/SKILL.md)** - Defining reusable business metrics in SQL +- **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** - ETL pipelines using SQL and materialized views +- **[databricks-genie](../databricks-genie/SKILL.md)** - Natural language SQL exploration via Genie Spaces + +## Resources + +- [Databricks SQL Language Reference](https://docs.databricks.com/sql/language-manual/index.html) +- [AI Functions](https://docs.databricks.com/large-language-models/ai-functions.html) +- [Pipe Syntax](https://docs.databricks.com/sql/language-manual/sql-ref-pipe-syntax.html) +- [Materialized Views](https://docs.databricks.com/sql/language-manual/sql-ref-materialized-views.html) +- [SQL Scripting](https://docs.databricks.com/sql/language-manual/sql-ref-sql-scripting.html) +- [Geospatial H3 Functions](https://docs.databricks.com/sql/language-manual/sql-ref-h3-geospatial-functions.html) diff --git a/databricks-skills/databricks-docs/SKILL.md b/databricks-skills/databricks-docs/SKILL.md index 54bb157f..3f3e71a2 100644 --- a/databricks-skills/databricks-docs/SKILL.md +++ b/databricks-skills/databricks-docs/SKILL.md @@ -55,6 +55,13 @@ The llms.txt file is organized by category: 2. Read the specific docs to understand the feature 3. Determine which skill/tools apply, then use them +## Resources + +- [Databricks Documentation](https://docs.databricks.com/) +- [Databricks llms.txt](https://docs.databricks.com/llms.txt) +- [Release Notes](https://docs.databricks.com/release-notes/index.html) +- [REST API Reference](https://docs.databricks.com/api/workspace/introduction) + ## Related Skills - **[databricks-python-sdk](../databricks-python-sdk/SKILL.md)** - SDK patterns for programmatic Databricks access diff --git a/databricks-skills/databricks-genie/SKILL.md b/databricks-skills/databricks-genie/SKILL.md index e5b32b6e..6d1c2619 100644 --- a/databricks-skills/databricks-genie/SKILL.md +++ b/databricks-skills/databricks-genie/SKILL.md @@ -124,3 +124,9 @@ Use these skills in sequence: - **[databricks-synthetic-data-gen](../databricks-synthetic-data-gen/SKILL.md)** - Generate raw parquet data to populate tables for Genie - **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** - Build bronze/silver/gold tables consumed by Genie Spaces - **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Manage the catalogs, schemas, and tables Genie queries + +## Resources + +- [Genie Spaces](https://docs.databricks.com/genie/index.html) +- [Genie Conversation API](https://docs.databricks.com/api/workspace/genie) +- [Genie Best Practices](https://docs.databricks.com/genie/best-practices.html) diff --git a/databricks-skills/databricks-lakebase-autoscale/SKILL.md b/databricks-skills/databricks-lakebase-autoscale/SKILL.md index 50ba1df7..1a805d88 100644 --- a/databricks-skills/databricks-lakebase-autoscale/SKILL.md +++ b/databricks-skills/databricks-lakebase-autoscale/SKILL.md @@ -292,3 +292,9 @@ These features are NOT yet supported in Lakebase Autoscaling: - **[databricks-python-sdk](../databricks-python-sdk/SKILL.md)** - SDK used for project management and token generation - **[databricks-asset-bundles](../databricks-asset-bundles/SKILL.md)** - deploying apps with Lakebase resources - **[databricks-jobs](../databricks-jobs/SKILL.md)** - scheduling reverse ETL sync jobs + +## Resources + +- [Lakebase Overview](https://docs.databricks.com/database/lakebase/index.html) +- [Lakebase Autoscale](https://docs.databricks.com/database/lakebase/autoscale.html) +- [Lakebase Branching](https://docs.databricks.com/database/lakebase/branching.html) diff --git a/databricks-skills/databricks-lakebase-provisioned/SKILL.md b/databricks-skills/databricks-lakebase-provisioned/SKILL.md index b2b404aa..8072de50 100644 --- a/databricks-skills/databricks-lakebase-provisioned/SKILL.md +++ b/databricks-skills/databricks-lakebase-provisioned/SKILL.md @@ -306,3 +306,9 @@ databricks database start-database-instance --name my-lakebase-instance - **[databricks-python-sdk](../databricks-python-sdk/SKILL.md)** - SDK used for instance management and token generation - **[databricks-asset-bundles](../databricks-asset-bundles/SKILL.md)** - deploying apps with Lakebase resources - **[databricks-jobs](../databricks-jobs/SKILL.md)** - scheduling reverse ETL sync jobs + +## Resources + +- [Lakebase Overview](https://docs.databricks.com/database/lakebase/index.html) +- [Lakebase Provisioned](https://docs.databricks.com/database/lakebase/provisioned.html) +- [Lakebase Sync (Reverse ETL)](https://docs.databricks.com/database/lakebase/sync.html) diff --git a/databricks-skills/databricks-metric-views/SKILL.md b/databricks-skills/databricks-metric-views/SKILL.md index d3f5834c..a87dbc13 100644 --- a/databricks-skills/databricks-metric-views/SKILL.md +++ b/databricks-skills/databricks-metric-views/SKILL.md @@ -219,6 +219,12 @@ Metric views work natively with: - **SQL Editor** - Direct SQL querying with MEASURE() - **Catalog Explorer UI** - Visual creation and browsing +## Related Skills + +- **[databricks-dbsql](../databricks-dbsql/SKILL.md)** - SQL features used in metric view definitions +- **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Catalog governance for metric view objects +- **[databricks-aibi-dashboards](../databricks-aibi-dashboards/SKILL.md)** - Dashboards that consume metric views + ## Resources - [Metric Views Documentation](https://docs.databricks.com/en/metric-views/) diff --git a/databricks-skills/databricks-mlflow-evaluation/SKILL.md b/databricks-skills/databricks-mlflow-evaluation/SKILL.md index 45db5f61..6a791f8f 100644 --- a/databricks-skills/databricks-mlflow-evaluation/SKILL.md +++ b/databricks-skills/databricks-mlflow-evaluation/SKILL.md @@ -146,3 +146,10 @@ See `GOTCHAS.md` for complete list. - **[databricks-agent-bricks](../databricks-agent-bricks/SKILL.md)** - Building agents that can be evaluated with this skill - **[databricks-python-sdk](../databricks-python-sdk/SKILL.md)** - SDK patterns used alongside MLflow APIs - **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Unity Catalog tables for managed evaluation datasets + +## Resources + +- [MLflow GenAI Evaluation](https://mlflow.org/docs/latest/llms/llm-evaluate/index.html) +- [MLflow Scorers](https://mlflow.org/docs/latest/llms/llm-evaluate/llm-as-judge.html) +- [MLflow Tracing](https://mlflow.org/docs/latest/llms/tracing/index.html) +- [Databricks Agent Evaluation](https://docs.databricks.com/generative-ai/agent-evaluation/index.html) diff --git a/databricks-skills/databricks-parsing/SKILL.md b/databricks-skills/databricks-parsing/SKILL.md index 96b00b32..c951ab70 100644 --- a/databricks-skills/databricks-parsing/SKILL.md +++ b/databricks-skills/databricks-parsing/SKILL.md @@ -273,3 +273,9 @@ error_status[] -- errors per page (if any) - **[databricks-agent-bricks](../databricks-agent-bricks/SKILL.md)** — Pre-built Knowledge Assistants (out-of-the-box RAG without custom parsing) - **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** — Production pipelines for batch document processing - **[databricks-dbsql](../databricks-dbsql/SKILL.md)** — Full AI functions reference including `ai_query`, `ai_extract`, `ai_classify` + +## Resources + +- [Document Parsing with Databricks](https://docs.databricks.com/generative-ai/tutorials/ai-cookbook/nbs/1-pdf-chunk-extract.html) +- [Unstructured Data in Unity Catalog Volumes](https://docs.databricks.com/ingestion/file-upload/upload-to-volume.html) +- [Foundation Model APIs](https://docs.databricks.com/machine-learning/foundation-models/index.html) diff --git a/databricks-skills/databricks-python-sdk/SKILL.md b/databricks-skills/databricks-python-sdk/SKILL.md index 1365666a..a4e75d3e 100644 --- a/databricks-skills/databricks-python-sdk/SKILL.md +++ b/databricks-skills/databricks-python-sdk/SKILL.md @@ -623,3 +623,10 @@ If I'm unsure about a method, I should: - **[databricks-model-serving](../databricks-model-serving/SKILL.md)** - serving endpoint management - **[databricks-vector-search](../databricks-vector-search/SKILL.md)** - vector index operations - **[databricks-lakebase-provisioned](../databricks-lakebase-provisioned/SKILL.md)** - managed PostgreSQL via SDK + +## Resources + +- [Databricks SDK for Python](https://docs.databricks.com/dev-tools/sdk-python.html) +- [SDK API Reference](https://databricks-sdk-py.readthedocs.io/en/latest/) +- [Databricks Connect](https://docs.databricks.com/dev-tools/databricks-connect/index.html) +- [REST API Reference](https://docs.databricks.com/api/workspace/introduction) diff --git a/databricks-skills/databricks-spark-declarative-pipelines/SKILL.md b/databricks-skills/databricks-spark-declarative-pipelines/SKILL.md index feea8afd..b8999f37 100644 --- a/databricks-skills/databricks-spark-declarative-pipelines/SKILL.md +++ b/databricks-skills/databricks-spark-declarative-pipelines/SKILL.md @@ -575,3 +575,11 @@ For advanced configuration options (development mode, continuous pipelines, cust - **[databricks-asset-bundles](../databricks-asset-bundles/SKILL.md)** - for multi-environment deployment of pipeline projects - **[databricks-synthetic-data-gen](../databricks-synthetic-data-gen/SKILL.md)** - for generating test data to feed into pipelines - **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - for catalog/schema/volume management and governance + +## Resources + +- [Lakeflow Declarative Pipelines](https://docs.databricks.com/delta-live-tables/index.html) +- [Streaming Tables](https://docs.databricks.com/sql/language-manual/sql-ref-streaming-tables.html) +- [Materialized Views](https://docs.databricks.com/sql/language-manual/sql-ref-materialized-views.html) +- [Change Data Capture](https://docs.databricks.com/delta-live-tables/cdc.html) +- [Auto Loader](https://docs.databricks.com/ingestion/auto-loader/index.html) diff --git a/databricks-skills/databricks-spark-structured-streaming/SKILL.md b/databricks-skills/databricks-spark-structured-streaming/SKILL.md index b1f59306..22a93b30 100644 --- a/databricks-skills/databricks-spark-structured-streaming/SKILL.md +++ b/databricks-skills/databricks-spark-structured-streaming/SKILL.md @@ -54,6 +54,21 @@ df.writeStream \ |-------|-------------|-----------| | **Production Checklist** | Comprehensive best practices | See [streaming-best-practices.md](streaming-best-practices.md) | +## Related Skills + +- **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** - Declarative streaming tables and DLT pipelines +- **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Catalog governance for streaming tables +- **[databricks-zerobus-ingest](../databricks-zerobus-ingest/SKILL.md)** - Near real-time ingestion into Delta tables via gRPC +- **[databricks-dbsql](../databricks-dbsql/SKILL.md)** - SQL queries on streaming table outputs + +## Resources + +- [Structured Streaming Guide](https://docs.databricks.com/structured-streaming/index.html) +- [Auto Loader](https://docs.databricks.com/ingestion/auto-loader/index.html) +- [Streaming Tables](https://docs.databricks.com/sql/language-manual/sql-ref-streaming-tables.html) +- [RocksDB State Store](https://docs.databricks.com/structured-streaming/rocksdb-state-store.html) +- [Stream-Stream Joins](https://docs.databricks.com/structured-streaming/stream-stream-joins.html) + ## Production Checklist - [ ] Checkpoint location is persistent (UC volumes, not DBFS) diff --git a/databricks-skills/databricks-synthetic-data-gen/SKILL.md b/databricks-skills/databricks-synthetic-data-gen/SKILL.md index 5bd95e58..acc4a125 100644 --- a/databricks-skills/databricks-synthetic-data-gen/SKILL.md +++ b/databricks-skills/databricks-synthetic-data-gen/SKILL.md @@ -256,3 +256,9 @@ See [references/5-output-formats.md](references/5-output-formats.md) for detaile | `F.window` vs `Window` confusion | Use `from pyspark.sql.window import Window` for `row_number()`, `rank()`, etc. `F.window` is for streaming only. | See [references/6-troubleshooting.md](references/6-troubleshooting.md) for full troubleshooting guide. + +## Resources + +- [Faker Library](https://faker.readthedocs.io/en/master/) +- [PySpark DataFrame API](https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/dataframe.html) +- [Unity Catalog Volumes](https://docs.databricks.com/connect/unity-catalog/volumes.html) diff --git a/databricks-skills/databricks-unstructured-pdf-generation/SKILL.md b/databricks-skills/databricks-unstructured-pdf-generation/SKILL.md index ee9abf05..291905e1 100644 --- a/databricks-skills/databricks-unstructured-pdf-generation/SKILL.md +++ b/databricks-skills/databricks-unstructured-pdf-generation/SKILL.md @@ -192,3 +192,9 @@ AZURE_OPENAI_DEPLOYMENT=gpt-4o - **[databricks-vector-search](../databricks-vector-search/SKILL.md)** - Index generated documents for semantic search and RAG - **[databricks-synthetic-data-gen](../databricks-synthetic-data-gen/SKILL.md)** - Generate structured tabular data (complement to unstructured PDFs) - **[databricks-mlflow-evaluation](../databricks-mlflow-evaluation/SKILL.md)** - Evaluate RAG systems using the generated question/guideline pairs + +## Resources + +- [Unity Catalog Volumes](https://docs.databricks.com/connect/unity-catalog/volumes.html) +- [ReportLab Documentation](https://docs.reportlab.com/) +- [Document Parsing with Databricks](https://docs.databricks.com/generative-ai/tutorials/ai-cookbook/nbs/1-pdf-chunk-extract.html) diff --git a/databricks-skills/databricks-vector-search/SKILL.md b/databricks-skills/databricks-vector-search/SKILL.md index 22bc276d..0cb90c92 100644 --- a/databricks-skills/databricks-vector-search/SKILL.md +++ b/databricks-skills/databricks-vector-search/SKILL.md @@ -360,3 +360,10 @@ The following MCP tools are available for managing Vector Search infrastructure. - **[databricks-unstructured-pdf-generation](../databricks-unstructured-pdf-generation/SKILL.md)** - Generate documents to index in Vector Search - **[databricks-unity-catalog](../databricks-unity-catalog/SKILL.md)** - Manage the catalogs and tables that back Delta Sync indexes - **[databricks-spark-declarative-pipelines](../databricks-spark-declarative-pipelines/SKILL.md)** - Build Delta tables used as Vector Search sources + +## Resources + +- [Vector Search Overview](https://docs.databricks.com/generative-ai/vector-search/index.html) +- [Vector Search API Reference](https://docs.databricks.com/api/workspace/vectorsearchindexes) +- [Embedding Models](https://docs.databricks.com/machine-learning/foundation-models/index.html) +- [RAG Tutorial](https://docs.databricks.com/generative-ai/tutorials/ai-cookbook/index.html) From 828b05efadf2c2e51f082d7ec55fc3b1b60dcf12 Mon Sep 17 00:00:00 2001 From: Steven Tan Date: Tue, 10 Mar 2026 10:31:11 +0800 Subject: [PATCH 2/2] Fix pre-existing ruff format issues in auth.py and test_sql.py These files were already failing the CI Lint & Format check on main. --- .../databricks_tools_core/auth.py | 8 +-- databricks-tools-core/tests/unit/test_sql.py | 54 ++++++++----------- 2 files changed, 23 insertions(+), 39 deletions(-) diff --git a/databricks-tools-core/databricks_tools_core/auth.py b/databricks-tools-core/databricks_tools_core/auth.py index 21913983..c3db9fb4 100644 --- a/databricks-tools-core/databricks_tools_core/auth.py +++ b/databricks-tools-core/databricks_tools_core/auth.py @@ -160,9 +160,7 @@ def get_workspace_client() -> WorkspaceClient: # Cross-workspace: explicit token overrides env OAuth so tool operations # target the caller-specified workspace instead of the app's own workspace if force and host and token: - return tag_client( - WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs) - ) + return tag_client(WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs)) # In Databricks Apps (OAuth credentials in env), explicitly use OAuth M2M. # Setting auth_type="oauth-m2m" prevents the SDK from also reading @@ -185,9 +183,7 @@ def get_workspace_client() -> WorkspaceClient: # Development mode: use explicit token if provided if host and token: - return tag_client( - WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs) - ) + return tag_client(WorkspaceClient(host=host, token=token, auth_type="pat", **product_kwargs)) if host: return tag_client(WorkspaceClient(host=host, **product_kwargs)) diff --git a/databricks-tools-core/tests/unit/test_sql.py b/databricks-tools-core/tests/unit/test_sql.py index d1b661c6..42137ba5 100644 --- a/databricks-tools-core/tests/unit/test_sql.py +++ b/databricks-tools-core/tests/unit/test_sql.py @@ -121,8 +121,7 @@ def test_executor_without_query_tags_omits_from_api(self, mock_get_client): assert "query_tags" not in call_kwargs -def _make_warehouse(id, name, state, creator_name="other@example.com", - enable_serverless_compute=False): +def _make_warehouse(id, name, state, creator_name="other@example.com", enable_serverless_compute=False): """Helper to create a mock warehouse object.""" w = mock.Mock() w.id = id @@ -141,33 +140,29 @@ class TestSortWithinTier: def test_serverless_first(self): """Serverless warehouses should come before classic ones.""" classic = _make_warehouse("c1", "Classic WH", State.RUNNING) - serverless = _make_warehouse("s1", "Serverless WH", State.RUNNING, - enable_serverless_compute=True) + serverless = _make_warehouse("s1", "Serverless WH", State.RUNNING, enable_serverless_compute=True) result = _sort_within_tier([classic, serverless], current_user=None) assert result[0].id == "s1" assert result[1].id == "c1" def test_serverless_before_user_owned(self): """Serverless should be preferred over user-owned classic.""" - classic_owned = _make_warehouse("c1", "My WH", State.RUNNING, - creator_name="me@example.com") - serverless_other = _make_warehouse("s1", "Other WH", State.RUNNING, - creator_name="other@example.com", - enable_serverless_compute=True) - result = _sort_within_tier([classic_owned, serverless_other], - current_user="me@example.com") + classic_owned = _make_warehouse("c1", "My WH", State.RUNNING, creator_name="me@example.com") + serverless_other = _make_warehouse( + "s1", "Other WH", State.RUNNING, creator_name="other@example.com", enable_serverless_compute=True + ) + result = _sort_within_tier([classic_owned, serverless_other], current_user="me@example.com") assert result[0].id == "s1" def test_serverless_user_owned_first(self): """Among serverless, user-owned should come first.""" - serverless_other = _make_warehouse("s1", "Other Serverless", State.RUNNING, - creator_name="other@example.com", - enable_serverless_compute=True) - serverless_owned = _make_warehouse("s2", "My Serverless", State.RUNNING, - creator_name="me@example.com", - enable_serverless_compute=True) - result = _sort_within_tier([serverless_other, serverless_owned], - current_user="me@example.com") + serverless_other = _make_warehouse( + "s1", "Other Serverless", State.RUNNING, creator_name="other@example.com", enable_serverless_compute=True + ) + serverless_owned = _make_warehouse( + "s2", "My Serverless", State.RUNNING, creator_name="me@example.com", enable_serverless_compute=True + ) + result = _sort_within_tier([serverless_other, serverless_owned], current_user="me@example.com") assert result[0].id == "s2" assert result[1].id == "s1" @@ -177,8 +172,7 @@ def test_empty_list(self): def test_no_current_user(self): """Without a current user, only serverless preference applies.""" classic = _make_warehouse("c1", "Classic", State.RUNNING) - serverless = _make_warehouse("s1", "Serverless", State.RUNNING, - enable_serverless_compute=True) + serverless = _make_warehouse("s1", "Serverless", State.RUNNING, enable_serverless_compute=True) result = _sort_within_tier([classic, serverless], current_user=None) assert result[0].id == "s1" @@ -186,14 +180,12 @@ def test_no_current_user(self): class TestGetBestWarehouseServerless: """Tests for serverless preference in get_best_warehouse.""" - @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", - return_value="me@example.com") + @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com") @mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client") def test_prefers_serverless_within_running_shared(self, mock_client_fn, mock_user): """Among running shared warehouses, serverless should be picked.""" classic_shared = _make_warehouse("c1", "Shared WH", State.RUNNING) - serverless_shared = _make_warehouse("s1", "Shared Serverless", State.RUNNING, - enable_serverless_compute=True) + serverless_shared = _make_warehouse("s1", "Shared Serverless", State.RUNNING, enable_serverless_compute=True) mock_client = mock.Mock() mock_client.warehouses.list.return_value = [classic_shared, serverless_shared] mock_client_fn.return_value = mock_client @@ -201,14 +193,12 @@ def test_prefers_serverless_within_running_shared(self, mock_client_fn, mock_use result = get_best_warehouse() assert result == "s1" - @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", - return_value="me@example.com") + @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com") @mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client") def test_prefers_serverless_within_running_other(self, mock_client_fn, mock_user): """Among running non-shared warehouses, serverless should be picked.""" classic = _make_warehouse("c1", "My WH", State.RUNNING) - serverless = _make_warehouse("s1", "Fast WH", State.RUNNING, - enable_serverless_compute=True) + serverless = _make_warehouse("s1", "Fast WH", State.RUNNING, enable_serverless_compute=True) mock_client = mock.Mock() mock_client.warehouses.list.return_value = [classic, serverless] mock_client_fn.return_value = mock_client @@ -216,14 +206,12 @@ def test_prefers_serverless_within_running_other(self, mock_client_fn, mock_user result = get_best_warehouse() assert result == "s1" - @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", - return_value="me@example.com") + @mock.patch("databricks_tools_core.sql.warehouse.get_current_username", return_value="me@example.com") @mock.patch("databricks_tools_core.sql.warehouse.get_workspace_client") def test_tier_order_preserved_over_serverless(self, mock_client_fn, mock_user): """A running shared classic should still beat a stopped serverless.""" running_shared_classic = _make_warehouse("c1", "Shared WH", State.RUNNING) - stopped_serverless = _make_warehouse("s1", "Fast WH", State.STOPPED, - enable_serverless_compute=True) + stopped_serverless = _make_warehouse("s1", "Fast WH", State.STOPPED, enable_serverless_compute=True) mock_client = mock.Mock() mock_client.warehouses.list.return_value = [stopped_serverless, running_shared_classic] mock_client_fn.return_value = mock_client