databricks-sdk 0.0.7__py3-none-any.whl → 0.1.1__py3-none-any.whl

databricks/sdk/__init__.py

@@ -1,33 +1,56 @@
 import databricks.sdk.core as client
 import databricks.sdk.dbutils as dbutils
-import databricks.sdk.mixins.compute as compute_mixin
-import databricks.sdk.mixins.dbfs as dbfs_mixin
-import databricks.sdk.service.billing as billing
-import databricks.sdk.service.clusterpolicies as clusterpolicies
-import databricks.sdk.service.clusters as clusters
-import databricks.sdk.service.commands as commands
-import databricks.sdk.service.dbfs as dbfs
-import databricks.sdk.service.deployment as deployment
-import databricks.sdk.service.endpoints as endpoints
-import databricks.sdk.service.gitcredentials as gitcredentials
-import databricks.sdk.service.globalinitscripts as globalinitscripts
-import databricks.sdk.service.instancepools as instancepools
-import databricks.sdk.service.ipaccesslists as ipaccesslists
-import databricks.sdk.service.jobs as jobs
-import databricks.sdk.service.libraries as libraries
-import databricks.sdk.service.mlflow as mlflow
-import databricks.sdk.service.oauth2 as oauth2
-import databricks.sdk.service.permissions as permissions
-import databricks.sdk.service.pipelines as pipelines
-import databricks.sdk.service.repos as repos
-import databricks.sdk.service.scim as scim
-import databricks.sdk.service.secrets as secrets
-import databricks.sdk.service.sql as sql
-import databricks.sdk.service.tokenmanagement as tokenmanagement
-import databricks.sdk.service.tokens as tokens
-import databricks.sdk.service.unitycatalog as unitycatalog
-import databricks.sdk.service.workspace as workspace
-import databricks.sdk.service.workspaceconf as workspaceconf
+from databricks.sdk.mixins.compute import ClustersExt
+from databricks.sdk.mixins.dbfs import DbfsExt
+from databricks.sdk.service.billing import (BillableUsageAPI, BudgetsAPI,
+                                            LogDeliveryAPI)
+from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI,
+                                            AccountMetastoresAPI,
+                                            AccountStorageCredentialsAPI,
+                                            CatalogsAPI, ExternalLocationsAPI,
+                                            FunctionsAPI, GrantsAPI,
+                                            MetastoresAPI, SchemasAPI,
+                                            StorageCredentialsAPI,
+                                            TableConstraintsAPI, TablesAPI,
+                                            VolumesAPI)
+from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI,
+                                            CommandExecutionAPI,
+                                            GlobalInitScriptsAPI,
+                                            InstancePoolsAPI,
+                                            InstanceProfilesAPI, LibrariesAPI,
+                                            PolicyFamiliesAPI)
+from databricks.sdk.service.files import DbfsAPI
+from databricks.sdk.service.iam import (AccountGroupsAPI,
+                                        AccountServicePrincipalsAPI,
+                                        AccountUsersAPI, CurrentUserAPI,
+                                        GroupsAPI, PermissionsAPI,
+                                        ServicePrincipalsAPI, UsersAPI,
+                                        WorkspaceAssignmentAPI)
+from databricks.sdk.service.jobs import JobsAPI
+from databricks.sdk.service.ml import ExperimentsAPI, ModelRegistryAPI
+from databricks.sdk.service.oauth2 import (CustomAppIntegrationAPI,
+                                           OAuthEnrollmentAPI,
+                                           PublishedAppIntegrationAPI)
+from databricks.sdk.service.pipelines import PipelinesAPI
+from databricks.sdk.service.provisioning import (CredentialsAPI,
+                                                 EncryptionKeysAPI,
+                                                 NetworksAPI, PrivateAccessAPI,
+                                                 StorageAPI, VpcEndpointsAPI,
+                                                 WorkspacesAPI)
+from databricks.sdk.service.serving import ServingEndpointsAPI
+from databricks.sdk.service.settings import (AccountIpAccessListsAPI,
+                                             IpAccessListsAPI,
+                                             TokenManagementAPI, TokensAPI,
+                                             WorkspaceConfAPI)
+from databricks.sdk.service.sharing import (ProvidersAPI,
+                                            RecipientActivationAPI,
+                                            RecipientsAPI, SharesAPI)
+from databricks.sdk.service.sql import (AlertsAPI, DashboardsAPI,
+                                        DataSourcesAPI, DbsqlPermissionsAPI,
+                                        QueriesAPI, QueryHistoryAPI,
+                                        StatementExecutionAPI, WarehousesAPI)
+from databricks.sdk.service.workspace import (GitCredentialsAPI, ReposAPI,
+                                              SecretsAPI, WorkspaceAPI)
 
 
 class WorkspaceClient:
