classifyre-cli 0.4.11__tar.gz → 0.4.12__tar.gz

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/.turbo/turbo-build.log

@@ -1,3 +1,3 @@
 $ uv sync
-Resolved 265 packages in 157ms
+Resolved 267 packages in 181ms
 Checked 50 packages in 1ms

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: classifyre-cli
-Version: 0.4.11
+Version: 0.4.12
 Summary: Classifyre CLI — scan and classify unstructured data sources
 License: MIT
 Keywords: data,ingestion,metadata,pii,secrets,unstructured

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@classifyre/cli",
-  "version": "0.4.11",
+  "version": "0.4.12",
   "private": true,
   "scripts": {
     "build": "uv sync",

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "classifyre-cli"
-version = "0.4.11"
+version = "0.4.12"
 description = "Classifyre CLI — scan and classify unstructured data sources"
 readme = "README.md"
 requires-python = ">=3.12"

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/src/models/generated_input.py

@@ -1078,9 +1078,10 @@ class DatabricksAuthMode(StrEnum):
 
     PAT_TOKEN = 'PAT_TOKEN'
     SERVICE_PRINCIPAL = 'SERVICE_PRINCIPAL'
+    AZURE_SERVICE_PRINCIPAL = 'AZURE_SERVICE_PRINCIPAL'
 
 
-class DatabricksRequiredPat(BaseModel):
+class PersonalAccessToken(BaseModel):
     model_config = ConfigDict(
         extra='forbid',
     )
@@ -1094,7 +1095,7 @@ class DatabricksRequiredPat(BaseModel):
     )
 
 
-class DatabricksRequiredServicePrincipal(BaseModel):
+class ServicePrincipalOAuthM2M(BaseModel):
     model_config = ConfigDict(
         extra='forbid',
     )
@@ -1109,6 +1110,24 @@ class DatabricksRequiredServicePrincipal(BaseModel):
     client_id: str = Field(..., description='Databricks service principal client ID')
 
 
