databricks-sdk 0.44.1__py3-none-any.whl → 0.46.0__py3-none-any.whl

databricks/sdk/__init__.py

@@ -8,6 +8,7 @@ import databricks.sdk.dbutils as dbutils
 import databricks.sdk.service as service
 from databricks.sdk import azure
 from databricks.sdk.credentials_provider import CredentialsStrategy
+from databricks.sdk.data_plane import DataPlaneTokenSource
 from databricks.sdk.mixins.compute import ClustersExt
 from databricks.sdk.mixins.files import DbfsExt, FilesExt
 from databricks.sdk.mixins.jobs import JobsExt
@@ -63,7 +64,8 @@ from databricks.sdk.service.marketplace import (
     ProviderExchangeFiltersAPI, ProviderExchangesAPI, ProviderFilesAPI,
     ProviderListingsAPI, ProviderPersonalizationRequestsAPI,
     ProviderProviderAnalyticsDashboardsAPI, ProviderProvidersAPI)
-from databricks.sdk.service.ml import ExperimentsAPI, ModelRegistryAPI
+from databricks.sdk.service.ml import (ExperimentsAPI, ForecastingAPI,
+                                       ModelRegistryAPI)
 from databricks.sdk.service.oauth2 import (AccountFederationPolicyAPI,
                                            CustomAppIntegrationAPI,
                                            OAuthPublishedAppsAPI,
@@ -122,6 +124,7 @@ def _make_dbutils(config: client.Config):
 
     # We are in runtime, so we can use the runtime dbutils
     from databricks.sdk.runtime import dbutils as runtime_dbutils
+
     return runtime_dbutils
 
 
@@ -138,58 +141,62 @@ class WorkspaceClient:
     The WorkspaceClient is a client for the workspace-level Databricks REST API.
     """
 
-    def __init__(self,
-                 *,
-                 host: Optional[str] = None,
-                 account_id: Optional[str] = None,
-                 username: Optional[str] = None,
-                 password: Optional[str] = None,
-                 client_id: Optional[str] = None,
-                 client_secret: Optional[str] = None,
-                 token: Optional[str] = None,
-                 profile: Optional[str] = None,
-                 config_file: Optional[str] = None,
-                 azure_workspace_resource_id: Optional[str] = None,
-                 azure_client_secret: Optional[str] = None,
-                 azure_client_id: Optional[str] = None,
-                 azure_tenant_id: Optional[str] = None,
-                 azure_environment: Optional[str] = None,
-                 auth_type: Optional[str] = None,
-                 cluster_id: Optional[str] = None,
-                 google_credentials: Optional[str] = None,
-                 google_service_account: Optional[str] = None,
-                 debug_truncate_bytes: Optional[int] = None,
-                 debug_headers: Optional[bool] = None,
-                 product="unknown",
-                 product_version="0.0.0",
-                 credentials_strategy: Optional[CredentialsStrategy] = None,
-                 credentials_provider: Optional[CredentialsStrategy] = None,
-                 config: Optional[client.Config] = None):
+    def __init__(
+        self,
+        *,
+        host: Optional[str] = None,
+        account_id: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        token: Optional[str] = None,
+        profile: Optional[str] = None,
+        config_file: Optional[str] = None,
+        azure_workspace_resource_id: Optional[str] = None,
+        azure_client_secret: Optional[str] = None,
+        azure_client_id: Optional[str] = None,
+        azure_tenant_id: Optional[str] = None,
+        azure_environment: Optional[str] = None,
+        auth_type: Optional[str] = None,
+        cluster_id: Optional[str] = None,
+        google_credentials: Optional[str] = None,
+        google_service_account: Optional[str] = None,
+        debug_truncate_bytes: Optional[int] = None,
+        debug_headers: Optional[bool] = None,
+        product="unknown",
+        product_version="0.0.0",
+        credentials_strategy: Optional[CredentialsStrategy] = None,
+        credentials_provider: Optional[CredentialsStrategy] = None,
+        config: Optional[client.Config] = None,
+    ):
         if not config:
-            config = client.Config(host=host,
-                                   account_id=account_id,
-                                   username=username,
-                                   password=password,
-                                   client_id=client_id,
-                                   client_secret=client_secret,
-                                   token=token,
-                                   profile=profile,
-                                   config_file=config_file,
-                                   azure_workspace_resource_id=azure_workspace_resource_id,
-                                   azure_client_secret=azure_client_secret,
-                                   azure_client_id=azure_client_id,
-                                   azure_tenant_id=azure_tenant_id,
-                                   azure_environment=azure_environment,
-                                   auth_type=auth_type,
-                                   cluster_id=cluster_id,
-                                   google_credentials=google_credentials,
-                                   google_service_account=google_service_account,
-                                   credentials_strategy=credentials_strategy,
-                                   credentials_provider=credentials_provider,
-                                   debug_truncate_bytes=debug_truncate_bytes,
-                                   debug_headers=debug_headers,
-                                   product=product,
-                                   product_version=product_version)
+            config = client.Config(
+                host=host,
+                account_id=account_id,
+                username=username,
+                password=password,
+                client_id=client_id,
+                client_secret=client_secret,
+                token=token,
+                profile=profile,
+                config_file=config_file,
+                azure_workspace_resource_id=azure_workspace_resource_id,
+                azure_client_secret=azure_client_secret,
+                azure_client_id=azure_client_id,
+                azure_tenant_id=azure_tenant_id,
+                azure_environment=azure_environment,
+                auth_type=auth_type,
+                cluster_id=cluster_id,
+                google_credentials=google_credentials,
+                google_service_account=google_service_account,
+                credentials_strategy=credentials_strategy,
+                credentials_provider=credentials_provider,
+                debug_truncate_bytes=debug_truncate_bytes,
+                debug_headers=debug_headers,
+                product=product,
+                product_version=product_version,
+            )
         self._config = config.copy()
         self._dbutils = _make_dbutils(self._config)
         self._api_client = client.ApiClient(self._config)
@@ -212,7 +219,8 @@ class WorkspaceClient:
         self._consumer_installations = service.marketplace.ConsumerInstallationsAPI(self._api_client)
         self._consumer_listings = service.marketplace.ConsumerListingsAPI(self._api_client)
         self._consumer_personalization_requests = service.marketplace.ConsumerPersonalizationRequestsAPI(
-            self._api_client)
+            self._api_client
+        )
         self._consumer_providers = service.marketplace.ConsumerProvidersAPI(self._api_client)
         self._credentials = service.catalog.CredentialsAPI(self._api_client)
         self._credentials_manager = service.settings.CredentialsManagerAPI(self._api_client)
@@ -246,8 +254,7 @@ class WorkspaceClient:
         self._permission_migration = service.iam.PermissionMigrationAPI(self._api_client)
         self._permissions = service.iam.PermissionsAPI(self._api_client)
         self._pipelines = service.pipelines.PipelinesAPI(self._api_client)
-        self._policy_compliance_for_clusters = service.compute.PolicyComplianceForClustersAPI(
-            self._api_client)
+        self._policy_compliance_for_clusters = service.compute.PolicyComplianceForClustersAPI(self._api_client)
         self._policy_compliance_for_jobs = service.jobs.PolicyComplianceForJobsAPI(self._api_client)
         self._policy_families = service.compute.PolicyFamiliesAPI(self._api_client)
         self._provider_exchange_filters = service.marketplace.ProviderExchangeFiltersAPI(self._api_client)
@@ -255,9 +262,11 @@ class WorkspaceClient:
         self._provider_files = service.marketplace.ProviderFilesAPI(self._api_client)
         self._provider_listings = service.marketplace.ProviderListingsAPI(self._api_client)
         self._provider_personalization_requests = service.marketplace.ProviderPersonalizationRequestsAPI(
-            self._api_client)
+            self._api_client
+        )
         self._provider_provider_analytics_dashboards = service.marketplace.ProviderProviderAnalyticsDashboardsAPI(
-            self._api_client)
+            self._api_client
+        )
         self._provider_providers = service.marketplace.ProviderProvidersAPI(self._api_client)
         self._providers = service.sharing.ProvidersAPI(self._api_client)
         self._quality_monitors = service.catalog.QualityMonitorsAPI(self._api_client)
@@ -277,8 +286,12 @@ class WorkspaceClient:
         self._secrets = service.workspace.SecretsAPI(self._api_client)
         self._service_principals = service.iam.ServicePrincipalsAPI(self._api_client)
         self._serving_endpoints = serving_endpoints
+        serving_endpoints_data_plane_token_source = DataPlaneTokenSource(
+            self._config.host, self._config.oauth_token, not self._config.enable_experimental_async_token_refresh
+        )
         self._serving_endpoints_data_plane = service.serving.ServingEndpointsDataPlaneAPI(
-            self._api_client, serving_endpoints)
+            self._api_client, serving_endpoints, serving_endpoints_data_plane_token_source
+        )
         self._settings = service.settings.SettingsAPI(self._api_client)
         self._shares = service.sharing.SharesAPI(self._api_client)
         self._statement_execution = service.sql.StatementExecutionAPI(self._api_client)
@@ -297,6 +310,7 @@ class WorkspaceClient:
         self._workspace = WorkspaceExt(self._api_client)
         self._workspace_bindings = service.catalog.WorkspaceBindingsAPI(self._api_client)
         self._workspace_conf = service.settings.WorkspaceConfAPI(self._api_client)
+        self._forecasting = service.ml.ForecastingAPI(self._api_client)
 
     @property
     def config(self) -> client.Config:
@@ -606,8 +620,7 @@ class WorkspaceClient:
         return self._provider_personalization_requests
 
     @property
-    def provider_provider_analytics_dashboards(
-            self) -> service.marketplace.ProviderProviderAnalyticsDashboardsAPI:
+    def provider_provider_analytics_dashboards(self) -> service.marketplace.ProviderProviderAnalyticsDashboardsAPI:
         """Manage templated analytics solution for providers."""
         return self._provider_provider_analytics_dashboards
 
@@ -801,11 +814,14 @@ class WorkspaceClient:
         """This API allows updating known workspace settings for advanced users."""
         return self._workspace_conf
 
+    @property
+    def forecasting(self) -> service.ml.ForecastingAPI:
+        """The Forecasting API allows you to create and get serverless forecasting experiments."""
+        return self._forecasting
+
     def get_workspace_id(self) -> int:
         """Get the workspace ID of the workspace that this client is connected to."""
-        response = self._api_client.do("GET",
-                                       "/api/2.0/preview/scim/v2/Me",
-                                       response_headers=['X-Databricks-Org-Id'])
+        response = self._api_client.do("GET", "/api/2.0/preview/scim/v2/Me", response_headers=["X-Databricks-Org-Id"])
         return int(response["X-Databricks-Org-Id"])
 
     def __repr__(self):
@@ -817,58 +833,62 @@ class AccountClient:
     The AccountClient is a client for the account-level Databricks REST API.
     """
 
-    def __init__(self,
-                 *,
-                 host: Optional[str] = None,
-                 account_id: Optional[str] = None,
-                 username: Optional[str] = None,
-                 password: Optional[str] = None,
-                 client_id: Optional[str] = None,
-                 client_secret: Optional[str] = None,
-                 token: Optional[str] = None,
-                 profile: Optional[str] = None,
-                 config_file: Optional[str] = None,
-                 azure_workspace_resource_id: Optional[str] = None,
-                 azure_client_secret: Optional[str] = None,
-                 azure_client_id: Optional[str] = None,
-                 azure_tenant_id: Optional[str] = None,
-                 azure_environment: Optional[str] = None,
-                 auth_type: Optional[str] = None,
-                 cluster_id: Optional[str] = None,
-                 google_credentials: Optional[str] = None,
-                 google_service_account: Optional[str] = None,
-                 debug_truncate_bytes: Optional[int] = None,
-                 debug_headers: Optional[bool] = None,
-                 product="unknown",
-                 product_version="0.0.0",
-                 credentials_strategy: Optional[CredentialsStrategy] = None,
-                 credentials_provider: Optional[CredentialsStrategy] = None,
-                 config: Optional[client.Config] = None):
+    def __init__(
+        self,
+        *,
+        host: Optional[str] = None,
+        account_id: Optional[str] = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        token: Optional[str] = None,
+        profile: Optional[str] = None,
+        config_file: Optional[str] = None,
+        azure_workspace_resource_id: Optional[str] = None,
+        azure_client_secret: Optional[str] = None,
+        azure_client_id: Optional[str] = None,
+        azure_tenant_id: Optional[str] = None,
+        azure_environment: Optional[str] = None,
+        auth_type: Optional[str] = None,
+        cluster_id: Optional[str] = None,
+        google_credentials: Optional[str] = None,
+        google_service_account: Optional[str] = None,
+        debug_truncate_bytes: Optional[int] = None,
+        debug_headers: Optional[bool] = None,
+        product="unknown",
+        product_version="0.0.0",
+        credentials_strategy: Optional[CredentialsStrategy] = None,
+        credentials_provider: Optional[CredentialsStrategy] = None,
+        config: Optional[client.Config] = None,
+    ):
         if not config:
-            config = client.Config(host=host,
-                                   account_id=account_id,
-                                   username=username,
-                                   password=password,
-                                   client_id=client_id,
-                                   client_secret=client_secret,
-                                   token=token,
-                                   profile=profile,
-                                   config_file=config_file,
-                                   azure_workspace_resource_id=azure_workspace_resource_id,
-                                   azure_client_secret=azure_client_secret,
-                                   azure_client_id=azure_client_id,
-                                   azure_tenant_id=azure_tenant_id,
-                                   azure_environment=azure_environment,
-                                   auth_type=auth_type,
-                                   cluster_id=cluster_id,
-                                   google_credentials=google_credentials,
-                                   google_service_account=google_service_account,
-                                   credentials_strategy=credentials_strategy,
-                                   credentials_provider=credentials_provider,
-                                   debug_truncate_bytes=debug_truncate_bytes,
-                                   debug_headers=debug_headers,
-                                   product=product,
-                                   product_version=product_version)
+            config = client.Config(
+                host=host,
+                account_id=account_id,
+                username=username,
+                password=password,
+                client_id=client_id,
+                client_secret=client_secret,
+                token=token,
+                profile=profile,
+                config_file=config_file,
+                azure_workspace_resource_id=azure_workspace_resource_id,
+                azure_client_secret=azure_client_secret,
+                azure_client_id=azure_client_id,
+                azure_tenant_id=azure_tenant_id,
+                azure_environment=azure_environment,
+                auth_type=auth_type,
+                cluster_id=cluster_id,
+                google_credentials=google_credentials,
+                google_service_account=google_service_account,
+                credentials_strategy=credentials_strategy,
+                credentials_provider=credentials_provider,
+                debug_truncate_bytes=debug_truncate_bytes,
+                debug_headers=debug_headers,
+                product=product,
+                product_version=product_version,
+            )
         self._config = config.copy()
         self._api_client = client.ApiClient(self._config)
         self._access_control = service.iam.AccountAccessControlAPI(self._api_client)
@@ -888,8 +908,7 @@ class AccountClient:
         self._o_auth_published_apps = service.oauth2.OAuthPublishedAppsAPI(self._api_client)
         self._private_access = service.provisioning.PrivateAccessAPI(self._api_client)
         self._published_app_integration = service.oauth2.PublishedAppIntegrationAPI(self._api_client)
-        self._service_principal_federation_policy = service.oauth2.ServicePrincipalFederationPolicyAPI(
-            self._api_client)
+        self._service_principal_federation_policy = service.oauth2.ServicePrincipalFederationPolicyAPI(self._api_client)
         self._service_principal_secrets = service.oauth2.ServicePrincipalSecretsAPI(self._api_client)
         self._service_principals = service.iam.AccountServicePrincipalsAPI(self._api_client)
         self._settings = service.settings.AccountSettingsAPI(self._api_client)