@@ -81,62 +104,55 @@ class WorkspaceClient:
         self.config = config
         self.dbutils = dbutils.RemoteDbUtils(self.config)
         self.api_client = client.ApiClient(self.config)
-        self.alerts = sql.AlertsAPI(self.api_client)
-        self.catalogs = unitycatalog.CatalogsAPI(self.api_client)
-        self.cluster_policies = clusterpolicies.ClusterPoliciesAPI(self.api_client)
-        self.clusters = compute_mixin.ClustersExt(self.api_client)
-        self.command_execution = commands.CommandExecutionAPI(self.api_client)
-        self.current_user = scim.CurrentUserAPI(self.api_client)
-        self.dashboards = sql.DashboardsAPI(self.api_client)
-        self.data_sources = sql.DataSourcesAPI(self.api_client)
-        self.dbfs = dbfs_mixin.DbfsExt(self.api_client)
-        self.dbsql_permissions = sql.DbsqlPermissionsAPI(self.api_client)
-        self.experiments = mlflow.ExperimentsAPI(self.api_client)
-        self.external_locations = unitycatalog.ExternalLocationsAPI(self.api_client)
-        self.functions = unitycatalog.FunctionsAPI(self.api_client)
-        self.git_credentials = gitcredentials.GitCredentialsAPI(self.api_client)
-        self.global_init_scripts = globalinitscripts.GlobalInitScriptsAPI(self.api_client)
-        self.grants = unitycatalog.GrantsAPI(self.api_client)
-        self.groups = scim.GroupsAPI(self.api_client)
-        self.instance_pools = instancepools.InstancePoolsAPI(self.api_client)
-        self.instance_profiles = clusters.InstanceProfilesAPI(self.api_client)
-        self.ip_access_lists = ipaccesslists.IpAccessListsAPI(self.api_client)
-        self.jobs = jobs.JobsAPI(self.api_client)
-        self.libraries = libraries.LibrariesAPI(self.api_client)
-        self.m_lflow_artifacts = mlflow.MLflowArtifactsAPI(self.api_client)
-        self.m_lflow_databricks = mlflow.MLflowDatabricksAPI(self.api_client)
-        self.m_lflow_metrics = mlflow.MLflowMetricsAPI(self.api_client)
-        self.m_lflow_runs = mlflow.MLflowRunsAPI(self.api_client)
-        self.metastores = unitycatalog.MetastoresAPI(self.api_client)
-        self.model_version_comments = mlflow.ModelVersionCommentsAPI(self.api_client)
-        self.model_versions = mlflow.ModelVersionsAPI(self.api_client)
-        self.permissions = permissions.PermissionsAPI(self.api_client)
-        self.pipelines = pipelines.PipelinesAPI(self.api_client)
-        self.policy_families = clusterpolicies.PolicyFamiliesAPI(self.api_client)
-        self.providers = unitycatalog.ProvidersAPI(self.api_client)
-        self.queries = sql.QueriesAPI(self.api_client)
-        self.query_history = sql.QueryHistoryAPI(self.api_client)
-        self.recipient_activation = unitycatalog.RecipientActivationAPI(self.api_client)
-        self.recipients = unitycatalog.RecipientsAPI(self.api_client)
-        self.registered_models = mlflow.RegisteredModelsAPI(self.api_client)
-        self.registry_webhooks = mlflow.RegistryWebhooksAPI(self.api_client)
-        self.repos = repos.ReposAPI(self.api_client)
-        self.schemas = unitycatalog.SchemasAPI(self.api_client)
-        self.secrets = secrets.SecretsAPI(self.api_client)
-        self.service_principals = scim.ServicePrincipalsAPI(self.api_client)
-        self.serving_endpoints = endpoints.ServingEndpointsAPI(self.api_client)
-        self.shares = unitycatalog.SharesAPI(self.api_client)
-        self.statement_execution = sql.StatementExecutionAPI(self.api_client)
-        self.storage_credentials = unitycatalog.StorageCredentialsAPI(self.api_client)
-        self.table_constraints = unitycatalog.TableConstraintsAPI(self.api_client)
-        self.tables = unitycatalog.TablesAPI(self.api_client)
-        self.token_management = tokenmanagement.TokenManagementAPI(self.api_client)
-        self.tokens = tokens.TokensAPI(self.api_client)
-        self.transition_requests = mlflow.TransitionRequestsAPI(self.api_client)
-        self.users = scim.UsersAPI(self.api_client)
-        self.warehouses = sql.WarehousesAPI(self.api_client)
-        self.workspace = workspace.WorkspaceAPI(self.api_client)
-        self.workspace_conf = workspaceconf.WorkspaceConfAPI(self.api_client)
+        self.alerts = AlertsAPI(self.api_client)
+        self.catalogs = CatalogsAPI(self.api_client)
+        self.cluster_policies = ClusterPoliciesAPI(self.api_client)
+        self.clusters = ClustersExt(self.api_client)
+        self.command_execution = CommandExecutionAPI(self.api_client)
+        self.current_user = CurrentUserAPI(self.api_client)
+        self.dashboards = DashboardsAPI(self.api_client)
+        self.data_sources = DataSourcesAPI(self.api_client)
+        self.dbfs = DbfsExt(self.api_client)
+        self.dbsql_permissions = DbsqlPermissionsAPI(self.api_client)
+        self.experiments = ExperimentsAPI(self.api_client)
+        self.external_locations = ExternalLocationsAPI(self.api_client)
+        self.functions = FunctionsAPI(self.api_client)
+        self.git_credentials = GitCredentialsAPI(self.api_client)
+        self.global_init_scripts = GlobalInitScriptsAPI(self.api_client)
+        self.grants = GrantsAPI(self.api_client)
+        self.groups = GroupsAPI(self.api_client)
+        self.instance_pools = InstancePoolsAPI(self.api_client)
+        self.instance_profiles = InstanceProfilesAPI(self.api_client)
+        self.ip_access_lists = IpAccessListsAPI(self.api_client)
+        self.jobs = JobsAPI(self.api_client)
+        self.libraries = LibrariesAPI(self.api_client)
+        self.metastores = MetastoresAPI(self.api_client)
+        self.model_registry = ModelRegistryAPI(self.api_client)
+        self.permissions = PermissionsAPI(self.api_client)
+        self.pipelines = PipelinesAPI(self.api_client)
+        self.policy_families = PolicyFamiliesAPI(self.api_client)
+        self.providers = ProvidersAPI(self.api_client)
+        self.queries = QueriesAPI(self.api_client)
+        self.query_history = QueryHistoryAPI(self.api_client)
+        self.recipient_activation = RecipientActivationAPI(self.api_client)
+        self.recipients = RecipientsAPI(self.api_client)
+        self.repos = ReposAPI(self.api_client)
+        self.schemas = SchemasAPI(self.api_client)
+        self.secrets = SecretsAPI(self.api_client)
+        self.service_principals = ServicePrincipalsAPI(self.api_client)
+        self.serving_endpoints = ServingEndpointsAPI(self.api_client)
+        self.shares = SharesAPI(self.api_client)
+        self.statement_execution = StatementExecutionAPI(self.api_client)
+        self.storage_credentials = StorageCredentialsAPI(self.api_client)
+        self.table_constraints = TableConstraintsAPI(self.api_client)
+        self.tables = TablesAPI(self.api_client)
+        self.token_management = TokenManagementAPI(self.api_client)
+        self.tokens = TokensAPI(self.api_client)
+        self.users = UsersAPI(self.api_client)
+        self.volumes = VolumesAPI(self.api_client)
+        self.warehouses = WarehousesAPI(self.api_client)
+        self.workspace = WorkspaceAPI(self.api_client)
+        self.workspace_conf = WorkspaceConfAPI(self.api_client)
 
 
 class AccountClient:
@@ -189,23 +205,24 @@ class AccountClient:
                                    product_version=product_version)
         self.config = config
         self.api_client = client.ApiClient(self.config)
-        self.billable_usage = billing.BillableUsageAPI(self.api_client)
-        self.budgets = billing.BudgetsAPI(self.api_client)
-        self.credentials = deployment.CredentialsAPI(self.api_client)
-        self.custom_app_integration = oauth2.CustomAppIntegrationAPI(self.api_client)
-        self.encryption_keys = deployment.EncryptionKeysAPI(self.api_client)
-        self.account_groups = scim.AccountGroupsAPI(self.api_client)
-        self.log_delivery = billing.LogDeliveryAPI(self.api_client)
-        self.account_metastore_assignments = unitycatalog.AccountMetastoreAssignmentsAPI(self.api_client)
-        self.account_metastores = unitycatalog.AccountMetastoresAPI(self.api_client)
-        self.networks = deployment.NetworksAPI(self.api_client)
-        self.o_auth_enrollment = oauth2.OAuthEnrollmentAPI(self.api_client)
-        self.private_access = deployment.PrivateAccessAPI(self.api_client)
-        self.published_app_integration = oauth2.PublishedAppIntegrationAPI(self.api_client)
-        self.account_service_principals = scim.AccountServicePrincipalsAPI(self.api_client)
-        self.storage = deployment.StorageAPI(self.api_client)
-        self.account_storage_credentials = unitycatalog.AccountStorageCredentialsAPI(self.api_client)
-        self.account_users = scim.AccountUsersAPI(self.api_client)
-        self.vpc_endpoints = deployment.VpcEndpointsAPI(self.api_client)
-        self.workspace_assignment = permissions.WorkspaceAssignmentAPI(self.api_client)
-        self.workspaces = deployment.WorkspacesAPI(self.api_client)
+        self.billable_usage = BillableUsageAPI(self.api_client)
+        self.budgets = BudgetsAPI(self.api_client)
+        self.credentials = CredentialsAPI(self.api_client)
+        self.custom_app_integration = CustomAppIntegrationAPI(self.api_client)
+        self.encryption_keys = EncryptionKeysAPI(self.api_client)
+        self.groups = AccountGroupsAPI(self.api_client)
+        self.ip_access_lists = AccountIpAccessListsAPI(self.api_client)
+        self.log_delivery = LogDeliveryAPI(self.api_client)
+        self.metastore_assignments = AccountMetastoreAssignmentsAPI(self.api_client)
+        self.metastores = AccountMetastoresAPI(self.api_client)
+        self.networks = NetworksAPI(self.api_client)
+        self.o_auth_enrollment = OAuthEnrollmentAPI(self.api_client)
+        self.private_access = PrivateAccessAPI(self.api_client)
+        self.published_app_integration = PublishedAppIntegrationAPI(self.api_client)
+        self.service_principals = AccountServicePrincipalsAPI(self.api_client)
+        self.storage = StorageAPI(self.api_client)
+        self.storage_credentials = AccountStorageCredentialsAPI(self.api_client)
+        self.users = AccountUsersAPI(self.api_client)
+        self.vpc_endpoints = VpcEndpointsAPI(self.api_client)
+        self.workspace_assignment = WorkspaceAssignmentAPI(self.api_client)
+        self.workspaces = WorkspacesAPI(self.api_client)