+class AzureServicePrincipal(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    auth_mode: Literal['AZURE_SERVICE_PRINCIPAL']
+    workspace_url: AnyUrl = Field(
+        ...,
+        description='Azure Databricks workspace URL (for example, https://adb-1234567890123456.7.azuredatabricks.net)',
+    )
+    warehouse_id: str = Field(
+        ..., description='Databricks SQL warehouse ID used for sampling queries'
+    )
+    client_id: str = Field(
+        ..., description='Azure AD application (client) ID for the service principal'
+    )
+    tenant_id: str = Field(..., description='Azure AD tenant ID')
+
+
 class DatabricksMaskedPat(BaseModel):
     model_config = ConfigDict(
         extra='forbid',
@@ -1125,6 +1144,15 @@ class DatabricksMaskedServicePrincipal(BaseModel):
     )
 
 
+class DatabricksMaskedAzureServicePrincipal(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    client_secret: str = Field(
+        ..., description='Azure AD client secret for the service principal'
+    )
+
+
 class DatabricksOptionalConnection(BaseModel):
     """
     Databricks API and SQL statement execution tuning options.
@@ -2020,8 +2048,8 @@ class DatabricksInput(CoreInput):
     type: Literal['DATABRICKS'] = Field(
         'DATABRICKS', description='Type of the asset or source'
     )
-    required: DatabricksRequiredPat | DatabricksRequiredServicePrincipal = Field(
-        ..., title='DatabricksRequired'
+    required: PersonalAccessToken | ServicePrincipalOAuthM2M | AzureServicePrincipal = (
+        Field(..., title='DatabricksRequired')
     )
     masked: DatabricksMaskedPat | DatabricksMaskedServicePrincipal = Field(
         ..., title='DatabricksMasked'

{classifyre_cli-0.4.11 → classifyre_cli-0.4.12}/src/sources/databricks/source.py

@@ -13,16 +13,18 @@ from urllib.parse import urlparse
 import requests
 
 from ...models.generated_input import (
+    AzureServicePrincipal,
     DatabricksInput,
+    DatabricksMaskedAzureServicePrincipal,
     DatabricksMaskedPat,
     DatabricksMaskedServicePrincipal,
     DatabricksOptionalConnection,
     DatabricksOptionalExtraction,
     DatabricksOptionalScope,
-    DatabricksRequiredPat,
-    DatabricksRequiredServicePrincipal,
+    PersonalAccessToken,
     SamplingConfig,
     SamplingStrategy,
+    ServicePrincipalOAuthM2M,
 )
 from ...models.generated_single_asset_scan_results import (
     AssetType as OutputAssetType,
@@ -106,14 +108,20 @@ class DatabricksSource(BaseTabularSource):
         required = self.config.required
         masked = self.config.masked
 
-        if isinstance(required, DatabricksRequiredPat):
+        if isinstance(required, PersonalAccessToken):
             if not isinstance(masked, DatabricksMaskedPat):
                 raise ValueError("DATABRICKS PAT_TOKEN auth requires masked.token")
             return
-        if isinstance(required, DatabricksRequiredServicePrincipal):
+        if isinstance(required, ServicePrincipalOAuthM2M):
             if not isinstance(masked, DatabricksMaskedServicePrincipal):
                 raise ValueError("DATABRICKS SERVICE_PRINCIPAL auth requires masked.client_secret")
             return
+        if isinstance(required, AzureServicePrincipal):
+            if not isinstance(masked, DatabricksMaskedAzureServicePrincipal):
+                raise ValueError(
+                    "DATABRICKS AZURE_SERVICE_PRINCIPAL auth requires masked.client_secret"
+                )
+            return
         raise ValueError("Unsupported DATABRICKS auth configuration")
 
     # ── Identity ─────────────────────────────────────────────────────────
@@ -163,7 +171,10 @@ class DatabricksSource(BaseTabularSource):
         return int(self._connection_options().statement_timeout_seconds or 60)
 
     def _is_pat_mode(self) -> bool:
-        return isinstance(self.config.required, DatabricksRequiredPat)
+        return isinstance(self.config.required, PersonalAccessToken)
+
+    def _is_azure_sp_mode(self) -> bool:
+        return isinstance(self.config.required, AzureServicePrincipal)
 
     def _masked_pat_token(self) -> str:
         masked = self.config.masked
@@ -174,12 +185,23 @@ class DatabricksSource(BaseTabularSource):
     def _service_principal_credentials(self) -> tuple[str, str]:
         required = self.config.required
         masked = self.config.masked
-        if not isinstance(required, DatabricksRequiredServicePrincipal):
+        if not isinstance(required, ServicePrincipalOAuthM2M):
             raise ValueError("SERVICE_PRINCIPAL auth mode is required")
         if not isinstance(masked, DatabricksMaskedServicePrincipal):
             raise ValueError("DATABRICKS SERVICE_PRINCIPAL auth requires masked.client_secret")
         return required.client_id, masked.client_secret
 
+    def _azure_sp_credentials(self) -> tuple[str, str, str]:
+        required = self.config.required
+        masked = self.config.masked
+        if not isinstance(required, AzureServicePrincipal):
+            raise ValueError("AZURE_SERVICE_PRINCIPAL auth mode is required")
+        if not isinstance(masked, DatabricksMaskedAzureServicePrincipal):
+            raise ValueError(
+                "DATABRICKS AZURE_SERVICE_PRINCIPAL auth requires masked.client_secret"
+            )
+        return required.tenant_id, required.client_id, masked.client_secret
+
     def _is_access_token_expired(self) -> bool:
         if self._access_token_expiry is None:
             return True
@@ -206,12 +228,38 @@ class DatabricksSource(BaseTabularSource):
         self._access_token_expiry = datetime.now(UTC) + timedelta(seconds=max(expires_in - 300, 0))
         return token.strip()
 
+    def _acquire_azure_token(self) -> str:
+        # Azure AD v1 token endpoint; resource ID is the fixed Databricks app in Azure
+        _databricks_azure_resource = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+        tenant_id, client_id, client_secret = self._azure_sp_credentials()
+        response = self.session.post(
+            f"https://login.microsoftonline.com/{tenant_id}/oauth2/token",
+            data={
+                "grant_type": "client_credentials",
+                "client_id": client_id,
+                "client_secret": client_secret,
+                "resource": _databricks_azure_resource,
+            },
+            timeout=self._timeout_seconds(),
+        )
+        response.raise_for_status()
+        payload = response.json()
+        token = payload.get("access_token")
+        if not isinstance(token, str) or not token.strip():
+            raise ValueError("Azure AD token response did not include access_token")
+        expires_in = int(payload.get("expires_in", 3600))
+        self._access_token_expiry = datetime.now(UTC) + timedelta(seconds=max(expires_in - 300, 0))
+        return token.strip()
+
     def _access_token_value(self) -> str:
         if self._is_pat_mode():
             return self._masked_pat_token().strip()
         if self._access_token and not self._is_access_token_expired():
             return self._access_token
-        self._access_token = self._acquire_service_principal_token()
+        if self._is_azure_sp_mode():
+            self._access_token = self._acquire_azure_token()
+        else:
+            self._access_token = self._acquire_service_principal_token()
         return self._access_token
 
     def _authorization_header(self) -> str:
@@ -860,7 +908,12 @@ class DatabricksSource(BaseTabularSource):
                 with conn.cursor() as cursor:
                     cursor.execute("SELECT 1")
                     cursor.fetchone()
-            auth_mode = "PAT_TOKEN" if self._is_pat_mode() else "SERVICE_PRINCIPAL"
+            if self._is_pat_mode():
+                auth_mode = "PAT_TOKEN"
+            elif self._is_azure_sp_mode():
+                auth_mode = "AZURE_SERVICE_PRINCIPAL"
+            else:
+                auth_mode = "SERVICE_PRINCIPAL"
             result["status"] = "SUCCESS"
             result["message"] = (
                 "Successfully connected to Databricks Unity Catalog "