Merge pull request #75 from databricks-solutions/feature/telemetry

Add User-Agent telemetry to all Databricks SDK calls

Author: hiydavid

backend/_telemetry.py

@@ -0,0 +1,4 @@
+"""Databricks SDK telemetry constants for the genie-workbench product."""
+
+PRODUCT_NAME = "genie-workbench"
+PRODUCT_VERSION = "0.1.0"  # Keep in sync with pyproject.toml

backend/services/auth.py

@@ -16,6 +16,8 @@
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.config import Config
 
+from backend._telemetry import PRODUCT_NAME, PRODUCT_VERSION
+
 logger = logging.getLogger(__name__)
 
 # Singleton client for local dev (or fallback when no user token is available)
@@ -55,7 +57,7 @@ def set_obo_user_token(token: str) -> None:
         client_id=None,
         client_secret=None,  # gitleaks:allow
     )
-    client = WorkspaceClient(config=cfg)
+    client = WorkspaceClient(config=cfg, product=PRODUCT_NAME, product_version=PRODUCT_VERSION)
     _obo_client.set(client)
     logger.debug("OBO client set for current request (host=%s, auth=%s)", host, cfg.auth_type)
 
@@ -70,7 +72,7 @@ def _get_default_client() -> WorkspaceClient:
     global _client, _auth_logged
 
     if _client is None:
-        _client = WorkspaceClient()
+        _client = WorkspaceClient(product=PRODUCT_NAME, product_version=PRODUCT_VERSION)
 
         if not _auth_logged:
             logger.info("=== Databricks SDK Authentication ===")

packages/genie-space-optimizer/src/genie_space_optimizer/_telemetry.py

@@ -0,0 +1,5 @@
+"""Databricks SDK telemetry constants for the genie-space-optimizer product."""
+
+from genie_space_optimizer import __version__ as PRODUCT_VERSION
+
+PRODUCT_NAME = "genie-space-optimizer"

packages/genie-space-optimizer/src/genie_space_optimizer/_workspace_client.py

@@ -0,0 +1,14 @@
+"""Factory for creating telemetry-tagged WorkspaceClient instances."""
+
+from databricks.sdk import WorkspaceClient
+
+from genie_space_optimizer._telemetry import PRODUCT_NAME, PRODUCT_VERSION
+
+
+def make_workspace_client(**kwargs) -> WorkspaceClient:
+    """Create a WorkspaceClient tagged with GSO product telemetry."""
+    return WorkspaceClient(
+        product=PRODUCT_NAME,
+        product_version=PRODUCT_VERSION,
+        **kwargs,
+    )

packages/genie-space-optimizer/src/genie_space_optimizer/backend/core/_defaults.py

@@ -6,6 +6,8 @@
 from databricks.sdk import WorkspaceClient
 from fastapi import Depends, FastAPI, Request
 
+from genie_space_optimizer._workspace_client import make_workspace_client
+
 from ._base import LifespanDependency
 from ._config import AppConfig, logger
 from ._headers import HeadersDependency
@@ -26,7 +28,7 @@ def __call__(request: Request) -> AppConfig:
 class _WorkspaceClientDependency(LifespanDependency):
     @asynccontextmanager
     async def lifespan(self, app: FastAPI) -> AsyncGenerator[None, None]:
-        app.state.workspace_client = WorkspaceClient()
+        app.state.workspace_client = make_workspace_client()
         yield
 
     @staticmethod
@@ -49,10 +51,10 @@ def _get_user_ws(
             "OBO token not available — falling back to service principal credentials. "
             "User-scoped permissions will not apply."
         )
-        return WorkspaceClient()
+        return make_workspace_client()
 
     host = os.environ.get("DATABRICKS_HOST", "")
-    return WorkspaceClient(
+    return make_workspace_client(
         host=host,
         token=headers.token.get_secret_value(),
         auth_type="pat",

packages/genie-space-optimizer/src/genie_space_optimizer/backend/core/lakebase.py

@@ -112,7 +112,8 @@ def validate_db(engine: Engine, db_config: DatabaseConfig) -> None:
             f"Validating database connection to instance {db_config.instance_name}"
         )
         try:
-            ws = WorkspaceClient()
+            from genie_space_optimizer._workspace_client import make_workspace_client
+            ws = make_workspace_client()
             ws.database.get_database_instance(db_config.instance_name)
         except NotFound:
             raise ValueError(

packages/genie-space-optimizer/src/genie_space_optimizer/jobs/run_baseline.py

@@ -200,7 +200,8 @@
 
 # COMMAND ----------
 
-w = WorkspaceClient()
+from genie_space_optimizer._workspace_client import make_workspace_client
+w = make_workspace_client()
 spark = SparkSession.builder.getOrCreate()
 
 # Read task values from preflight

packages/genie-space-optimizer/src/genie_space_optimizer/jobs/run_cross_env_deploy.py

@@ -152,8 +152,9 @@
 
 _banner("Connecting to Target Workspace")
 
+from genie_space_optimizer._workspace_client import make_workspace_client
 target_url = target_workspace_url.rstrip("/")
-target_ws = WorkspaceClient(host=target_url)
+target_ws = make_workspace_client(host=target_url)
 
 current_user = target_ws.current_user.me()
 _log(

packages/genie-space-optimizer/src/genie_space_optimizer/jobs/run_deploy.py

@@ -125,7 +125,8 @@
 
 # COMMAND ----------
 
-w = WorkspaceClient()
+from genie_space_optimizer._workspace_client import make_workspace_client
+w = make_workspace_client()
 spark = SparkSession.builder.getOrCreate()
 
 # Read task values from upstream

packages/genie-space-optimizer/src/genie_space_optimizer/jobs/run_enrichment.py

@@ -66,7 +66,8 @@
 
 # COMMAND ----------
 
-w = WorkspaceClient()
+from genie_space_optimizer._workspace_client import make_workspace_client
+w = make_workspace_client()
 spark = SparkSession.builder.getOrCreate()
 
 run_id = dbutils.jobs.taskValues.get(taskKey="preflight", key="run_id")