databricks/sdk/core.py

@@ -114,7 +114,9 @@ def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
 def external_browser(cfg: 'Config') -> Optional[HeaderFactory]:
     if cfg.auth_type != 'external-browser':
         return None
-    if cfg.is_aws:
+    if cfg.client_id:
+        client_id = cfg.client_id
+    elif cfg.is_aws:
         client_id = 'databricks-cli'
     elif cfg.is_azure:
         # Use Azure AD app for cases when Azure CLI is not available on the machine.
@@ -123,7 +125,10 @@ def external_browser(cfg: 'Config') -> Optional[HeaderFactory]:
         client_id = '6128a518-99a9-425b-8333-4cc94f04cacd'
     else:
         raise ValueError(f'local browser SSO is not supported')
-    oauth_client = OAuthClient(cfg.host, client_id, 'http://localhost:8020')
+    oauth_client = OAuthClient(host=cfg.host,
+                               client_id=client_id,
+                               redirect_url='http://localhost:8020',
+                               client_secret=cfg.client_secret)
     consent = oauth_client.initiate_consent()
     if not consent:
         return None
@@ -177,16 +182,18 @@ def azure_service_principal(cfg: 'Config') -> HeaderFactory:
     return refreshed_headers
 
 
-class AzureCliTokenSource(Refreshable):
-    """ Obtain the token granted by `az login` CLI command """
+class CliTokenSource(Refreshable):
 
-    def __init__(self, resource: str):
+    def __init__(self, cmd: List[str], token_type_field: str, access_token_field: str, expiry_field: str):
         super().__init__()
-        self.resource = resource
+        self._cmd = cmd
+        self._token_type_field = token_type_field
+        self._access_token_field = access_token_field
+        self._expiry_field = expiry_field
 
     @staticmethod
     def _parse_expiry(expiry: str) -> datetime:
-        for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S"):
+        for fmt in ("%Y-%m-%d %H:%M:%S.%f", "%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f%z"):
             try:
                 return datetime.strptime(expiry, fmt)
             except ValueError as e:
@@ -196,18 +203,28 @@ class AzureCliTokenSource(Refreshable):
 
     def refresh(self) -> Token:
         try:
-            cmd = ["az", "account", "get-access-token", "--resource", self.resource, "--output", "json"]
-            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
+            out = subprocess.check_output(self._cmd, stderr=subprocess.STDOUT)
             it = json.loads(out.decode())
-            expires_on = self._parse_expiry(it["expiresOn"])
-            return Token(access_token=it["accessToken"],
-                         refresh_token=it.get('refreshToken', None),
-                         token_type=it["tokenType"],
+            expires_on = self._parse_expiry(it[self._expiry_field])
+            return Token(access_token=it[self._access_token_field],
+                         token_type=it[self._token_type_field],
                          expiry=expires_on)
         except ValueError as e:
             raise ValueError(f"cannot unmarshal CLI result: {e}")
         except subprocess.CalledProcessError as e:
-            raise IOError(f'cannot get access token: {e.output.decode()}') from e
+            message = e.output.decode().strip()
+            raise IOError(f'cannot get access token: {message}') from e
+
+
+class AzureCliTokenSource(CliTokenSource):
+    """ Obtain the token granted by `az login` CLI command """
+
+    def __init__(self, resource: str):
+        cmd = ["az", "account", "get-access-token", "--resource", resource, "--output", "json"]
+        super().__init__(cmd=cmd,
+                         token_type_field='tokenType',
+                         access_token_field='accessToken',
+                         expiry_field='expiresOn')
 
 
 @credentials_provider('azure-cli', ['is_azure'])
@@ -231,6 +248,45 @@ def azure_cli(cfg: 'Config') -> Optional[HeaderFactory]:
     return inner
 
 
+class BricksCliTokenSource(CliTokenSource):
+    """ Obtain the token granted by `bricks auth login` CLI command """
+
+    def __init__(self, cfg: 'Config'):
+        cli_path = cfg.bricks_cli_path
+        if not cli_path:
+            cli_path = 'bricks'
+        cmd = [cli_path, 'auth', 'token', '--host', cfg.host]
+        if cfg.is_account_client:
+            cmd += ['--account-id', cfg.account_id]
+        super().__init__(cmd=cmd,
+                         token_type_field='token_type',
+                         access_token_field='access_token',
+                         expiry_field='expiry')
+
+
+@credentials_provider('bricks-cli', ['host', 'is_aws'])
+def bricks_cli(cfg: 'Config') -> Optional[HeaderFactory]:
+    token_source = BricksCliTokenSource(cfg)
+    try:
+        token_source.token()
+    except FileNotFoundError:
+        logger.debug(f'Most likely Bricks CLI is not installed.')
+        return None
+    except IOError as e:
+        if 'databricks OAuth is not' in str(e):
+            logger.debug(f'OAuth not configured or not available: {e}')
+            return None
+        raise e
+
+    logger.info("Using Bricks CLI authentication")
+
+    def inner() -> Dict[str, str]:
+        token = token_source.token()
+        return {'Authorization': f'{token.token_type} {token.access_token}'}
+
+    return inner
+
+
 class DefaultCredentials:
     """ Select the first applicable credential provider from the chain """
 
@@ -243,7 +299,7 @@ class DefaultCredentials:
     def __call__(self, cfg: 'Config') -> HeaderFactory:
         auth_providers = [
             pat_auth, basic_auth, oauth_service_principal, azure_service_principal, azure_cli,
-            external_browser
+            external_browser, bricks_cli
         ]
         for provider in auth_providers:
             auth_type = provider.auth_type()
@@ -372,7 +428,11 @@ class Config:
     def is_azure(self) -> bool:
         has_resource_id = self.azure_workspace_resource_id is not None
         has_host = self.host is not None
-        return has_resource_id or (has_host and ".azuredatabricks.net" in self.host)
+        is_public_cloud = has_host and ".azuredatabricks.net" in self.host
+        is_china_cloud = has_host and ".databricks.azure.cn" in self.host
+        is_gov_cloud = has_host and ".databricks.azure.us" in self.host
+        is_valid_cloud = is_public_cloud or is_china_cloud or is_gov_cloud
+        return has_resource_id or (has_host and is_valid_cloud)
 
     @property
     def is_gcp(self) -> bool:

databricks/sdk/dbutils.py

@@ -6,8 +6,9 @@ import typing
 from collections import namedtuple
 
 from .core import ApiClient, Config
-from .mixins import compute, dbfs
-from .service import commands, secrets
+from .mixins import compute as compute_ext
+from .mixins import dbfs as dbfs_ext
+from .service import compute, workspace
 
 
 class FileInfo(namedtuple('FileInfo', ['path', 'name', 'size', "modificationTime"])):
@@ -31,7 +32,7 @@ class SecretMetadata(namedtuple('SecretMetadata', ['key'])):
 class _FsUtil:
     """ Manipulates the Databricks filesystem (DBFS) """
 
-    def __init__(self, dbfs_ext: dbfs.DbfsExt, proxy_factory: typing.Callable[[str], '_ProxyUtil']):
+    def __init__(self, dbfs_ext: dbfs_ext.DbfsExt, proxy_factory: typing.Callable[[str], '_ProxyUtil']):
         self._dbfs = dbfs_ext
         self._proxy_factory = proxy_factory
 
@@ -177,7 +178,7 @@ _FILTER = _RedactingFilter()
 class _SecretsUtil:
     """Remote equivalent of secrets util"""
 
-    def __init__(self, secrets_api: secrets.SecretsAPI):
+    def __init__(self, secrets_api: workspace.SecretsAPI):
         self._api = secrets_api # nolint
 
     def getBytes(self, scope: str, key: str) -> bytes:
@@ -215,13 +216,13 @@ class RemoteDbUtils:
     def __init__(self, config: 'Config' = None):
         self._config = Config() if not config else config
         self._client = ApiClient(self._config)
-        self._clusters = compute.ClustersExt(self._client)
-        self._commands = commands.CommandExecutionAPI(self._client)
+        self._clusters = compute_ext.ClustersExt(self._client)
+        self._commands = compute.CommandExecutionAPI(self._client)
         self._lock = threading.Lock()
         self._ctx = None
 
-        self.fs = _FsUtil(dbfs.DbfsExt(self._client), self.__getattr__)
-        self.secrets = _SecretsUtil(secrets.SecretsAPI(self._client))
+        self.fs = _FsUtil(dbfs_ext.DbfsExt(self._client), self.__getattr__)
+        self.secrets = _SecretsUtil(workspace.SecretsAPI(self._client))
 
     @property
     def _cluster_id(self) -> str:
@@ -231,7 +232,7 @@ class RemoteDbUtils:
             raise ValueError(self._config.wrap_debug_info(message))
         return cluster_id
 
-    def _running_command_context(self) -> commands.ContextStatusResponse:
+    def _running_command_context(self) -> compute.ContextStatusResponse:
         if self._ctx:
             return self._ctx
         with self._lock:
@@ -239,7 +240,7 @@ class RemoteDbUtils:
                 return self._ctx
             self._clusters.ensure_cluster_is_running(self._cluster_id)
             self._ctx = self._commands.create(cluster_id=self._cluster_id,
-                                              language=commands.Language.python).result()
+                                              language=compute.Language.python).result()
         return self._ctx
 
     def __getattr__(self, util) -> '_ProxyUtil':
@@ -252,9 +253,9 @@ class RemoteDbUtils:
 class _ProxyUtil:
     """Enables temporary workaround to call remote in-REPL dbutils without having to re-implement them"""
 
-    def __init__(self, *, command_execution: commands.CommandExecutionAPI,
-                 context_factory: typing.Callable[[], commands.ContextStatusResponse], cluster_id: str,
-                 name: str):
+    def __init__(self, *, command_execution: compute.CommandExecutionAPI,
+                 context_factory: typing.Callable[[],
+                                                  compute.ContextStatusResponse], cluster_id: str, name: str):
         self._commands = command_execution
         self._cluster_id = cluster_id
         self._context_factory = context_factory
@@ -270,8 +271,8 @@ class _ProxyUtil:
 
 class _ProxyCall:
 
-    def __init__(self, *, command_execution: commands.CommandExecutionAPI,
-                 context_factory: typing.Callable[[], commands.ContextStatusResponse], cluster_id: str,
+    def __init__(self, *, command_execution: compute.CommandExecutionAPI,
+                 context_factory: typing.Callable[[], compute.ContextStatusResponse], cluster_id: str,
                  util: str, method: str):
         self._commands = command_execution
         self._cluster_id = cluster_id
@@ -289,10 +290,10 @@ class _ProxyCall:
         '''
         ctx = self._context_factory()
         result = self._commands.execute(cluster_id=self._cluster_id,
-                                        language=commands.Language.python,
+                                        language=compute.Language.python,
                                         context_id=ctx.id,
                                         command=code).result()
-        if result.status == commands.CommandStatus.Finished:
+        if result.status == compute.CommandStatus.Finished:
             raw = result.results.data
             return json.loads(raw)
         else:

databricks/sdk/mixins/compute.py

@@ -2,7 +2,7 @@ import re
 from dataclasses import dataclass
 from typing import Optional
 
-from databricks.sdk.service import clusters
+from databricks.sdk.service import compute
 
 
 @dataclass
@@ -54,7 +54,7 @@ class SemVer:
         return self.build < other.build
 
 
-class ClustersExt(clusters.ClustersAPI):
+class ClustersExt(compute.ClustersAPI):
 
     def select_spark_version(self,
                              long_term_support: bool = False,
@@ -92,7 +92,7 @@ class ClustersExt(clusters.ClustersAPI):
         return versions[0]
 
     @staticmethod
-    def _node_sorting_tuple(item: clusters.NodeType) -> tuple:
+    def _node_sorting_tuple(item: compute.NodeType) -> tuple:
         local_disks = local_disk_size_gb = local_nvme_disk = local_nvme_disk_size_gb = 0
         if item.node_instance_type is not None:
             local_disks = item.node_instance_type.local_disks
@@ -103,12 +103,12 @@ class ClustersExt(clusters.ClustersAPI):
                 local_nvme_disk, local_nvme_disk_size_gb, item.num_gpus, item.instance_type_id)
 
     @staticmethod
-    def _should_node_be_skipped(nt: clusters.NodeType) -> bool:
+    def _should_node_be_skipped(nt: compute.NodeType) -> bool:
         if not nt.node_info:
             return False
         if not nt.node_info.status:
             return False
-        val = clusters.CloudProviderNodeStatus
+        val = compute.CloudProviderNodeStatus
         for st in nt.node_info.status:
             if st in (val.NotAvailableInRegion, val.NotEnabledOnSubscription):
                 return True
@@ -173,7 +173,7 @@ class ClustersExt(clusters.ClustersAPI):
         raise ValueError("cannot determine smallest node type")
 
     def ensure_cluster_is_running(self, cluster_id: str):
-        state = clusters.State
+        state = compute.State
         info = self.get(cluster_id)
         if info.state == state.TERMINATED:
             self.start(cluster_id).result()

databricks/sdk/mixins/dbfs.py

@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, AnyStr, BinaryIO, Iterable, Iterator, Type
 
 from databricks.sdk.core import DatabricksError
 
-from ..service import dbfs
+from ..service import files
 
 if TYPE_CHECKING:
     from _typeshed import Self
@@ -18,13 +18,13 @@ if TYPE_CHECKING:
 class _DbfsIO(BinaryIO):
     MAX_CHUNK_SIZE = 1024 * 1024
 
-    _status: dbfs.FileInfo = None
-    _created: dbfs.CreateResponse = None
+    _status: files.FileInfo = None
+    _created: files.CreateResponse = None
     _offset = 0
     _closed = False
 
     def __init__(self,
-                 api: dbfs.DbfsAPI,
+                 api: files.DbfsAPI,
                  path: str,
                  *,
                  read: bool = False,
@@ -308,12 +308,12 @@ class _DbfsPath(_Path):
         return f'<_DbfsPath {self._path}>'
 
 
-class DbfsExt(dbfs.DbfsAPI):
+class DbfsExt(files.DbfsAPI):
 
     def open(self, path: str, *, read: bool = False, write: bool = False, overwrite: bool = False) -> _DbfsIO:
         return _DbfsIO(self, path, read=read, write=write, overwrite=overwrite)
 
-    def list(self, path: str, *, recursive=False) -> Iterator[dbfs.FileInfo]:
+    def list(self, path: str, *, recursive=False) -> Iterator[files.FileInfo]:
         """List directory contents or file details.
 
         List the contents of a directory, or details of the file. If the file or directory does not exist,