diff --git a/.codegen.json b/.codegen.json index f98f968ab..516edbc7c 100644 --- a/.codegen.json +++ b/.codegen.json @@ -1,5 +1,5 @@ { - "mode": "py_v0", + "mode": "py_mod", "api_changelog": true, "version": { "databricks/sdk/version.py": "__version__ = \"$VERSION\"" diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 12fb465ab..f5d482394 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -94dc3e7289a19a90b167adf27316bd703a86f0eb \ No newline at end of file +e6971360e2752b3513d44a25d25f6cc5448056c8 \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index a4029b165..d240f2758 100755 --- a/.gitattributes +++ b/.gitattributes @@ -1,23 +1,80 @@ -databricks/sdk/__init__.py linguist-generated=true -databricks/sdk/databricks/errors/overrides.py linguist-generated=true -databricks/sdk/databricks/errors/platform.py linguist-generated=true -databricks/sdk/service/apps.py linguist-generated=true -databricks/sdk/service/billing.py linguist-generated=true -databricks/sdk/service/catalog.py linguist-generated=true -databricks/sdk/service/cleanrooms.py linguist-generated=true -databricks/sdk/service/compute.py linguist-generated=true -databricks/sdk/service/dashboards.py linguist-generated=true -databricks/sdk/service/files.py linguist-generated=true -databricks/sdk/service/iam.py linguist-generated=true -databricks/sdk/service/jobs.py linguist-generated=true -databricks/sdk/service/marketplace.py linguist-generated=true -databricks/sdk/service/ml.py linguist-generated=true -databricks/sdk/service/oauth2.py linguist-generated=true -databricks/sdk/service/pipelines.py linguist-generated=true -databricks/sdk/service/provisioning.py linguist-generated=true -databricks/sdk/service/serving.py linguist-generated=true -databricks/sdk/service/settings.py linguist-generated=true -databricks/sdk/service/sharing.py linguist-generated=true -databricks/sdk/service/sql.py linguist-generated=true -databricks/sdk/service/vectorsearch.py linguist-generated=true -databricks/sdk/service/workspace.py linguist-generated=true +databricks/sdk/apps/__init__.py linguist-generated=true +databricks/sdk/apps/v2/__init__.py linguist-generated=true +databricks/sdk/apps/v2/apps.py linguist-generated=true +databricks/sdk/apps/v2/client.py linguist-generated=true +databricks/sdk/billing/__init__.py linguist-generated=true +databricks/sdk/billing/v2/__init__.py linguist-generated=true +databricks/sdk/billing/v2/billing.py linguist-generated=true +databricks/sdk/billing/v2/client.py linguist-generated=true +databricks/sdk/catalog/__init__.py linguist-generated=true +databricks/sdk/catalog/v2/__init__.py linguist-generated=true +databricks/sdk/catalog/v2/catalog.py linguist-generated=true +databricks/sdk/catalog/v2/client.py linguist-generated=true +databricks/sdk/cleanrooms/__init__.py linguist-generated=true +databricks/sdk/cleanrooms/v2/__init__.py linguist-generated=true +databricks/sdk/cleanrooms/v2/cleanrooms.py linguist-generated=true +databricks/sdk/cleanrooms/v2/client.py linguist-generated=true +databricks/sdk/compute/__init__.py linguist-generated=true +databricks/sdk/compute/v2/__init__.py linguist-generated=true +databricks/sdk/compute/v2/client.py linguist-generated=true +databricks/sdk/compute/v2/compute.py linguist-generated=true +databricks/sdk/dashboards/__init__.py linguist-generated=true +databricks/sdk/dashboards/v2/__init__.py linguist-generated=true +databricks/sdk/dashboards/v2/client.py linguist-generated=true +databricks/sdk/dashboards/v2/dashboards.py linguist-generated=true +databricks/sdk/files/__init__.py linguist-generated=true +databricks/sdk/files/v2/__init__.py linguist-generated=true +databricks/sdk/files/v2/client.py linguist-generated=true +databricks/sdk/files/v2/files.py linguist-generated=true +databricks/sdk/iam/__init__.py linguist-generated=true +databricks/sdk/iam/v2/__init__.py linguist-generated=true +databricks/sdk/iam/v2/client.py linguist-generated=true +databricks/sdk/iam/v2/iam.py linguist-generated=true +databricks/sdk/jobs/__init__.py linguist-generated=true +databricks/sdk/jobs/v2/__init__.py linguist-generated=true +databricks/sdk/jobs/v2/client.py linguist-generated=true +databricks/sdk/jobs/v2/jobs.py linguist-generated=true +databricks/sdk/marketplace/__init__.py linguist-generated=true +databricks/sdk/marketplace/v2/__init__.py linguist-generated=true +databricks/sdk/marketplace/v2/client.py linguist-generated=true +databricks/sdk/marketplace/v2/marketplace.py linguist-generated=true +databricks/sdk/ml/__init__.py linguist-generated=true +databricks/sdk/ml/v2/__init__.py linguist-generated=true +databricks/sdk/ml/v2/client.py linguist-generated=true +databricks/sdk/ml/v2/ml.py linguist-generated=true +databricks/sdk/oauth2/__init__.py linguist-generated=true +databricks/sdk/oauth2/v2/__init__.py linguist-generated=true +databricks/sdk/oauth2/v2/client.py linguist-generated=true +databricks/sdk/oauth2/v2/oauth2.py linguist-generated=true +databricks/sdk/pipelines/__init__.py linguist-generated=true +databricks/sdk/pipelines/v2/__init__.py linguist-generated=true +databricks/sdk/pipelines/v2/client.py linguist-generated=true +databricks/sdk/pipelines/v2/pipelines.py linguist-generated=true +databricks/sdk/provisioning/__init__.py linguist-generated=true +databricks/sdk/provisioning/v2/__init__.py linguist-generated=true +databricks/sdk/provisioning/v2/client.py linguist-generated=true +databricks/sdk/provisioning/v2/provisioning.py linguist-generated=true +databricks/sdk/serving/__init__.py linguist-generated=true +databricks/sdk/serving/v2/__init__.py linguist-generated=true +databricks/sdk/serving/v2/client.py linguist-generated=true +databricks/sdk/serving/v2/serving.py linguist-generated=true +databricks/sdk/settings/__init__.py linguist-generated=true +databricks/sdk/settings/v2/__init__.py linguist-generated=true +databricks/sdk/settings/v2/client.py linguist-generated=true +databricks/sdk/settings/v2/settings.py linguist-generated=true +databricks/sdk/sharing/__init__.py linguist-generated=true +databricks/sdk/sharing/v2/__init__.py linguist-generated=true +databricks/sdk/sharing/v2/client.py linguist-generated=true +databricks/sdk/sharing/v2/sharing.py linguist-generated=true +databricks/sdk/sql/__init__.py linguist-generated=true +databricks/sdk/sql/v2/__init__.py linguist-generated=true +databricks/sdk/sql/v2/client.py linguist-generated=true +databricks/sdk/sql/v2/sql.py linguist-generated=true +databricks/sdk/vectorsearch/__init__.py linguist-generated=true +databricks/sdk/vectorsearch/v2/__init__.py linguist-generated=true +databricks/sdk/vectorsearch/v2/client.py linguist-generated=true +databricks/sdk/vectorsearch/v2/vectorsearch.py linguist-generated=true +databricks/sdk/workspace/__init__.py linguist-generated=true +databricks/sdk/workspace/v2/__init__.py linguist-generated=true +databricks/sdk/workspace/v2/client.py linguist-generated=true +databricks/sdk/workspace/v2/workspace.py linguist-generated=true diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 4cbfc92c1..c885c5d3a 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -9,104 +9,6 @@ from databricks.sdk.databricks import azure from databricks.sdk.databricks.credentials_provider import CredentialsStrategy from databricks.sdk.databricks.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 -from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt -from databricks.sdk.mixins.workspace import WorkspaceExt -from databricks.sdk.service.apps import AppsAPI -from databricks.sdk.service.billing import (BillableUsageAPI, BudgetPolicyAPI, - BudgetsAPI, LogDeliveryAPI, - UsageDashboardsAPI) -from databricks.sdk.service.catalog import (AccountMetastoreAssignmentsAPI, - AccountMetastoresAPI, - AccountStorageCredentialsAPI, - ArtifactAllowlistsAPI, CatalogsAPI, - ConnectionsAPI, CredentialsAPI, - ExternalLocationsAPI, FunctionsAPI, - GrantsAPI, MetastoresAPI, - ModelVersionsAPI, OnlineTablesAPI, - QualityMonitorsAPI, - RegisteredModelsAPI, - ResourceQuotasAPI, SchemasAPI, - StorageCredentialsAPI, - SystemSchemasAPI, - TableConstraintsAPI, TablesAPI, - TemporaryTableCredentialsAPI, - VolumesAPI, WorkspaceBindingsAPI) -from databricks.sdk.service.cleanrooms import (CleanRoomAssetsAPI, - CleanRoomsAPI, - CleanRoomTaskRunsAPI) -from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI, - CommandExecutionAPI, - GlobalInitScriptsAPI, - InstancePoolsAPI, - InstanceProfilesAPI, LibrariesAPI, - PolicyComplianceForClustersAPI, - PolicyFamiliesAPI) -from databricks.sdk.service.dashboards import (GenieAPI, LakeviewAPI, - LakeviewEmbeddedAPI, - QueryExecutionAPI) -from databricks.sdk.service.files import DbfsAPI, FilesAPI -from databricks.sdk.service.iam import (AccessControlAPI, - AccountAccessControlAPI, - AccountAccessControlProxyAPI, - AccountGroupsAPI, - AccountServicePrincipalsAPI, - AccountUsersAPI, CurrentUserAPI, - GroupsAPI, PermissionMigrationAPI, - PermissionsAPI, ServicePrincipalsAPI, - UsersAPI, WorkspaceAssignmentAPI) -from databricks.sdk.service.jobs import JobsAPI, PolicyComplianceForJobsAPI -from databricks.sdk.service.marketplace import ( - ConsumerFulfillmentsAPI, ConsumerInstallationsAPI, ConsumerListingsAPI, - ConsumerPersonalizationRequestsAPI, ConsumerProvidersAPI, - ProviderExchangeFiltersAPI, ProviderExchangesAPI, ProviderFilesAPI, - ProviderListingsAPI, ProviderPersonalizationRequestsAPI, - ProviderProviderAnalyticsDashboardsAPI, ProviderProvidersAPI) -from databricks.sdk.service.ml import (ExperimentsAPI, ForecastingAPI, - ModelRegistryAPI) -from databricks.sdk.service.oauth2 import (AccountFederationPolicyAPI, - CustomAppIntegrationAPI, - OAuthPublishedAppsAPI, - PublishedAppIntegrationAPI, - ServicePrincipalFederationPolicyAPI, - ServicePrincipalSecretsAPI) -from databricks.sdk.service.pipelines import PipelinesAPI -from databricks.sdk.service.provisioning import (CredentialsAPI, - EncryptionKeysAPI, - NetworksAPI, PrivateAccessAPI, - StorageAPI, VpcEndpointsAPI, - Workspace, WorkspacesAPI) -from databricks.sdk.service.serving import (ServingEndpointsAPI, - ServingEndpointsDataPlaneAPI) -from databricks.sdk.service.settings import ( - AccountIpAccessListsAPI, AccountSettingsAPI, - AibiDashboardEmbeddingAccessPolicyAPI, - AibiDashboardEmbeddingApprovedDomainsAPI, AutomaticClusterUpdateAPI, - ComplianceSecurityProfileAPI, CredentialsManagerAPI, - CspEnablementAccountAPI, DefaultNamespaceAPI, DisableLegacyAccessAPI, - DisableLegacyDbfsAPI, DisableLegacyFeaturesAPI, EnableIpAccessListsAPI, - EnhancedSecurityMonitoringAPI, EsmEnablementAccountAPI, IpAccessListsAPI, - NetworkConnectivityAPI, NotificationDestinationsAPI, PersonalComputeAPI, - RestrictWorkspaceAdminsAPI, SettingsAPI, TokenManagementAPI, TokensAPI, - WorkspaceConfAPI) -from databricks.sdk.service.sharing import (ProvidersAPI, - RecipientActivationAPI, - RecipientsAPI, SharesAPI) -from databricks.sdk.service.sql import (AlertsAPI, AlertsLegacyAPI, - DashboardsAPI, DashboardWidgetsAPI, - DataSourcesAPI, DbsqlPermissionsAPI, - QueriesAPI, QueriesLegacyAPI, - QueryHistoryAPI, - QueryVisualizationsAPI, - QueryVisualizationsLegacyAPI, - RedashConfigAPI, StatementExecutionAPI, - WarehousesAPI) -from databricks.sdk.service.vectorsearch import (VectorSearchEndpointsAPI, - VectorSearchIndexesAPI) -from databricks.sdk.service.workspace import (GitCredentialsAPI, ReposAPI, - SecretsAPI, WorkspaceAPI) _LOG = logging.getLogger(__name__) @@ -126,982 +28,3 @@ def _make_dbutils(config: client.Config): from databricks.sdk.databricks.runtime import dbutils as runtime_dbutils return runtime_dbutils - - -def _make_files_client(apiClient: client.ApiClient, config: client.Config): - if config.enable_experimental_files_api_client: - _LOG.info("Experimental Files API client is enabled") - return FilesExt(apiClient, config) - else: - return FilesAPI(apiClient) - - -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, - ): - 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, - ) - self._config = config.copy() - self._dbutils = _make_dbutils(self._config) - self._api_client = client.ApiClient(self._config) - serving_endpoints = ServingEndpointsExt(self._api_client) - self._access_control = service.iam.AccessControlAPI(self._api_client) - self._account_access_control_proxy = service.iam.AccountAccessControlProxyAPI(self._api_client) - self._alerts = service.sql.AlertsAPI(self._api_client) - self._alerts_legacy = service.sql.AlertsLegacyAPI(self._api_client) - self._apps = service.apps.AppsAPI(self._api_client) - self._artifact_allowlists = service.catalog.ArtifactAllowlistsAPI(self._api_client) - self._catalogs = service.catalog.CatalogsAPI(self._api_client) - self._clean_room_assets = service.cleanrooms.CleanRoomAssetsAPI(self._api_client) - self._clean_room_task_runs = service.cleanrooms.CleanRoomTaskRunsAPI(self._api_client) - self._clean_rooms = service.cleanrooms.CleanRoomsAPI(self._api_client) - self._cluster_policies = service.compute.ClusterPoliciesAPI(self._api_client) - self._clusters = ClustersExt(self._api_client) - self._command_execution = service.compute.CommandExecutionAPI(self._api_client) - self._connections = service.catalog.ConnectionsAPI(self._api_client) - self._consumer_fulfillments = service.marketplace.ConsumerFulfillmentsAPI(self._api_client) - 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._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) - self._current_user = service.iam.CurrentUserAPI(self._api_client) - self._dashboard_widgets = service.sql.DashboardWidgetsAPI(self._api_client) - self._dashboards = service.sql.DashboardsAPI(self._api_client) - self._data_sources = service.sql.DataSourcesAPI(self._api_client) - self._dbfs = DbfsExt(self._api_client) - self._dbsql_permissions = service.sql.DbsqlPermissionsAPI(self._api_client) - self._experiments = service.ml.ExperimentsAPI(self._api_client) - self._external_locations = service.catalog.ExternalLocationsAPI(self._api_client) - self._files = _make_files_client(self._api_client, self._config) - self._functions = service.catalog.FunctionsAPI(self._api_client) - self._genie = service.dashboards.GenieAPI(self._api_client) - self._git_credentials = service.workspace.GitCredentialsAPI(self._api_client) - self._global_init_scripts = service.compute.GlobalInitScriptsAPI(self._api_client) - self._grants = service.catalog.GrantsAPI(self._api_client) - self._groups = service.iam.GroupsAPI(self._api_client) - self._instance_pools = service.compute.InstancePoolsAPI(self._api_client) - self._instance_profiles = service.compute.InstanceProfilesAPI(self._api_client) - self._ip_access_lists = service.settings.IpAccessListsAPI(self._api_client) - self._jobs = JobsExt(self._api_client) - self._lakeview = service.dashboards.LakeviewAPI(self._api_client) - self._lakeview_embedded = service.dashboards.LakeviewEmbeddedAPI(self._api_client) - self._libraries = service.compute.LibrariesAPI(self._api_client) - self._metastores = service.catalog.MetastoresAPI(self._api_client) - self._model_registry = service.ml.ModelRegistryAPI(self._api_client) - self._model_versions = service.catalog.ModelVersionsAPI(self._api_client) - self._notification_destinations = service.settings.NotificationDestinationsAPI(self._api_client) - self._online_tables = service.catalog.OnlineTablesAPI(self._api_client) - 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_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) - self._provider_exchanges = service.marketplace.ProviderExchangesAPI(self._api_client) - 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._provider_provider_analytics_dashboards = service.marketplace.ProviderProviderAnalyticsDashboardsAPI( - 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) - self._queries = service.sql.QueriesAPI(self._api_client) - self._queries_legacy = service.sql.QueriesLegacyAPI(self._api_client) - self._query_execution = service.dashboards.QueryExecutionAPI(self._api_client) - self._query_history = service.sql.QueryHistoryAPI(self._api_client) - self._query_visualizations = service.sql.QueryVisualizationsAPI(self._api_client) - self._query_visualizations_legacy = service.sql.QueryVisualizationsLegacyAPI(self._api_client) - self._recipient_activation = service.sharing.RecipientActivationAPI(self._api_client) - self._recipients = service.sharing.RecipientsAPI(self._api_client) - self._redash_config = service.sql.RedashConfigAPI(self._api_client) - self._registered_models = service.catalog.RegisteredModelsAPI(self._api_client) - self._repos = service.workspace.ReposAPI(self._api_client) - self._resource_quotas = service.catalog.ResourceQuotasAPI(self._api_client) - self._schemas = service.catalog.SchemasAPI(self._api_client) - 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, 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) - self._storage_credentials = service.catalog.StorageCredentialsAPI(self._api_client) - self._system_schemas = service.catalog.SystemSchemasAPI(self._api_client) - self._table_constraints = service.catalog.TableConstraintsAPI(self._api_client) - self._tables = service.catalog.TablesAPI(self._api_client) - self._temporary_table_credentials = service.catalog.TemporaryTableCredentialsAPI(self._api_client) - self._token_management = service.settings.TokenManagementAPI(self._api_client) - self._tokens = service.settings.TokensAPI(self._api_client) - self._users = service.iam.UsersAPI(self._api_client) - self._vector_search_endpoints = service.vectorsearch.VectorSearchEndpointsAPI(self._api_client) - self._vector_search_indexes = service.vectorsearch.VectorSearchIndexesAPI(self._api_client) - self._volumes = service.catalog.VolumesAPI(self._api_client) - self._warehouses = service.sql.WarehousesAPI(self._api_client) - 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: - return self._config - - @property - def api_client(self) -> client.ApiClient: - return self._api_client - - @property - def dbutils(self) -> dbutils.RemoteDbUtils: - return self._dbutils - - @property - def access_control(self) -> service.iam.AccessControlAPI: - """Rule based Access Control for Databricks Resources.""" - return self._access_control - - @property - def account_access_control_proxy(self) -> service.iam.AccountAccessControlProxyAPI: - """These APIs manage access rules on resources in an account.""" - return self._account_access_control_proxy - - @property - def alerts(self) -> service.sql.AlertsAPI: - """The alerts API can be used to perform CRUD operations on alerts.""" - return self._alerts - - @property - def alerts_legacy(self) -> service.sql.AlertsLegacyAPI: - """The alerts API can be used to perform CRUD operations on alerts.""" - return self._alerts_legacy - - @property - def apps(self) -> service.apps.AppsAPI: - """Apps run directly on a customer’s Databricks instance, integrate with their data, use and extend Databricks services, and enable users to interact through single sign-on.""" - return self._apps - - @property - def artifact_allowlists(self) -> service.catalog.ArtifactAllowlistsAPI: - """In Databricks Runtime 13.3 and above, you can add libraries and init scripts to the `allowlist` in UC so that users can leverage these artifacts on compute configured with shared access mode.""" - return self._artifact_allowlists - - @property - def catalogs(self) -> service.catalog.CatalogsAPI: - """A catalog is the first layer of Unity Catalog’s three-level namespace.""" - return self._catalogs - - @property - def clean_room_assets(self) -> service.cleanrooms.CleanRoomAssetsAPI: - """Clean room assets are data and code objects — Tables, volumes, and notebooks that are shared with the clean room.""" - return self._clean_room_assets - - @property - def clean_room_task_runs(self) -> service.cleanrooms.CleanRoomTaskRunsAPI: - """Clean room task runs are the executions of notebooks in a clean room.""" - return self._clean_room_task_runs - - @property - def clean_rooms(self) -> service.cleanrooms.CleanRoomsAPI: - """A clean room uses Delta Sharing and serverless compute to provide a secure and privacy-protecting environment where multiple parties can work together on sensitive enterprise data without direct access to each other’s data.""" - return self._clean_rooms - - @property - def cluster_policies(self) -> service.compute.ClusterPoliciesAPI: - """You can use cluster policies to control users' ability to configure clusters based on a set of rules.""" - return self._cluster_policies - - @property - def clusters(self) -> ClustersExt: - """The Clusters API allows you to create, start, edit, list, terminate, and delete clusters.""" - return self._clusters - - @property - def command_execution(self) -> service.compute.CommandExecutionAPI: - """This API allows execution of Python, Scala, SQL, or R commands on running Databricks Clusters.""" - return self._command_execution - - @property - def connections(self) -> service.catalog.ConnectionsAPI: - """Connections allow for creating a connection to an external data source.""" - return self._connections - - @property - def consumer_fulfillments(self) -> service.marketplace.ConsumerFulfillmentsAPI: - """Fulfillments are entities that allow consumers to preview installations.""" - return self._consumer_fulfillments - - @property - def consumer_installations(self) -> service.marketplace.ConsumerInstallationsAPI: - """Installations are entities that allow consumers to interact with Databricks Marketplace listings.""" - return self._consumer_installations - - @property - def consumer_listings(self) -> service.marketplace.ConsumerListingsAPI: - """Listings are the core entities in the Marketplace.""" - return self._consumer_listings - - @property - def consumer_personalization_requests(self) -> service.marketplace.ConsumerPersonalizationRequestsAPI: - """Personalization Requests allow customers to interact with the individualized Marketplace listing flow.""" - return self._consumer_personalization_requests - - @property - def consumer_providers(self) -> service.marketplace.ConsumerProvidersAPI: - """Providers are the entities that publish listings to the Marketplace.""" - return self._consumer_providers - - @property - def credentials(self) -> service.catalog.CredentialsAPI: - """A credential represents an authentication and authorization mechanism for accessing services on your cloud tenant.""" - return self._credentials - - @property - def credentials_manager(self) -> service.settings.CredentialsManagerAPI: - """Credentials manager interacts with with Identity Providers to to perform token exchanges using stored credentials and refresh tokens.""" - return self._credentials_manager - - @property - def current_user(self) -> service.iam.CurrentUserAPI: - """This API allows retrieving information about currently authenticated user or service principal.""" - return self._current_user - - @property - def dashboard_widgets(self) -> service.sql.DashboardWidgetsAPI: - """This is an evolving API that facilitates the addition and removal of widgets from existing dashboards within the Databricks Workspace.""" - return self._dashboard_widgets - - @property - def dashboards(self) -> service.sql.DashboardsAPI: - """In general, there is little need to modify dashboards using the API.""" - return self._dashboards - - @property - def data_sources(self) -> service.sql.DataSourcesAPI: - """This API is provided to assist you in making new query objects.""" - return self._data_sources - - @property - def dbfs(self) -> DbfsExt: - """DBFS API makes it simple to interact with various data sources without having to include a users credentials every time to read a file.""" - return self._dbfs - - @property - def dbsql_permissions(self) -> service.sql.DbsqlPermissionsAPI: - """The SQL Permissions API is similar to the endpoints of the :method:permissions/set.""" - return self._dbsql_permissions - - @property - def experiments(self) -> service.ml.ExperimentsAPI: - """Experiments are the primary unit of organization in MLflow; all MLflow runs belong to an experiment.""" - return self._experiments - - @property - def external_locations(self) -> service.catalog.ExternalLocationsAPI: - """An external location is an object that combines a cloud storage path with a storage credential that authorizes access to the cloud storage path.""" - return self._external_locations - - @property - def files(self) -> service.files.FilesAPI: - """The Files API is a standard HTTP API that allows you to read, write, list, and delete files and directories by referring to their URI.""" - return self._files - - @property - def functions(self) -> service.catalog.FunctionsAPI: - """Functions implement User-Defined Functions (UDFs) in Unity Catalog.""" - return self._functions - - @property - def genie(self) -> service.dashboards.GenieAPI: - """Genie provides a no-code experience for business users, powered by AI/BI.""" - return self._genie - - @property - def git_credentials(self) -> service.workspace.GitCredentialsAPI: - """Registers personal access token for Databricks to do operations on behalf of the user.""" - return self._git_credentials - - @property - def global_init_scripts(self) -> service.compute.GlobalInitScriptsAPI: - """The Global Init Scripts API enables Workspace administrators to configure global initialization scripts for their workspace.""" - return self._global_init_scripts - - @property - def grants(self) -> service.catalog.GrantsAPI: - """In Unity Catalog, data is secure by default.""" - return self._grants - - @property - def groups(self) -> service.iam.GroupsAPI: - """Groups simplify identity management, making it easier to assign access to Databricks workspace, data, and other securable objects.""" - return self._groups - - @property - def instance_pools(self) -> service.compute.InstancePoolsAPI: - """Instance Pools API are used to create, edit, delete and list instance pools by using ready-to-use cloud instances which reduces a cluster start and auto-scaling times.""" - return self._instance_pools - - @property - def instance_profiles(self) -> service.compute.InstanceProfilesAPI: - """The Instance Profiles API allows admins to add, list, and remove instance profiles that users can launch clusters with.""" - return self._instance_profiles - - @property - def ip_access_lists(self) -> service.settings.IpAccessListsAPI: - """IP Access List enables admins to configure IP access lists.""" - return self._ip_access_lists - - @property - def jobs(self) -> JobsExt: - """The Jobs API allows you to create, edit, and delete jobs.""" - return self._jobs - - @property - def lakeview(self) -> service.dashboards.LakeviewAPI: - """These APIs provide specific management operations for Lakeview dashboards.""" - return self._lakeview - - @property - def lakeview_embedded(self) -> service.dashboards.LakeviewEmbeddedAPI: - """Token-based Lakeview APIs for embedding dashboards in external applications.""" - return self._lakeview_embedded - - @property - def libraries(self) -> service.compute.LibrariesAPI: - """The Libraries API allows you to install and uninstall libraries and get the status of libraries on a cluster.""" - return self._libraries - - @property - def metastores(self) -> service.catalog.MetastoresAPI: - """A metastore is the top-level container of objects in Unity Catalog.""" - return self._metastores - - @property - def model_registry(self) -> service.ml.ModelRegistryAPI: - """Note: This API reference documents APIs for the Workspace Model Registry.""" - return self._model_registry - - @property - def model_versions(self) -> service.catalog.ModelVersionsAPI: - """Databricks provides a hosted version of MLflow Model Registry in Unity Catalog.""" - return self._model_versions - - @property - def notification_destinations(self) -> service.settings.NotificationDestinationsAPI: - """The notification destinations API lets you programmatically manage a workspace's notification destinations.""" - return self._notification_destinations - - @property - def online_tables(self) -> service.catalog.OnlineTablesAPI: - """Online tables provide lower latency and higher QPS access to data from Delta tables.""" - return self._online_tables - - @property - def permission_migration(self) -> service.iam.PermissionMigrationAPI: - """APIs for migrating acl permissions, used only by the ucx tool: https://github.com/databrickslabs/ucx.""" - return self._permission_migration - - @property - def permissions(self) -> service.iam.PermissionsAPI: - """Permissions API are used to create read, write, edit, update and manage access for various users on different objects and endpoints.""" - return self._permissions - - @property - def pipelines(self) -> service.pipelines.PipelinesAPI: - """The Delta Live Tables API allows you to create, edit, delete, start, and view details about pipelines.""" - return self._pipelines - - @property - def policy_compliance_for_clusters(self) -> service.compute.PolicyComplianceForClustersAPI: - """The policy compliance APIs allow you to view and manage the policy compliance status of clusters in your workspace.""" - return self._policy_compliance_for_clusters - - @property - def policy_compliance_for_jobs(self) -> service.jobs.PolicyComplianceForJobsAPI: - """The compliance APIs allow you to view and manage the policy compliance status of jobs in your workspace.""" - return self._policy_compliance_for_jobs - - @property - def policy_families(self) -> service.compute.PolicyFamiliesAPI: - """View available policy families.""" - return self._policy_families - - @property - def provider_exchange_filters(self) -> service.marketplace.ProviderExchangeFiltersAPI: - """Marketplace exchanges filters curate which groups can access an exchange.""" - return self._provider_exchange_filters - - @property - def provider_exchanges(self) -> service.marketplace.ProviderExchangesAPI: - """Marketplace exchanges allow providers to share their listings with a curated set of customers.""" - return self._provider_exchanges - - @property - def provider_files(self) -> service.marketplace.ProviderFilesAPI: - """Marketplace offers a set of file APIs for various purposes such as preview notebooks and provider icons.""" - return self._provider_files - - @property - def provider_listings(self) -> service.marketplace.ProviderListingsAPI: - """Listings are the core entities in the Marketplace.""" - return self._provider_listings - - @property - def provider_personalization_requests(self) -> service.marketplace.ProviderPersonalizationRequestsAPI: - """Personalization requests are an alternate to instantly available listings.""" - return self._provider_personalization_requests - - @property - def provider_provider_analytics_dashboards(self) -> service.marketplace.ProviderProviderAnalyticsDashboardsAPI: - """Manage templated analytics solution for providers.""" - return self._provider_provider_analytics_dashboards - - @property - def provider_providers(self) -> service.marketplace.ProviderProvidersAPI: - """Providers are entities that manage assets in Marketplace.""" - return self._provider_providers - - @property - def providers(self) -> service.sharing.ProvidersAPI: - """A data provider is an object representing the organization in the real world who shares the data.""" - return self._providers - - @property - def quality_monitors(self) -> service.catalog.QualityMonitorsAPI: - """A monitor computes and monitors data or model quality metrics for a table over time.""" - return self._quality_monitors - - @property - def queries(self) -> service.sql.QueriesAPI: - """The queries API can be used to perform CRUD operations on queries.""" - return self._queries - - @property - def queries_legacy(self) -> service.sql.QueriesLegacyAPI: - """These endpoints are used for CRUD operations on query definitions.""" - return self._queries_legacy - - @property - def query_execution(self) -> service.dashboards.QueryExecutionAPI: - """Query execution APIs for AI / BI Dashboards.""" - return self._query_execution - - @property - def query_history(self) -> service.sql.QueryHistoryAPI: - """A service responsible for storing and retrieving the list of queries run against SQL endpoints and serverless compute.""" - return self._query_history - - @property - def query_visualizations(self) -> service.sql.QueryVisualizationsAPI: - """This is an evolving API that facilitates the addition and removal of visualizations from existing queries in the Databricks Workspace.""" - return self._query_visualizations - - @property - def query_visualizations_legacy(self) -> service.sql.QueryVisualizationsLegacyAPI: - """This is an evolving API that facilitates the addition and removal of vizualisations from existing queries within the Databricks Workspace.""" - return self._query_visualizations_legacy - - @property - def recipient_activation(self) -> service.sharing.RecipientActivationAPI: - """The Recipient Activation API is only applicable in the open sharing model where the recipient object has the authentication type of `TOKEN`.""" - return self._recipient_activation - - @property - def recipients(self) -> service.sharing.RecipientsAPI: - """A recipient is an object you create using :method:recipients/create to represent an organization which you want to allow access shares.""" - return self._recipients - - @property - def redash_config(self) -> service.sql.RedashConfigAPI: - """Redash V2 service for workspace configurations (internal).""" - return self._redash_config - - @property - def registered_models(self) -> service.catalog.RegisteredModelsAPI: - """Databricks provides a hosted version of MLflow Model Registry in Unity Catalog.""" - return self._registered_models - - @property - def repos(self) -> service.workspace.ReposAPI: - """The Repos API allows users to manage their git repos.""" - return self._repos - - @property - def resource_quotas(self) -> service.catalog.ResourceQuotasAPI: - """Unity Catalog enforces resource quotas on all securable objects, which limits the number of resources that can be created.""" - return self._resource_quotas - - @property - def schemas(self) -> service.catalog.SchemasAPI: - """A schema (also called a database) is the second layer of Unity Catalog’s three-level namespace.""" - return self._schemas - - @property - def secrets(self) -> service.workspace.SecretsAPI: - """The Secrets API allows you to manage secrets, secret scopes, and access permissions.""" - return self._secrets - - @property - def service_principals(self) -> service.iam.ServicePrincipalsAPI: - """Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD platforms.""" - return self._service_principals - - @property - def serving_endpoints(self) -> ServingEndpointsExt: - """The Serving Endpoints API allows you to create, update, and delete model serving endpoints.""" - return self._serving_endpoints - - @property - def serving_endpoints_data_plane(self) -> service.serving.ServingEndpointsDataPlaneAPI: - """Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints for Serving endpoints service.""" - return self._serving_endpoints_data_plane - - @property - def settings(self) -> service.settings.SettingsAPI: - """Workspace Settings API allows users to manage settings at the workspace level.""" - return self._settings - - @property - def shares(self) -> service.sharing.SharesAPI: - """A share is a container instantiated with :method:shares/create.""" - return self._shares - - @property - def statement_execution(self) -> service.sql.StatementExecutionAPI: - """The Databricks SQL Statement Execution API can be used to execute SQL statements on a SQL warehouse and fetch the result.""" - return self._statement_execution - - @property - def storage_credentials(self) -> service.catalog.StorageCredentialsAPI: - """A storage credential represents an authentication and authorization mechanism for accessing data stored on your cloud tenant.""" - return self._storage_credentials - - @property - def system_schemas(self) -> service.catalog.SystemSchemasAPI: - """A system schema is a schema that lives within the system catalog.""" - return self._system_schemas - - @property - def table_constraints(self) -> service.catalog.TableConstraintsAPI: - """Primary key and foreign key constraints encode relationships between fields in tables.""" - return self._table_constraints - - @property - def tables(self) -> service.catalog.TablesAPI: - """A table resides in the third layer of Unity Catalog’s three-level namespace.""" - return self._tables - - @property - def temporary_table_credentials(self) -> service.catalog.TemporaryTableCredentialsAPI: - """Temporary Table Credentials refer to short-lived, downscoped credentials used to access cloud storage locationswhere table data is stored in Databricks.""" - return self._temporary_table_credentials - - @property - def token_management(self) -> service.settings.TokenManagementAPI: - """Enables administrators to get all tokens and delete tokens for other users.""" - return self._token_management - - @property - def tokens(self) -> service.settings.TokensAPI: - """The Token API allows you to create, list, and revoke tokens that can be used to authenticate and access Databricks REST APIs.""" - return self._tokens - - @property - def users(self) -> service.iam.UsersAPI: - """User identities recognized by Databricks and represented by email addresses.""" - return self._users - - @property - def vector_search_endpoints(self) -> service.vectorsearch.VectorSearchEndpointsAPI: - """**Endpoint**: Represents the compute resources to host vector search indexes.""" - return self._vector_search_endpoints - - @property - def vector_search_indexes(self) -> service.vectorsearch.VectorSearchIndexesAPI: - """**Index**: An efficient representation of your embedding vectors that supports real-time and efficient approximate nearest neighbor (ANN) search queries.""" - return self._vector_search_indexes - - @property - def volumes(self) -> service.catalog.VolumesAPI: - """Volumes are a Unity Catalog (UC) capability for accessing, storing, governing, organizing and processing files.""" - return self._volumes - - @property - def warehouses(self) -> service.sql.WarehousesAPI: - """A SQL warehouse is a compute resource that lets you run SQL commands on data objects within Databricks SQL.""" - return self._warehouses - - @property - def workspace(self) -> WorkspaceExt: - """The Workspace API allows you to list, import, export, and delete notebooks and folders.""" - return self._workspace - - @property - def workspace_bindings(self) -> service.catalog.WorkspaceBindingsAPI: - """A securable in Databricks can be configured as __OPEN__ or __ISOLATED__.""" - return self._workspace_bindings - - @property - def workspace_conf(self) -> service.settings.WorkspaceConfAPI: - """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"]) - return int(response["X-Databricks-Org-Id"]) - - def __repr__(self): - return f"WorkspaceClient(host='{self._config.host}', auth_type='{self._config.auth_type}', ...)" - - -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, - ): - 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, - ) - self._config = config.copy() - self._api_client = client.ApiClient(self._config) - self._access_control = service.iam.AccountAccessControlAPI(self._api_client) - self._billable_usage = service.billing.BillableUsageAPI(self._api_client) - self._budget_policy = service.billing.BudgetPolicyAPI(self._api_client) - self._credentials = service.provisioning.CredentialsAPI(self._api_client) - self._custom_app_integration = service.oauth2.CustomAppIntegrationAPI(self._api_client) - self._encryption_keys = service.provisioning.EncryptionKeysAPI(self._api_client) - self._federation_policy = service.oauth2.AccountFederationPolicyAPI(self._api_client) - self._groups = service.iam.AccountGroupsAPI(self._api_client) - self._ip_access_lists = service.settings.AccountIpAccessListsAPI(self._api_client) - self._log_delivery = service.billing.LogDeliveryAPI(self._api_client) - self._metastore_assignments = service.catalog.AccountMetastoreAssignmentsAPI(self._api_client) - self._metastores = service.catalog.AccountMetastoresAPI(self._api_client) - self._network_connectivity = service.settings.NetworkConnectivityAPI(self._api_client) - self._networks = service.provisioning.NetworksAPI(self._api_client) - 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_secrets = service.oauth2.ServicePrincipalSecretsAPI(self._api_client) - self._service_principals = service.iam.AccountServicePrincipalsAPI(self._api_client) - self._settings = service.settings.AccountSettingsAPI(self._api_client) - self._storage = service.provisioning.StorageAPI(self._api_client) - self._storage_credentials = service.catalog.AccountStorageCredentialsAPI(self._api_client) - self._usage_dashboards = service.billing.UsageDashboardsAPI(self._api_client) - self._users = service.iam.AccountUsersAPI(self._api_client) - self._vpc_endpoints = service.provisioning.VpcEndpointsAPI(self._api_client) - self._workspace_assignment = service.iam.WorkspaceAssignmentAPI(self._api_client) - self._workspaces = service.provisioning.WorkspacesAPI(self._api_client) - self._budgets = service.billing.BudgetsAPI(self._api_client) - - @property - def config(self) -> client.Config: - return self._config - - @property - def api_client(self) -> client.ApiClient: - return self._api_client - - @property - def access_control(self) -> service.iam.AccountAccessControlAPI: - """These APIs manage access rules on resources in an account.""" - return self._access_control - - @property - def billable_usage(self) -> service.billing.BillableUsageAPI: - """This API allows you to download billable usage logs for the specified account and date range.""" - return self._billable_usage - - @property - def budget_policy(self) -> service.billing.BudgetPolicyAPI: - """A service serves REST API about Budget policies.""" - return self._budget_policy - - @property - def credentials(self) -> service.provisioning.CredentialsAPI: - """These APIs manage credential configurations for this workspace.""" - return self._credentials - - @property - def custom_app_integration(self) -> service.oauth2.CustomAppIntegrationAPI: - """These APIs enable administrators to manage custom OAuth app integrations, which is required for adding/using Custom OAuth App Integration like Tableau Cloud for Databricks in AWS cloud.""" - return self._custom_app_integration - - @property - def encryption_keys(self) -> service.provisioning.EncryptionKeysAPI: - """These APIs manage encryption key configurations for this workspace (optional).""" - return self._encryption_keys - - @property - def federation_policy(self) -> service.oauth2.AccountFederationPolicyAPI: - """These APIs manage account federation policies.""" - return self._federation_policy - - @property - def groups(self) -> service.iam.AccountGroupsAPI: - """Groups simplify identity management, making it easier to assign access to Databricks account, data, and other securable objects.""" - return self._groups - - @property - def ip_access_lists(self) -> service.settings.AccountIpAccessListsAPI: - """The Accounts IP Access List API enables account admins to configure IP access lists for access to the account console.""" - return self._ip_access_lists - - @property - def log_delivery(self) -> service.billing.LogDeliveryAPI: - """These APIs manage log delivery configurations for this account.""" - return self._log_delivery - - @property - def metastore_assignments(self) -> service.catalog.AccountMetastoreAssignmentsAPI: - """These APIs manage metastore assignments to a workspace.""" - return self._metastore_assignments - - @property - def metastores(self) -> service.catalog.AccountMetastoresAPI: - """These APIs manage Unity Catalog metastores for an account.""" - return self._metastores - - @property - def network_connectivity(self) -> service.settings.NetworkConnectivityAPI: - """These APIs provide configurations for the network connectivity of your workspaces for serverless compute resources.""" - return self._network_connectivity - - @property - def networks(self) -> service.provisioning.NetworksAPI: - """These APIs manage network configurations for customer-managed VPCs (optional).""" - return self._networks - - @property - def o_auth_published_apps(self) -> service.oauth2.OAuthPublishedAppsAPI: - """These APIs enable administrators to view all the available published OAuth applications in Databricks.""" - return self._o_auth_published_apps - - @property - def private_access(self) -> service.provisioning.PrivateAccessAPI: - """These APIs manage private access settings for this account.""" - return self._private_access - - @property - def published_app_integration(self) -> service.oauth2.PublishedAppIntegrationAPI: - """These APIs enable administrators to manage published OAuth app integrations, which is required for adding/using Published OAuth App Integration like Tableau Desktop for Databricks in AWS cloud.""" - return self._published_app_integration - - @property - def service_principal_federation_policy(self) -> service.oauth2.ServicePrincipalFederationPolicyAPI: - """These APIs manage service principal federation policies.""" - return self._service_principal_federation_policy - - @property - def service_principal_secrets(self) -> service.oauth2.ServicePrincipalSecretsAPI: - """These APIs enable administrators to manage service principal secrets.""" - return self._service_principal_secrets - - @property - def service_principals(self) -> service.iam.AccountServicePrincipalsAPI: - """Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD platforms.""" - return self._service_principals - - @property - def settings(self) -> service.settings.AccountSettingsAPI: - """Accounts Settings API allows users to manage settings at the account level.""" - return self._settings - - @property - def storage(self) -> service.provisioning.StorageAPI: - """These APIs manage storage configurations for this workspace.""" - return self._storage - - @property - def storage_credentials(self) -> service.catalog.AccountStorageCredentialsAPI: - """These APIs manage storage credentials for a particular metastore.""" - return self._storage_credentials - - @property - def usage_dashboards(self) -> service.billing.UsageDashboardsAPI: - """These APIs manage usage dashboards for this account.""" - return self._usage_dashboards - - @property - def users(self) -> service.iam.AccountUsersAPI: - """User identities recognized by Databricks and represented by email addresses.""" - return self._users - - @property - def vpc_endpoints(self) -> service.provisioning.VpcEndpointsAPI: - """These APIs manage VPC endpoint configurations for this account.""" - return self._vpc_endpoints - - @property - def workspace_assignment(self) -> service.iam.WorkspaceAssignmentAPI: - """The Workspace Permission Assignment API allows you to manage workspace permissions for principals in your account.""" - return self._workspace_assignment - - @property - def workspaces(self) -> service.provisioning.WorkspacesAPI: - """These APIs manage workspaces for this account.""" - return self._workspaces - - @property - def budgets(self) -> service.billing.BudgetsAPI: - """These APIs manage budget configurations for this account.""" - return self._budgets - - def get_workspace_client(self, workspace: Workspace) -> WorkspaceClient: - """Constructs a ``WorkspaceClient`` for the given workspace. - - Returns a ``WorkspaceClient`` that is configured to use the same - credentials as this ``AccountClient``. The underlying config is - copied from this ``AccountClient``, but the ``host`` and - ``azure_workspace_resource_id`` are overridden to match the - given workspace, and the ``account_id`` field is cleared. - - Usage: - - .. code-block:: - - wss = list(a.workspaces.list()) - if len(wss) == 0: - pytest.skip("no workspaces") - w = a.get_workspace_client(wss[0]) - assert w.current_user.me().active - - :param workspace: The workspace to construct a client for. - :return: A ``WorkspaceClient`` for the given workspace. - """ - config = self._config.deep_copy() - config.host = config.environment.deployment_url(workspace.deployment_name) - config.azure_workspace_resource_id = azure.get_azure_resource_id(workspace) - config.account_id = None - config.init_auth() - return WorkspaceClient(config=config) - - def __repr__(self): - return f"AccountClient(account_id='{self._config.account_id}', auth_type='{self._config.auth_type}', ...)" diff --git a/databricks/sdk/mixins/__init__.py b/databricks/sdk/apps/__init__.py similarity index 100% rename from databricks/sdk/mixins/__init__.py rename to databricks/sdk/apps/__init__.py diff --git a/databricks/sdk/apps/v2/__init__.py b/databricks/sdk/apps/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/service/apps.py b/databricks/sdk/apps/v2/apps.py old mode 100755 new mode 100644 similarity index 87% rename from databricks/sdk/service/apps.py rename to databricks/sdk/apps/v2/apps.py index 766f99581..0beb5434c --- a/databricks/sdk/service/apps.py +++ b/databricks/sdk/apps/v2/apps.py @@ -3,19 +3,14 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -50,9 +45,16 @@ class App: effective_budget_policy_id: Optional[str] = None + effective_user_api_scopes: Optional[List[str]] = None + """The effective api scopes granted to the user access token.""" + id: Optional[str] = None """The unique identifier of the app.""" + oauth2_app_client_id: Optional[str] = None + + oauth2_app_integration_id: Optional[str] = None + pending_deployment: Optional[AppDeployment] = None """The pending deployment of the app. A deployment is considered pending when it is being prepared for deployment to the app compute.""" @@ -75,6 +77,8 @@ class App: url: Optional[str] = None """The URL of the app once it is deployed.""" + user_api_scopes: Optional[List[str]] = None + def as_dict(self) -> dict: """Serializes the App into a dictionary suitable for use as a JSON request body.""" body = {} @@ -96,10 +100,16 @@ def as_dict(self) -> dict: body["description"] = self.description if self.effective_budget_policy_id is not None: body["effective_budget_policy_id"] = self.effective_budget_policy_id + if self.effective_user_api_scopes: + body["effective_user_api_scopes"] = [v for v in self.effective_user_api_scopes] if self.id is not None: body["id"] = self.id if self.name is not None: body["name"] = self.name + if self.oauth2_app_client_id is not None: + body["oauth2_app_client_id"] = self.oauth2_app_client_id + if self.oauth2_app_integration_id is not None: + body["oauth2_app_integration_id"] = self.oauth2_app_integration_id if self.pending_deployment: body["pending_deployment"] = self.pending_deployment.as_dict() if self.resources: @@ -116,6 +126,8 @@ def as_dict(self) -> dict: body["updater"] = self.updater if self.url is not None: body["url"] = self.url + if self.user_api_scopes: + body["user_api_scopes"] = [v for v in self.user_api_scopes] return body def as_shallow_dict(self) -> dict: @@ -139,10 +151,16 @@ def as_shallow_dict(self) -> dict: body["description"] = self.description if self.effective_budget_policy_id is not None: body["effective_budget_policy_id"] = self.effective_budget_policy_id + if self.effective_user_api_scopes: + body["effective_user_api_scopes"] = self.effective_user_api_scopes if self.id is not None: body["id"] = self.id if self.name is not None: body["name"] = self.name + if self.oauth2_app_client_id is not None: + body["oauth2_app_client_id"] = self.oauth2_app_client_id + if self.oauth2_app_integration_id is not None: + body["oauth2_app_integration_id"] = self.oauth2_app_integration_id if self.pending_deployment: body["pending_deployment"] = self.pending_deployment if self.resources: @@ -159,6 +177,8 @@ def as_shallow_dict(self) -> dict: body["updater"] = self.updater if self.url is not None: body["url"] = self.url + if self.user_api_scopes: + body["user_api_scopes"] = self.user_api_scopes return body @classmethod @@ -174,8 +194,11 @@ def from_dict(cls, d: Dict[str, Any]) -> App: default_source_code_path=d.get("default_source_code_path", None), description=d.get("description", None), effective_budget_policy_id=d.get("effective_budget_policy_id", None), + effective_user_api_scopes=d.get("effective_user_api_scopes", None), id=d.get("id", None), name=d.get("name", None), + oauth2_app_client_id=d.get("oauth2_app_client_id", None), + oauth2_app_integration_id=d.get("oauth2_app_integration_id", None), pending_deployment=_from_dict(d, "pending_deployment", AppDeployment), resources=_repeated_dict(d, "resources", AppResource), service_principal_client_id=d.get("service_principal_client_id", None), @@ -184,6 +207,7 @@ def from_dict(cls, d: Dict[str, Any]) -> App: update_time=d.get("update_time", None), updater=d.get("updater", None), url=d.get("url", None), + user_api_scopes=d.get("user_api_scopes", None), ) @@ -1044,107 +1068,7 @@ class AppsAPI: def __init__(self, api_client): self._api = api_client - def wait_get_app_active( - self, name: str, timeout=timedelta(minutes=20), callback: Optional[Callable[[App], None]] = None - ) -> App: - deadline = time.time() + timeout.total_seconds() - target_states = (ComputeState.ACTIVE,) - failure_states = ( - ComputeState.ERROR, - ComputeState.STOPPED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(name=name) - status = poll.compute_status.state - status_message = f"current status: {status}" - if poll.compute_status: - status_message = poll.compute_status.message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach ACTIVE, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"name={name}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_get_deployment_app_succeeded( - self, - app_name: str, - deployment_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[AppDeployment], None]] = None, - ) -> AppDeployment: - deadline = time.time() + timeout.total_seconds() - target_states = (AppDeploymentState.SUCCEEDED,) - failure_states = (AppDeploymentState.FAILED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get_deployment(app_name=app_name, deployment_id=deployment_id) - status = poll.status.state - status_message = f"current status: {status}" - if poll.status: - status_message = poll.status.message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach SUCCEEDED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"app_name={app_name}, deployment_id={deployment_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_get_app_stopped( - self, name: str, timeout=timedelta(minutes=20), callback: Optional[Callable[[App], None]] = None - ) -> App: - deadline = time.time() + timeout.total_seconds() - target_states = (ComputeState.STOPPED,) - failure_states = (ComputeState.ERROR,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(name=name) - status = poll.compute_status.state - status_message = f"current status: {status}" - if poll.compute_status: - status_message = poll.compute_status.message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach STOPPED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"name={name}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def create(self, *, app: Optional[App] = None, no_compute: Optional[bool] = None) -> Wait[App]: + def create(self, *, app: Optional[App] = None, no_compute: Optional[bool] = None) -> App: """Create an app. Creates a new app. @@ -1155,7 +1079,7 @@ def create(self, *, app: Optional[App] = None, no_compute: Optional[bool] = None :returns: Long-running operation waiter for :class:`App`. - See :method:wait_get_app_active for more details. + See :method:WaitGetAppActive for more details. """ body = app.as_dict() query = {} @@ -1166,13 +1090,8 @@ def create(self, *, app: Optional[App] = None, no_compute: Optional[bool] = None "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/apps", query=query, body=body, headers=headers) - return Wait(self.wait_get_app_active, response=App.from_dict(op_response), name=op_response["name"]) - - def create_and_wait( - self, *, app: Optional[App] = None, no_compute: Optional[bool] = None, timeout=timedelta(minutes=20) - ) -> App: - return self.create(app=app, no_compute=no_compute).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/apps", query=query, body=body, headers=headers) + return App.from_dict(res) def delete(self, name: str) -> App: """Delete an app. @@ -1192,7 +1111,7 @@ def delete(self, name: str) -> App: res = self._api.do("DELETE", f"/api/2.0/apps/{name}", headers=headers) return App.from_dict(res) - def deploy(self, app_name: str, *, app_deployment: Optional[AppDeployment] = None) -> Wait[AppDeployment]: + def deploy(self, app_name: str, *, app_deployment: Optional[AppDeployment] = None) -> AppDeployment: """Create an app deployment. Creates an app deployment for the app with the supplied name. @@ -1203,7 +1122,7 @@ def deploy(self, app_name: str, *, app_deployment: Optional[AppDeployment] = Non :returns: Long-running operation waiter for :class:`AppDeployment`. - See :method:wait_get_deployment_app_succeeded for more details. + See :method:WaitGetDeploymentAppSucceeded for more details. """ body = app_deployment.as_dict() headers = { @@ -1211,18 +1130,8 @@ def deploy(self, app_name: str, *, app_deployment: Optional[AppDeployment] = Non "Content-Type": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/apps/{app_name}/deployments", body=body, headers=headers) - return Wait( - self.wait_get_deployment_app_succeeded, - response=AppDeployment.from_dict(op_response), - app_name=app_name, - deployment_id=op_response["deployment_id"], - ) - - def deploy_and_wait( - self, app_name: str, *, app_deployment: Optional[AppDeployment] = None, timeout=timedelta(minutes=20) - ) -> AppDeployment: - return self.deploy(app_deployment=app_deployment, app_name=app_name).result(timeout=timeout) + res = self._api.do("POST", f"/api/2.0/apps/{app_name}/deployments", body=body, headers=headers) + return AppDeployment.from_dict(res) def get(self, name: str) -> App: """Get an app. @@ -1389,7 +1298,7 @@ def set_permissions( res = self._api.do("PUT", f"/api/2.0/permissions/apps/{app_name}", body=body, headers=headers) return AppPermissions.from_dict(res) - def start(self, name: str) -> Wait[App]: + def start(self, name: str) -> App: """Start an app. Start the last active deployment of the app in the workspace. @@ -1399,7 +1308,7 @@ def start(self, name: str) -> Wait[App]: :returns: Long-running operation waiter for :class:`App`. - See :method:wait_get_app_active for more details. + See :method:WaitGetAppActive for more details. """ headers = { @@ -1407,13 +1316,10 @@ def start(self, name: str) -> Wait[App]: "Content-Type": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/apps/{name}/start", headers=headers) - return Wait(self.wait_get_app_active, response=App.from_dict(op_response), name=op_response["name"]) - - def start_and_wait(self, name: str, timeout=timedelta(minutes=20)) -> App: - return self.start(name=name).result(timeout=timeout) + res = self._api.do("POST", f"/api/2.0/apps/{name}/start", headers=headers) + return App.from_dict(res) - def stop(self, name: str) -> Wait[App]: + def stop(self, name: str) -> App: """Stop an app. Stops the active deployment of the app in the workspace. @@ -1423,7 +1329,7 @@ def stop(self, name: str) -> Wait[App]: :returns: Long-running operation waiter for :class:`App`. - See :method:wait_get_app_stopped for more details. + See :method:WaitGetAppStopped for more details. """ headers = { @@ -1431,11 +1337,8 @@ def stop(self, name: str) -> Wait[App]: "Content-Type": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/apps/{name}/stop", headers=headers) - return Wait(self.wait_get_app_stopped, response=App.from_dict(op_response), name=op_response["name"]) - - def stop_and_wait(self, name: str, timeout=timedelta(minutes=20)) -> App: - return self.stop(name=name).result(timeout=timeout) + res = self._api.do("POST", f"/api/2.0/apps/{name}/stop", headers=headers) + return App.from_dict(res) def update(self, name: str, *, app: Optional[App] = None) -> App: """Update an app. diff --git a/databricks/sdk/apps/v2/client.py b/databricks/sdk/apps/v2/client.py new file mode 100644 index 000000000..ff2cc6315 --- /dev/null +++ b/databricks/sdk/apps/v2/client.py @@ -0,0 +1,78 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .apps import AppsAPI + +_LOG = logging.getLogger(__name__) + + +class AppsClient(AppsAPI): + """ + Apps run directly on a customer’s Databricks instance, integrate with their data, use and + extend Databricks services, and enable users to interact through single sign-on. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/billing/__init__.py b/databricks/sdk/billing/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/billing/v2/__init__.py b/databricks/sdk/billing/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/service/billing.py b/databricks/sdk/billing/v2/billing.py similarity index 97% rename from databricks/sdk/service/billing.py rename to databricks/sdk/billing/v2/billing.py index 4e8e4bca6..7edf243d7 100755 --- a/databricks/sdk/service/billing.py +++ b/databricks/sdk/billing/v2/billing.py @@ -7,13 +7,10 @@ from enum import Enum from typing import Any, BinaryIO, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - -from databricks.sdk.service import compute - # all definitions in this file are in alphabetical order @@ -364,7 +361,11 @@ def from_dict(cls, d: Dict[str, Any]) -> BudgetConfigurationFilterWorkspaceIdCla class BudgetPolicy: """Contains the BudgetPolicy details.""" - custom_tags: Optional[List[compute.CustomPolicyTag]] = None + binding_workspace_ids: Optional[List[int]] = None + """List of workspaces that this budget policy will be exclusively bound to. An empty binding + implies that this budget policy is open to any workspace in the account.""" + + custom_tags: Optional[List[CustomPolicyTag]] = None """A list of tags defined by the customer. At most 20 entries are allowed per policy.""" policy_id: Optional[str] = None @@ -372,11 +373,14 @@ class BudgetPolicy: policy_name: Optional[str] = None """The name of the policy. - Must be unique among active policies. - Can contain only characters - from the ISO 8859-1 (latin1) set.""" + from the ISO 8859-1 (latin1) set. - Can't start with reserved keywords such as + `databricks:default-policy`.""" def as_dict(self) -> dict: """Serializes the BudgetPolicy into a dictionary suitable for use as a JSON request body.""" body = {} + if self.binding_workspace_ids: + body["binding_workspace_ids"] = [v for v in self.binding_workspace_ids] if self.custom_tags: body["custom_tags"] = [v.as_dict() for v in self.custom_tags] if self.policy_id is not None: @@ -388,6 +392,8 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the BudgetPolicy into a shallow dictionary of its immediate attributes.""" body = {} + if self.binding_workspace_ids: + body["binding_workspace_ids"] = self.binding_workspace_ids if self.custom_tags: body["custom_tags"] = self.custom_tags if self.policy_id is not None: @@ -400,7 +406,8 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> BudgetPolicy: """Deserializes the BudgetPolicy from a dictionary.""" return cls( - custom_tags=_repeated_dict(d, "custom_tags", compute.CustomPolicyTag), + binding_workspace_ids=d.get("binding_workspace_ids", None), + custom_tags=_repeated_dict(d, "custom_tags", CustomPolicyTag), policy_id=d.get("policy_id", None), policy_name=d.get("policy_name", None), ) @@ -838,6 +845,46 @@ def from_dict(cls, d: Dict[str, Any]) -> CreateLogDeliveryConfigurationParams: ) +@dataclass +class CustomPolicyTag: + key: str + """The key of the tag. - Must be unique among all custom tags of the same policy - Cannot be + “budget-policy-name”, “budget-policy-id” or "budget-policy-resolution-result" - these + tags are preserved. + + - Follows the regex pattern defined in cluster-common/conf/src/ClusterTagConstraints.scala + (https://src.dev.databricks.com/databricks/universe@1647196627c8dc7b4152ad098a94b86484b93a6c/-/blob/cluster-common/conf/src/ClusterTagConstraints.scala?L17)""" + + value: Optional[str] = None + """The value of the tag. + + - Follows the regex pattern defined in cluster-common/conf/src/ClusterTagConstraints.scala + (https://src.dev.databricks.com/databricks/universe@1647196627c8dc7b4152ad098a94b86484b93a6c/-/blob/cluster-common/conf/src/ClusterTagConstraints.scala?L24)""" + + def as_dict(self) -> dict: + """Serializes the CustomPolicyTag into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.key is not None: + body["key"] = self.key + if self.value is not None: + body["value"] = self.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CustomPolicyTag into a shallow dictionary of its immediate attributes.""" + body = {} + if self.key is not None: + body["key"] = self.key + if self.value is not None: + body["value"] = self.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CustomPolicyTag: + """Deserializes the CustomPolicyTag from a dictionary.""" + return cls(key=d.get("key", None), value=d.get("value", None)) + + @dataclass class DeleteBudgetConfigurationResponse: def as_dict(self) -> dict: diff --git a/databricks/sdk/billing/v2/client.py b/databricks/sdk/billing/v2/client.py new file mode 100755 index 000000000..440feb981 --- /dev/null +++ b/databricks/sdk/billing/v2/client.py @@ -0,0 +1,396 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .billing import (BillableUsageAPI, BudgetPolicyAPI, BudgetsAPI, + LogDeliveryAPI, UsageDashboardsAPI) + +_LOG = logging.getLogger(__name__) + + +class BillableUsageClient(BillableUsageAPI): + """ + This API allows you to download billable usage logs for the specified account and date range. + This feature works with all account types. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class BudgetPolicyClient(BudgetPolicyAPI): + """ + A service serves REST API about Budget policies + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class BudgetsClient(BudgetsAPI): + """ + These APIs manage budget configurations for this account. Budgets enable you to monitor usage + across your account. You can set up budgets to either track account-wide spending, or apply + filters to track the spending of specific teams, projects, or workspaces. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class LogDeliveryClient(LogDeliveryAPI): + """ + These APIs manage log delivery configurations for this account. The two supported log types for + this API are _billable usage logs_ and _audit logs_. This feature is in Public Preview. This + feature works with all account ID types. + + Log delivery works with all account types. However, if your account is on the E2 version of the + platform or on a select custom plan that allows multiple workspaces per account, you can + optionally configure different storage destinations for each workspace. Log delivery status is + also provided to know the latest status of log delivery attempts. The high-level flow of + billable usage delivery: + + 1. **Create storage**: In AWS, [create a new AWS S3 bucket] with a specific bucket policy. Using + Databricks APIs, call the Account API to create a [storage configuration + object](:method:Storage/Create) that uses the bucket name. 2. **Create credentials**: In AWS, + create the appropriate AWS IAM role. For full details, including the required IAM role policies + and trust relationship, see [Billable usage log delivery]. Using Databricks APIs, call the + Account API to create a [credential configuration object](:method:Credentials/Create) that uses + the IAM role"s ARN. 3. **Create log delivery configuration**: Using Databricks APIs, call the + Account API to [create a log delivery configuration](:method:LogDelivery/Create) that uses the + credential and storage configuration objects from previous steps. You can specify if the logs + should include all events of that log type in your account (_Account level_ delivery) or only + events for a specific set of workspaces (_workspace level_ delivery). Account level log delivery + applies to all current and future workspaces plus account level logs, while workspace level log + delivery solely delivers logs related to the specified workspaces. You can create multiple types + of delivery configurations per account. + + For billable usage delivery: * For more information about billable usage logs, see [Billable + usage log delivery]. For the CSV schema, see the [Usage page]. * The delivery location is + `//billable-usage/csv/`, where `` is the name of the optional + delivery path prefix you set up during log delivery configuration. Files are named + `workspaceId=-usageMonth=.csv`. * All billable usage logs apply to specific + workspaces (_workspace level_ logs). You can aggregate usage for your entire account by creating + an _account level_ delivery configuration that delivers logs for all current and future + workspaces in your account. * The files are delivered daily by overwriting the month's CSV file + for each workspace. + + For audit log delivery: * For more information about about audit log delivery, see [Audit log + delivery], which includes information about the used JSON schema. * The delivery location is + `//workspaceId=/date=/auditlogs_.json`. + Files may get overwritten with the same content multiple times to achieve exactly-once delivery. + * If the audit log delivery configuration included specific workspace IDs, only + _workspace-level_ audit logs for those workspaces are delivered. If the log delivery + configuration applies to the entire account (_account level_ delivery configuration), the audit + log delivery includes workspace-level audit logs for all workspaces in the account as well as + account-level audit logs. See [Audit log delivery] for details. * Auditable events are typically + available in logs within 15 minutes. + + [Audit log delivery]: https://docs.databricks.com/administration-guide/account-settings/audit-logs.html + [Billable usage log delivery]: https://docs.databricks.com/administration-guide/account-settings/billable-usage-delivery.html + [Usage page]: https://docs.databricks.com/administration-guide/account-settings/usage.html + [create a new AWS S3 bucket]: https://docs.databricks.com/administration-guide/account-api/aws-storage.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class UsageDashboardsClient(UsageDashboardsAPI): + """ + These APIs manage usage dashboards for this account. Usage dashboards enable you to gain + insights into your usage with pre-built dashboards: visualize breakdowns, analyze tag + attributions, and identify cost drivers. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/catalog/__init__.py b/databricks/sdk/catalog/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/catalog/v2/__init__.py b/databricks/sdk/catalog/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/catalog/v2/catalog.py similarity index 99% rename from databricks/sdk/service/catalog.py rename to databricks/sdk/catalog/v2/catalog.py index 3124937d6..05c75356d 100755 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/catalog/v2/catalog.py @@ -3,19 +3,15 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -1065,6 +1061,7 @@ class CatalogType(Enum): """The type of the catalog.""" DELTASHARING_CATALOG = "DELTASHARING_CATALOG" + FOREIGN_CATALOG = "FOREIGN_CATALOG" MANAGED_CATALOG = "MANAGED_CATALOG" SYSTEM_CATALOG = "SYSTEM_CATALOG" @@ -1270,6 +1267,8 @@ class ColumnTypeName(Enum): DECIMAL = "DECIMAL" DOUBLE = "DOUBLE" FLOAT = "FLOAT" + GEOGRAPHY = "GEOGRAPHY" + GEOMETRY = "GEOMETRY" INT = "INT" INTERVAL = "INTERVAL" LONG = "LONG" @@ -2591,6 +2590,11 @@ class CreateVolumeRequestContent: """The name of the volume""" volume_type: VolumeType + """The type of the volume. An external volume is located in the specified external location. A + managed volume is located in the default location which is specified by the parent schema, or + the parent catalog, or the Metastore. [Learn more] + + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external""" comment: Optional[str] = None """The comment attached to the volume""" @@ -6664,6 +6668,7 @@ class Privilege(Enum): ACCESS = "ACCESS" ALL_PRIVILEGES = "ALL_PRIVILEGES" APPLY_TAG = "APPLY_TAG" + BROWSE = "BROWSE" CREATE = "CREATE" CREATE_CATALOG = "CREATE_CATALOG" CREATE_CONNECTION = "CREATE_CONNECTION" @@ -8132,38 +8137,6 @@ class TableType(Enum): VIEW = "VIEW" -@dataclass -class TagKeyValue: - key: Optional[str] = None - """name of the tag""" - - value: Optional[str] = None - """value of the tag associated with the key, could be optional""" - - def as_dict(self) -> dict: - """Serializes the TagKeyValue into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.key is not None: - body["key"] = self.key - if self.value is not None: - body["value"] = self.value - return body - - def as_shallow_dict(self) -> dict: - """Serializes the TagKeyValue into a shallow dictionary of its immediate attributes.""" - body = {} - if self.key is not None: - body["key"] = self.key - if self.value is not None: - body["value"] = self.value - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> TagKeyValue: - """Deserializes the TagKeyValue from a dictionary.""" - return cls(key=d.get("key", None), value=d.get("value", None)) - - @dataclass class TemporaryCredentials: aws_temp_credentials: Optional[AwsCredentials] = None @@ -9833,6 +9806,11 @@ class VolumeInfo: """The unique identifier of the volume""" volume_type: Optional[VolumeType] = None + """The type of the volume. An external volume is located in the specified external location. A + managed volume is located in the default location which is specified by the parent schema, or + the parent catalog, or the Metastore. [Learn more] + + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external""" def as_dict(self) -> dict: """Serializes the VolumeInfo into a dictionary suitable for use as a JSON request body.""" @@ -9937,6 +9915,11 @@ def from_dict(cls, d: Dict[str, Any]) -> VolumeInfo: class VolumeType(Enum): + """The type of the volume. An external volume is located in the specified external location. A + managed volume is located in the default location which is specified by the parent schema, or + the parent catalog, or the Metastore. [Learn more] + + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external""" EXTERNAL = "EXTERNAL" MANAGED = "MANAGED" @@ -12306,36 +12289,7 @@ class OnlineTablesAPI: def __init__(self, api_client): self._api = api_client - def wait_get_online_table_active( - self, name: str, timeout=timedelta(minutes=20), callback: Optional[Callable[[OnlineTable], None]] = None - ) -> OnlineTable: - deadline = time.time() + timeout.total_seconds() - target_states = (ProvisioningInfoState.ACTIVE,) - failure_states = (ProvisioningInfoState.FAILED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(name=name) - status = poll.unity_catalog_provisioning_state - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach ACTIVE, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"name={name}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def create(self, *, table: Optional[OnlineTable] = None) -> Wait[OnlineTable]: + def create(self, *, table: Optional[OnlineTable] = None) -> OnlineTable: """Create an Online Table. Create a new Online Table. @@ -12345,7 +12299,7 @@ def create(self, *, table: Optional[OnlineTable] = None) -> Wait[OnlineTable]: :returns: Long-running operation waiter for :class:`OnlineTable`. - See :method:wait_get_online_table_active for more details. + See :method:WaitGetOnlineTableActive for more details. """ body = table.as_dict() headers = { @@ -12353,13 +12307,8 @@ def create(self, *, table: Optional[OnlineTable] = None) -> Wait[OnlineTable]: "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/online-tables", body=body, headers=headers) - return Wait( - self.wait_get_online_table_active, response=OnlineTable.from_dict(op_response), name=op_response["name"] - ) - - def create_and_wait(self, *, table: Optional[OnlineTable] = None, timeout=timedelta(minutes=20)) -> OnlineTable: - return self.create(table=table).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/online-tables", body=body, headers=headers) + return OnlineTable.from_dict(res) def delete(self, name: str): """Delete an Online Table. @@ -14268,6 +14217,11 @@ def create( :param name: str The name of the volume :param volume_type: :class:`VolumeType` + The type of the volume. An external volume is located in the specified external location. A managed + volume is located in the default location which is specified by the parent schema, or the parent + catalog, or the Metastore. [Learn more] + + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external :param comment: str (optional) The comment attached to the volume :param storage_location: str (optional) diff --git a/databricks/sdk/catalog/v2/client.py b/databricks/sdk/catalog/v2/client.py new file mode 100755 index 000000000..a5994ecde --- /dev/null +++ b/databricks/sdk/catalog/v2/client.py @@ -0,0 +1,1764 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .catalog import (AccountMetastoreAssignmentsAPI, AccountMetastoresAPI, + AccountStorageCredentialsAPI, ArtifactAllowlistsAPI, + CatalogsAPI, ConnectionsAPI, CredentialsAPI, + ExternalLocationsAPI, FunctionsAPI, GrantsAPI, + MetastoresAPI, ModelVersionsAPI, OnlineTablesAPI, + QualityMonitorsAPI, RegisteredModelsAPI, + ResourceQuotasAPI, SchemasAPI, StorageCredentialsAPI, + SystemSchemasAPI, TableConstraintsAPI, TablesAPI, + TemporaryTableCredentialsAPI, VolumesAPI, + WorkspaceBindingsAPI) + +_LOG = logging.getLogger(__name__) + + +class AccountMetastoreAssignmentsClient(AccountMetastoreAssignmentsAPI): + """ + These APIs manage metastore assignments to a workspace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountMetastoresClient(AccountMetastoresAPI): + """ + These APIs manage Unity Catalog metastores for an account. A metastore contains catalogs that + can be associated with workspaces + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountStorageCredentialsClient(AccountStorageCredentialsAPI): + """ + These APIs manage storage credentials for a particular metastore. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ArtifactAllowlistsClient(ArtifactAllowlistsAPI): + """ + In Databricks Runtime 13.3 and above, you can add libraries and init scripts to the `allowlist` + in UC so that users can leverage these artifacts on compute configured with shared access mode. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CatalogsClient(CatalogsAPI): + """ + A catalog is the first layer of Unity Catalog’s three-level namespace. It’s used to organize + your data assets. Users can see all catalogs on which they have been assigned the USE_CATALOG + data permission. + + In Unity Catalog, admins and data stewards manage users and their access to data centrally + across all of the workspaces in a Databricks account. Users in different workspaces can share + access to the same data, depending on privileges granted centrally in Unity Catalog. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ConnectionsClient(ConnectionsAPI): + """ + Connections allow for creating a connection to an external data source. + + A connection is an abstraction of an external data source that can be connected from Databricks + Compute. Creating a connection object is the first step to managing external data sources within + Unity Catalog, with the second step being creating a data object (catalog, schema, or table) + using the connection. Data objects derived from a connection can be written to or read from + similar to other Unity Catalog data objects based on cloud storage. Users may create different + types of connections with each connection having a unique set of configuration options to + support credential management and other settings. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CredentialsClient(CredentialsAPI): + """ + A credential represents an authentication and authorization mechanism for accessing services on + your cloud tenant. Each credential is subject to Unity Catalog access-control policies that + control which users and groups can access the credential. + + To create credentials, you must be a Databricks account admin or have the `CREATE SERVICE + CREDENTIAL` privilege. The user who creates the credential can delegate ownership to another + user or group to manage permissions on it. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ExternalLocationsClient(ExternalLocationsAPI): + """ + An external location is an object that combines a cloud storage path with a storage credential + that authorizes access to the cloud storage path. Each external location is subject to Unity + Catalog access-control policies that control which users and groups can access the credential. + If a user does not have access to an external location in Unity Catalog, the request fails and + Unity Catalog does not attempt to authenticate to your cloud tenant on the user’s behalf. + + Databricks recommends using external locations rather than using storage credentials directly. + + To create external locations, you must be a metastore admin or a user with the + **CREATE_EXTERNAL_LOCATION** privilege. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class FunctionsClient(FunctionsAPI): + """ + Functions implement User-Defined Functions (UDFs) in Unity Catalog. + + The function implementation can be any SQL expression or Query, and it can be invoked wherever a + table reference is allowed in a query. In Unity Catalog, a function resides at the same level as + a table, so it can be referenced with the form + __catalog_name__.__schema_name__.__function_name__. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class GrantsClient(GrantsAPI): + """ + In Unity Catalog, data is secure by default. Initially, users have no access to data in a + metastore. Access can be granted by either a metastore admin, the owner of an object, or the + owner of the catalog or schema that contains the object. Securable objects in Unity Catalog are + hierarchical and privileges are inherited downward. + + Securable objects in Unity Catalog are hierarchical and privileges are inherited downward. This + means that granting a privilege on the catalog automatically grants the privilege to all current + and future objects within the catalog. Similarly, privileges granted on a schema are inherited + by all current and future objects within that schema. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class MetastoresClient(MetastoresAPI): + """ + A metastore is the top-level container of objects in Unity Catalog. It stores data assets + (tables and views) and the permissions that govern access to them. Databricks account admins can + create metastores and assign them to Databricks workspaces to control which workloads use each + metastore. For a workspace to use Unity Catalog, it must have a Unity Catalog metastore + attached. + + Each metastore is configured with a root storage location in a cloud storage account. This + storage location is used for metadata and managed tables data. + + NOTE: This metastore is distinct from the metastore included in Databricks workspaces created + before Unity Catalog was released. If your workspace includes a legacy Hive metastore, the data + in that metastore is available in a catalog named hive_metastore. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ModelVersionsClient(ModelVersionsAPI): + """ + Databricks provides a hosted version of MLflow Model Registry in Unity Catalog. Models in Unity + Catalog provide centralized access control, auditing, lineage, and discovery of ML models across + Databricks workspaces. + + This API reference documents the REST endpoints for managing model versions in Unity Catalog. + For more details, see the [registered models API docs](/api/workspace/registeredmodels). + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class OnlineTablesClient(OnlineTablesAPI): + """ + Online tables provide lower latency and higher QPS access to data from Delta tables. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QualityMonitorsClient(QualityMonitorsAPI): + """ + A monitor computes and monitors data or model quality metrics for a table over time. It + generates metrics tables and a dashboard that you can use to monitor table health and set + alerts. + + Most write operations require the user to be the owner of the table (or its parent schema or + parent catalog). Viewing the dashboard, computed metrics, or monitor configuration only requires + the user to have **SELECT** privileges on the table (along with **USE_SCHEMA** and + **USE_CATALOG**). + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class RegisteredModelsClient(RegisteredModelsAPI): + """ + Databricks provides a hosted version of MLflow Model Registry in Unity Catalog. Models in Unity + Catalog provide centralized access control, auditing, lineage, and discovery of ML models across + Databricks workspaces. + + An MLflow registered model resides in the third layer of Unity Catalog’s three-level + namespace. Registered models contain model versions, which correspond to actual ML models + (MLflow models). Creating new model versions currently requires use of the MLflow Python client. + Once model versions are created, you can load them for batch inference using MLflow Python + client APIs, or deploy them for real-time serving using Databricks Model Serving. + + All operations on registered models and model versions require USE_CATALOG permissions on the + enclosing catalog and USE_SCHEMA permissions on the enclosing schema. In addition, the following + additional privileges are required for various operations: + + * To create a registered model, users must additionally have the CREATE_MODEL permission on the + target schema. * To view registered model or model version metadata, model version data files, + or invoke a model version, users must additionally have the EXECUTE permission on the registered + model * To update registered model or model version tags, users must additionally have APPLY TAG + permissions on the registered model * To update other registered model or model version metadata + (comments, aliases) create a new model version, or update permissions on the registered model, + users must be owners of the registered model. + + Note: The securable type for models is "FUNCTION". When using REST APIs (e.g. tagging, grants) + that specify a securable type, use "FUNCTION" as the securable type. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ResourceQuotasClient(ResourceQuotasAPI): + """ + Unity Catalog enforces resource quotas on all securable objects, which limits the number of + resources that can be created. Quotas are expressed in terms of a resource type and a parent + (for example, tables per metastore or schemas per catalog). The resource quota APIs enable you + to monitor your current usage and limits. For more information on resource quotas see the [Unity + Catalog documentation]. + + [Unity Catalog documentation]: https://docs.databricks.com/en/data-governance/unity-catalog/index.html#resource-quotas + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class SchemasClient(SchemasAPI): + """ + A schema (also called a database) is the second layer of Unity Catalog’s three-level + namespace. A schema organizes tables, views and functions. To access (or list) a table or view + in a schema, users must have the USE_SCHEMA data permission on the schema and its parent + catalog, and they must have the SELECT permission on the table or view. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class StorageCredentialsClient(StorageCredentialsAPI): + """ + A storage credential represents an authentication and authorization mechanism for accessing data + stored on your cloud tenant. Each storage credential is subject to Unity Catalog access-control + policies that control which users and groups can access the credential. If a user does not have + access to a storage credential in Unity Catalog, the request fails and Unity Catalog does not + attempt to authenticate to your cloud tenant on the user’s behalf. + + Databricks recommends using external locations rather than using storage credentials directly. + + To create storage credentials, you must be a Databricks account admin. The account admin who + creates the storage credential can delegate ownership to another user or group to manage + permissions on it. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class SystemSchemasClient(SystemSchemasAPI): + """ + A system schema is a schema that lives within the system catalog. A system schema may contain + information about customer usage of Unity Catalog such as audit-logs, billing-logs, lineage + information, etc. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class TableConstraintsClient(TableConstraintsAPI): + """ + Primary key and foreign key constraints encode relationships between fields in tables. + + Primary and foreign keys are informational only and are not enforced. Foreign keys must + reference a primary key in another table. This primary key is the parent constraint of the + foreign key and the table this primary key is on is the parent table of the foreign key. + Similarly, the foreign key is the child constraint of its referenced primary key; the table of + the foreign key is the child table of the primary key. + + You can declare primary keys and foreign keys as part of the table specification during table + creation. You can also add or drop constraints on existing tables. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class TablesClient(TablesAPI): + """ + A table resides in the third layer of Unity Catalog’s three-level namespace. It contains rows + of data. To create a table, users must have CREATE_TABLE and USE_SCHEMA permissions on the + schema, and they must have the USE_CATALOG permission on its parent catalog. To query a table, + users must have the SELECT permission on the table, and they must have the USE_CATALOG + permission on its parent catalog and the USE_SCHEMA permission on its parent schema. + + A table can be managed or external. From an API perspective, a __VIEW__ is a particular kind of + table (rather than a managed or external table). + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class TemporaryTableCredentialsClient(TemporaryTableCredentialsAPI): + """ + Temporary Table Credentials refer to short-lived, downscoped credentials used to access cloud + storage locationswhere table data is stored in Databricks. These credentials are employed to + provide secure and time-limitedaccess to data in cloud environments such as AWS, Azure, and + Google Cloud. Each cloud provider has its own typeof credentials: AWS uses temporary session + tokens via AWS Security Token Service (STS), Azure utilizesShared Access Signatures (SAS) for + its data storage services, and Google Cloud supports temporary credentialsthrough OAuth + 2.0.Temporary table credentials ensure that data access is limited in scope and duration, + reducing the risk ofunauthorized access or misuse. To use the temporary table credentials API, a + metastore admin needs to enable the external_access_enabled flag (off by default) at the + metastore level, and user needs to be granted the EXTERNAL USE SCHEMA permission at the schema + level by catalog admin. Note that EXTERNAL USE SCHEMA is a schema level permission that can only + be granted by catalog admin explicitly and is not included in schema ownership or ALL PRIVILEGES + on the schema for security reason. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class VolumesClient(VolumesAPI): + """ + Volumes are a Unity Catalog (UC) capability for accessing, storing, governing, organizing and + processing files. Use cases include running machine learning on unstructured data such as image, + audio, video, or PDF files, organizing data sets during the data exploration stages in data + science, working with libraries that require access to the local file system on cluster + machines, storing library and config files of arbitrary formats such as .whl or .txt centrally + and providing secure access across workspaces to it, or transforming and querying non-tabular + data files in ETL. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WorkspaceBindingsClient(WorkspaceBindingsAPI): + """ + A securable in Databricks can be configured as __OPEN__ or __ISOLATED__. An __OPEN__ securable + can be accessed from any workspace, while an __ISOLATED__ securable can only be accessed from a + configured list of workspaces. This API allows you to configure (bind) securables to workspaces. + + NOTE: The __isolation_mode__ is configured for the securable itself (using its Update method) + and the workspace bindings are only consulted when the securable's __isolation_mode__ is set to + __ISOLATED__. + + A securable's workspace bindings can be configured by a metastore admin or the owner of the + securable. + + The original path (/api/2.1/unity-catalog/workspace-bindings/catalogs/{name}) is deprecated. + Please use the new path (/api/2.1/unity-catalog/bindings/{securable_type}/{securable_name}) + which introduces the ability to bind a securable in READ_ONLY mode (catalogs only). + + Securable types that support binding: - catalog - storage_credential - external_location + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/cleanrooms/__init__.py b/databricks/sdk/cleanrooms/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/cleanrooms/v2/__init__.py b/databricks/sdk/cleanrooms/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/service/cleanrooms.py b/databricks/sdk/cleanrooms/v2/cleanrooms.py similarity index 67% rename from databricks/sdk/service/cleanrooms.py rename to databricks/sdk/cleanrooms/v2/cleanrooms.py index e551d82be..2e4bff8b9 100755 --- a/databricks/sdk/service/cleanrooms.py +++ b/databricks/sdk/cleanrooms/v2/cleanrooms.py @@ -7,13 +7,11 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - -from databricks.sdk.service import catalog, jobs, settings, sharing - # all definitions in this file are in alphabetical order @@ -280,7 +278,7 @@ class CleanRoomAssetAssetType(Enum): @dataclass class CleanRoomAssetForeignTable: - columns: Optional[List[catalog.ColumnInfo]] = None + columns: Optional[List[ColumnInfo]] = None """The metadata information of the columns in the foreign table""" def as_dict(self) -> dict: @@ -300,7 +298,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> CleanRoomAssetForeignTable: """Deserializes the CleanRoomAssetForeignTable from a dictionary.""" - return cls(columns=_repeated_dict(d, "columns", catalog.ColumnInfo)) + return cls(columns=_repeated_dict(d, "columns", ColumnInfo)) @dataclass @@ -371,7 +369,7 @@ class CleanRoomAssetStatusEnum(Enum): @dataclass class CleanRoomAssetTable: - columns: Optional[List[catalog.ColumnInfo]] = None + columns: Optional[List[ColumnInfo]] = None """The metadata information of the columns in the table""" def as_dict(self) -> dict: @@ -391,7 +389,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> CleanRoomAssetTable: """Deserializes the CleanRoomAssetTable from a dictionary.""" - return cls(columns=_repeated_dict(d, "columns", catalog.ColumnInfo)) + return cls(columns=_repeated_dict(d, "columns", ColumnInfo)) @dataclass @@ -400,7 +398,7 @@ class CleanRoomAssetTableLocalDetails: """The fully qualified name of the table in its owner's local metastore, in the format of *catalog*.*schema*.*table_name*""" - partitions: Optional[List[sharing.Partition]] = None + partitions: Optional[List[Partition]] = None """Partition filtering specification for a shared table.""" def as_dict(self) -> dict: @@ -424,12 +422,12 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> CleanRoomAssetTableLocalDetails: """Deserializes the CleanRoomAssetTableLocalDetails from a dictionary.""" - return cls(local_name=d.get("local_name", None), partitions=_repeated_dict(d, "partitions", sharing.Partition)) + return cls(local_name=d.get("local_name", None), partitions=_repeated_dict(d, "partitions", Partition)) @dataclass class CleanRoomAssetView: - columns: Optional[List[catalog.ColumnInfo]] = None + columns: Optional[List[ColumnInfo]] = None """The metadata information of the columns in the view""" def as_dict(self) -> dict: @@ -449,7 +447,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> CleanRoomAssetView: """Deserializes the CleanRoomAssetView from a dictionary.""" - return cls(columns=_repeated_dict(d, "columns", catalog.ColumnInfo)) + return cls(columns=_repeated_dict(d, "columns", ColumnInfo)) @dataclass @@ -594,7 +592,7 @@ class CleanRoomNotebookTaskRun: LIST API. if the task was run within the same workspace the API is being called. If the task run was in a different workspace under the same metastore, only the workspace_id is included.""" - notebook_job_run_state: Optional[jobs.CleanRoomTaskRunState] = None + notebook_job_run_state: Optional[CleanRoomTaskRunState] = None """State of the task run.""" notebook_name: Optional[str] = None @@ -655,7 +653,7 @@ def from_dict(cls, d: Dict[str, Any]) -> CleanRoomNotebookTaskRun: """Deserializes the CleanRoomNotebookTaskRun from a dictionary.""" return cls( collaborator_job_run_info=_from_dict(d, "collaborator_job_run_info", CollaboratorJobRunInfo), - notebook_job_run_state=_from_dict(d, "notebook_job_run_state", jobs.CleanRoomTaskRunState), + notebook_job_run_state=_from_dict(d, "notebook_job_run_state", CleanRoomTaskRunState), notebook_name=d.get("notebook_name", None), output_schema_expiration_time=d.get("output_schema_expiration_time", None), output_schema_name=d.get("output_schema_name", None), @@ -732,7 +730,7 @@ class CleanRoomRemoteDetail: creator: Optional[CleanRoomCollaborator] = None """Collaborator who creates the clean room.""" - egress_network_policy: Optional[settings.EgressNetworkPolicy] = None + egress_network_policy: Optional[EgressNetworkPolicy] = None """Egress network policy to apply to the central clean room workspace.""" region: Optional[str] = None @@ -785,7 +783,7 @@ def from_dict(cls, d: Dict[str, Any]) -> CleanRoomRemoteDetail: collaborators=_repeated_dict(d, "collaborators", CleanRoomCollaborator), compliance_security_profile=_from_dict(d, "compliance_security_profile", ComplianceSecurityProfile), creator=_from_dict(d, "creator", CleanRoomCollaborator), - egress_network_policy=_from_dict(d, "egress_network_policy", settings.EgressNetworkPolicy), + egress_network_policy=_from_dict(d, "egress_network_policy", EgressNetworkPolicy), region=d.get("region", None), ) @@ -798,6 +796,79 @@ class CleanRoomStatusEnum(Enum): PROVISIONING = "PROVISIONING" +class CleanRoomTaskRunLifeCycleState(Enum): + """Copied from elastic-spark-common/api/messages/runs.proto. Using the original definition to + remove coupling with jobs API definition""" + + BLOCKED = "BLOCKED" + INTERNAL_ERROR = "INTERNAL_ERROR" + PENDING = "PENDING" + QUEUED = "QUEUED" + RUNNING = "RUNNING" + RUN_LIFE_CYCLE_STATE_UNSPECIFIED = "RUN_LIFE_CYCLE_STATE_UNSPECIFIED" + SKIPPED = "SKIPPED" + TERMINATED = "TERMINATED" + TERMINATING = "TERMINATING" + WAITING_FOR_RETRY = "WAITING_FOR_RETRY" + + +class CleanRoomTaskRunResultState(Enum): + """Copied from elastic-spark-common/api/messages/runs.proto. Using the original definition to avoid + cyclic dependency.""" + + CANCELED = "CANCELED" + DISABLED = "DISABLED" + EVICTED = "EVICTED" + EXCLUDED = "EXCLUDED" + FAILED = "FAILED" + MAXIMUM_CONCURRENT_RUNS_REACHED = "MAXIMUM_CONCURRENT_RUNS_REACHED" + RUN_RESULT_STATE_UNSPECIFIED = "RUN_RESULT_STATE_UNSPECIFIED" + SUCCESS = "SUCCESS" + SUCCESS_WITH_FAILURES = "SUCCESS_WITH_FAILURES" + TIMEDOUT = "TIMEDOUT" + UPSTREAM_CANCELED = "UPSTREAM_CANCELED" + UPSTREAM_EVICTED = "UPSTREAM_EVICTED" + UPSTREAM_FAILED = "UPSTREAM_FAILED" + + +@dataclass +class CleanRoomTaskRunState: + """Stores the run state of the clean rooms notebook task.""" + + life_cycle_state: Optional[CleanRoomTaskRunLifeCycleState] = None + """A value indicating the run's current lifecycle state. This field is always available in the + response.""" + + result_state: Optional[CleanRoomTaskRunResultState] = None + """A value indicating the run's result. This field is only available for terminal lifecycle states.""" + + def as_dict(self) -> dict: + """Serializes the CleanRoomTaskRunState into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.life_cycle_state is not None: + body["life_cycle_state"] = self.life_cycle_state.value + if self.result_state is not None: + body["result_state"] = self.result_state.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the CleanRoomTaskRunState into a shallow dictionary of its immediate attributes.""" + body = {} + if self.life_cycle_state is not None: + body["life_cycle_state"] = self.life_cycle_state + if self.result_state is not None: + body["result_state"] = self.result_state + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> CleanRoomTaskRunState: + """Deserializes the CleanRoomTaskRunState from a dictionary.""" + return cls( + life_cycle_state=_enum(d, "life_cycle_state", CleanRoomTaskRunLifeCycleState), + result_state=_enum(d, "result_state", CleanRoomTaskRunResultState), + ) + + @dataclass class CollaboratorJobRunInfo: collaborator_alias: Optional[str] = None @@ -857,11 +928,186 @@ def from_dict(cls, d: Dict[str, Any]) -> CollaboratorJobRunInfo: ) +@dataclass +class ColumnInfo: + comment: Optional[str] = None + """User-provided free-form text description.""" + + mask: Optional[ColumnMask] = None + + name: Optional[str] = None + """Name of Column.""" + + nullable: Optional[bool] = None + """Whether field may be Null (default: true).""" + + partition_index: Optional[int] = None + """Partition index for column.""" + + position: Optional[int] = None + """Ordinal position of column (starting at position 0).""" + + type_interval_type: Optional[str] = None + """Format of IntervalType.""" + + type_json: Optional[str] = None + """Full data type specification, JSON-serialized.""" + + type_name: Optional[ColumnTypeName] = None + + type_precision: Optional[int] = None + """Digits of precision; required for DecimalTypes.""" + + type_scale: Optional[int] = None + """Digits to right of decimal; Required for DecimalTypes.""" + + type_text: Optional[str] = None + """Full data type specification as SQL/catalogString text.""" + + def as_dict(self) -> dict: + """Serializes the ColumnInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.comment is not None: + body["comment"] = self.comment + if self.mask: + body["mask"] = self.mask.as_dict() + if self.name is not None: + body["name"] = self.name + if self.nullable is not None: + body["nullable"] = self.nullable + if self.partition_index is not None: + body["partition_index"] = self.partition_index + if self.position is not None: + body["position"] = self.position + if self.type_interval_type is not None: + body["type_interval_type"] = self.type_interval_type + if self.type_json is not None: + body["type_json"] = self.type_json + if self.type_name is not None: + body["type_name"] = self.type_name.value + if self.type_precision is not None: + body["type_precision"] = self.type_precision + if self.type_scale is not None: + body["type_scale"] = self.type_scale + if self.type_text is not None: + body["type_text"] = self.type_text + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ColumnInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.comment is not None: + body["comment"] = self.comment + if self.mask: + body["mask"] = self.mask + if self.name is not None: + body["name"] = self.name + if self.nullable is not None: + body["nullable"] = self.nullable + if self.partition_index is not None: + body["partition_index"] = self.partition_index + if self.position is not None: + body["position"] = self.position + if self.type_interval_type is not None: + body["type_interval_type"] = self.type_interval_type + if self.type_json is not None: + body["type_json"] = self.type_json + if self.type_name is not None: + body["type_name"] = self.type_name + if self.type_precision is not None: + body["type_precision"] = self.type_precision + if self.type_scale is not None: + body["type_scale"] = self.type_scale + if self.type_text is not None: + body["type_text"] = self.type_text + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ColumnInfo: + """Deserializes the ColumnInfo from a dictionary.""" + return cls( + comment=d.get("comment", None), + mask=_from_dict(d, "mask", ColumnMask), + name=d.get("name", None), + nullable=d.get("nullable", None), + partition_index=d.get("partition_index", None), + position=d.get("position", None), + type_interval_type=d.get("type_interval_type", None), + type_json=d.get("type_json", None), + type_name=_enum(d, "type_name", ColumnTypeName), + type_precision=d.get("type_precision", None), + type_scale=d.get("type_scale", None), + type_text=d.get("type_text", None), + ) + + +@dataclass +class ColumnMask: + function_name: Optional[str] = None + """The full name of the column mask SQL UDF.""" + + using_column_names: Optional[List[str]] = None + """The list of additional table columns to be passed as input to the column mask function. The + first arg of the mask function should be of the type of the column being masked and the types of + the rest of the args should match the types of columns in 'using_column_names'.""" + + def as_dict(self) -> dict: + """Serializes the ColumnMask into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.function_name is not None: + body["function_name"] = self.function_name + if self.using_column_names: + body["using_column_names"] = [v for v in self.using_column_names] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ColumnMask into a shallow dictionary of its immediate attributes.""" + body = {} + if self.function_name is not None: + body["function_name"] = self.function_name + if self.using_column_names: + body["using_column_names"] = self.using_column_names + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ColumnMask: + """Deserializes the ColumnMask from a dictionary.""" + return cls(function_name=d.get("function_name", None), using_column_names=d.get("using_column_names", None)) + + +class ColumnTypeName(Enum): + + ARRAY = "ARRAY" + BINARY = "BINARY" + BOOLEAN = "BOOLEAN" + BYTE = "BYTE" + CHAR = "CHAR" + DATE = "DATE" + DECIMAL = "DECIMAL" + DOUBLE = "DOUBLE" + FLOAT = "FLOAT" + GEOGRAPHY = "GEOGRAPHY" + GEOMETRY = "GEOMETRY" + INT = "INT" + INTERVAL = "INTERVAL" + LONG = "LONG" + MAP = "MAP" + NULL = "NULL" + SHORT = "SHORT" + STRING = "STRING" + STRUCT = "STRUCT" + TABLE_TYPE = "TABLE_TYPE" + TIMESTAMP = "TIMESTAMP" + TIMESTAMP_NTZ = "TIMESTAMP_NTZ" + USER_DEFINED_TYPE = "USER_DEFINED_TYPE" + VARIANT = "VARIANT" + + @dataclass class ComplianceSecurityProfile: """The compliance security profile used to process regulated data following compliance standards.""" - compliance_standards: Optional[List[settings.ComplianceStandard]] = None + compliance_standards: Optional[List[ComplianceStandard]] = None """The list of compliance standards that the compliance security profile is configured to enforce.""" is_enabled: Optional[bool] = None @@ -871,7 +1117,7 @@ def as_dict(self) -> dict: """Serializes the ComplianceSecurityProfile into a dictionary suitable for use as a JSON request body.""" body = {} if self.compliance_standards: - body["compliance_standards"] = [v.as_dict() for v in self.compliance_standards] + body["compliance_standards"] = [v.value for v in self.compliance_standards] if self.is_enabled is not None: body["is_enabled"] = self.is_enabled return body @@ -889,11 +1135,28 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> ComplianceSecurityProfile: """Deserializes the ComplianceSecurityProfile from a dictionary.""" return cls( - compliance_standards=_repeated_dict(d, "compliance_standards", settings.ComplianceStandard), + compliance_standards=_repeated_enum(d, "compliance_standards", ComplianceStandard), is_enabled=d.get("is_enabled", None), ) +class ComplianceStandard(Enum): + """Compliance stardard for SHIELD customers""" + + CANADA_PROTECTED_B = "CANADA_PROTECTED_B" + CYBER_ESSENTIAL_PLUS = "CYBER_ESSENTIAL_PLUS" + FEDRAMP_HIGH = "FEDRAMP_HIGH" + FEDRAMP_IL5 = "FEDRAMP_IL5" + FEDRAMP_MODERATE = "FEDRAMP_MODERATE" + HIPAA = "HIPAA" + HITRUST = "HITRUST" + IRAP_PROTECTED = "IRAP_PROTECTED" + ISMAP = "ISMAP" + ITAR_EAR = "ITAR_EAR" + NONE = "NONE" + PCI_DSS = "PCI_DSS" + + @dataclass class CreateCleanRoomOutputCatalogResponse: output_catalog: Optional[CleanRoomOutputCatalog] = None @@ -957,6 +1220,305 @@ def from_dict(cls, d: Dict[str, Any]) -> DeleteResponse: return cls() +@dataclass +class EgressNetworkPolicy: + """The network policies applying for egress traffic. This message is used by the UI/REST API. We + translate this message to the format expected by the dataplane in Lakehouse Network Manager (for + the format expected by the dataplane, see networkconfig.textproto).""" + + internet_access: Optional[EgressNetworkPolicyInternetAccessPolicy] = None + """The access policy enforced for egress traffic to the internet.""" + + def as_dict(self) -> dict: + """Serializes the EgressNetworkPolicy into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.internet_access: + body["internet_access"] = self.internet_access.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EgressNetworkPolicy into a shallow dictionary of its immediate attributes.""" + body = {} + if self.internet_access: + body["internet_access"] = self.internet_access + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicy: + """Deserializes the EgressNetworkPolicy from a dictionary.""" + return cls(internet_access=_from_dict(d, "internet_access", EgressNetworkPolicyInternetAccessPolicy)) + + +@dataclass +class EgressNetworkPolicyInternetAccessPolicy: + allowed_internet_destinations: Optional[List[EgressNetworkPolicyInternetAccessPolicyInternetDestination]] = None + + allowed_storage_destinations: Optional[List[EgressNetworkPolicyInternetAccessPolicyStorageDestination]] = None + + log_only_mode: Optional[EgressNetworkPolicyInternetAccessPolicyLogOnlyMode] = None + """Optional. If not specified, assume the policy is enforced for all workloads.""" + + restriction_mode: Optional[EgressNetworkPolicyInternetAccessPolicyRestrictionMode] = None + """At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: + Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can + only access explicitly allowed internet and storage destinations, as well as UC connections and + external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via + private link.""" + + def as_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicy into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.allowed_internet_destinations: + body["allowed_internet_destinations"] = [v.as_dict() for v in self.allowed_internet_destinations] + if self.allowed_storage_destinations: + body["allowed_storage_destinations"] = [v.as_dict() for v in self.allowed_storage_destinations] + if self.log_only_mode: + body["log_only_mode"] = self.log_only_mode.as_dict() + if self.restriction_mode is not None: + body["restriction_mode"] = self.restriction_mode.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicy into a shallow dictionary of its immediate attributes.""" + body = {} + if self.allowed_internet_destinations: + body["allowed_internet_destinations"] = self.allowed_internet_destinations + if self.allowed_storage_destinations: + body["allowed_storage_destinations"] = self.allowed_storage_destinations + if self.log_only_mode: + body["log_only_mode"] = self.log_only_mode + if self.restriction_mode is not None: + body["restriction_mode"] = self.restriction_mode + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicy: + """Deserializes the EgressNetworkPolicyInternetAccessPolicy from a dictionary.""" + return cls( + allowed_internet_destinations=_repeated_dict( + d, "allowed_internet_destinations", EgressNetworkPolicyInternetAccessPolicyInternetDestination + ), + allowed_storage_destinations=_repeated_dict( + d, "allowed_storage_destinations", EgressNetworkPolicyInternetAccessPolicyStorageDestination + ), + log_only_mode=_from_dict(d, "log_only_mode", EgressNetworkPolicyInternetAccessPolicyLogOnlyMode), + restriction_mode=_enum(d, "restriction_mode", EgressNetworkPolicyInternetAccessPolicyRestrictionMode), + ) + + +@dataclass +class EgressNetworkPolicyInternetAccessPolicyInternetDestination: + """Users can specify accessible internet destinations when outbound access is restricted. We only + support domain name (FQDN) destinations for the time being, though going forwards we want to + support host names and IP addresses.""" + + destination: Optional[str] = None + + protocol: Optional[ + EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol + ] = None + """The filtering protocol used by the DP. For private and public preview, SEG will only support TCP + filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be + set to TCP by default and hidden from the user. In the future, users may be able to select HTTP + filtering (i.e. SNI based filtering, filtering by FQDN).""" + + type: Optional[EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType] = None + + def as_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + if self.protocol is not None: + body["protocol"] = self.protocol.value + if self.type is not None: + body["type"] = self.type.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + if self.protocol is not None: + body["protocol"] = self.protocol + if self.type is not None: + body["type"] = self.type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyInternetDestination: + """Deserializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination from a dictionary.""" + return cls( + destination=d.get("destination", None), + protocol=_enum( + d, + "protocol", + EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol, + ), + type=_enum(d, "type", EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType), + ) + + +class EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol(Enum): + """The filtering protocol used by the DP. For private and public preview, SEG will only support TCP + filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be + set to TCP by default and hidden from the user. In the future, users may be able to select HTTP + filtering (i.e. SNI based filtering, filtering by FQDN).""" + + TCP = "TCP" + + +class EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType(Enum): + + FQDN = "FQDN" + + +@dataclass +class EgressNetworkPolicyInternetAccessPolicyLogOnlyMode: + log_only_mode_type: Optional[EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType] = None + + workloads: Optional[List[EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType]] = None + + def as_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.log_only_mode_type is not None: + body["log_only_mode_type"] = self.log_only_mode_type.value + if self.workloads: + body["workloads"] = [v.value for v in self.workloads] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode into a shallow dictionary of its immediate attributes.""" + body = {} + if self.log_only_mode_type is not None: + body["log_only_mode_type"] = self.log_only_mode_type + if self.workloads: + body["workloads"] = self.workloads + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyLogOnlyMode: + """Deserializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode from a dictionary.""" + return cls( + log_only_mode_type=_enum( + d, "log_only_mode_type", EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType + ), + workloads=_repeated_enum(d, "workloads", EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType), + ) + + +class EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType(Enum): + + ALL_SERVICES = "ALL_SERVICES" + SELECTED_SERVICES = "SELECTED_SERVICES" + + +class EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType(Enum): + """The values should match the list of workloads used in networkconfig.proto""" + + DBSQL = "DBSQL" + ML_SERVING = "ML_SERVING" + + +class EgressNetworkPolicyInternetAccessPolicyRestrictionMode(Enum): + """At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: + Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can + only access explicitly allowed internet and storage destinations, as well as UC connections and + external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via + private link.""" + + FULL_ACCESS = "FULL_ACCESS" + PRIVATE_ACCESS_ONLY = "PRIVATE_ACCESS_ONLY" + RESTRICTED_ACCESS = "RESTRICTED_ACCESS" + + +@dataclass +class EgressNetworkPolicyInternetAccessPolicyStorageDestination: + """Users can specify accessible storage destinations.""" + + allowed_paths: Optional[List[str]] = None + + azure_container: Optional[str] = None + + azure_dns_zone: Optional[str] = None + + azure_storage_account: Optional[str] = None + + azure_storage_service: Optional[str] = None + + bucket_name: Optional[str] = None + + region: Optional[str] = None + + type: Optional[EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType] = None + + def as_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.allowed_paths: + body["allowed_paths"] = [v for v in self.allowed_paths] + if self.azure_container is not None: + body["azure_container"] = self.azure_container + if self.azure_dns_zone is not None: + body["azure_dns_zone"] = self.azure_dns_zone + if self.azure_storage_account is not None: + body["azure_storage_account"] = self.azure_storage_account + if self.azure_storage_service is not None: + body["azure_storage_service"] = self.azure_storage_service + if self.bucket_name is not None: + body["bucket_name"] = self.bucket_name + if self.region is not None: + body["region"] = self.region + if self.type is not None: + body["type"] = self.type.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination into a shallow dictionary of its immediate attributes.""" + body = {} + if self.allowed_paths: + body["allowed_paths"] = self.allowed_paths + if self.azure_container is not None: + body["azure_container"] = self.azure_container + if self.azure_dns_zone is not None: + body["azure_dns_zone"] = self.azure_dns_zone + if self.azure_storage_account is not None: + body["azure_storage_account"] = self.azure_storage_account + if self.azure_storage_service is not None: + body["azure_storage_service"] = self.azure_storage_service + if self.bucket_name is not None: + body["bucket_name"] = self.bucket_name + if self.region is not None: + body["region"] = self.region + if self.type is not None: + body["type"] = self.type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyStorageDestination: + """Deserializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination from a dictionary.""" + return cls( + allowed_paths=d.get("allowed_paths", None), + azure_container=d.get("azure_container", None), + azure_dns_zone=d.get("azure_dns_zone", None), + azure_storage_account=d.get("azure_storage_account", None), + azure_storage_service=d.get("azure_storage_service", None), + bucket_name=d.get("bucket_name", None), + region=d.get("region", None), + type=_enum(d, "type", EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType), + ) + + +class EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType(Enum): + + AWS_S3 = "AWS_S3" + AZURE_STORAGE = "AZURE_STORAGE" + CLOUDFLARE_R2 = "CLOUDFLARE_R2" + GOOGLE_CLOUD_STORAGE = "GOOGLE_CLOUD_STORAGE" + + @dataclass class ListCleanRoomAssetsResponse: assets: Optional[List[CleanRoomAsset]] = None @@ -1059,6 +1621,90 @@ def from_dict(cls, d: Dict[str, Any]) -> ListCleanRoomsResponse: ) +@dataclass +class Partition: + values: Optional[List[PartitionValue]] = None + """An array of partition values.""" + + def as_dict(self) -> dict: + """Serializes the Partition into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.values: + body["values"] = [v.as_dict() for v in self.values] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Partition into a shallow dictionary of its immediate attributes.""" + body = {} + if self.values: + body["values"] = self.values + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Partition: + """Deserializes the Partition from a dictionary.""" + return cls(values=_repeated_dict(d, "values", PartitionValue)) + + +@dataclass +class PartitionValue: + name: Optional[str] = None + """The name of the partition column.""" + + op: Optional[PartitionValueOp] = None + """The operator to apply for the value.""" + + recipient_property_key: Optional[str] = None + """The key of a Delta Sharing recipient's property. For example "databricks-account-id". When this + field is set, field `value` can not be set.""" + + value: Optional[str] = None + """The value of the partition column. When this value is not set, it means `null` value. When this + field is set, field `recipient_property_key` can not be set.""" + + def as_dict(self) -> dict: + """Serializes the PartitionValue into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.op is not None: + body["op"] = self.op.value + if self.recipient_property_key is not None: + body["recipient_property_key"] = self.recipient_property_key + if self.value is not None: + body["value"] = self.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PartitionValue into a shallow dictionary of its immediate attributes.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.op is not None: + body["op"] = self.op + if self.recipient_property_key is not None: + body["recipient_property_key"] = self.recipient_property_key + if self.value is not None: + body["value"] = self.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PartitionValue: + """Deserializes the PartitionValue from a dictionary.""" + return cls( + name=d.get("name", None), + op=_enum(d, "op", PartitionValueOp), + recipient_property_key=d.get("recipient_property_key", None), + value=d.get("value", None), + ) + + +class PartitionValueOp(Enum): + + EQUAL = "EQUAL" + LIKE = "LIKE" + + @dataclass class UpdateCleanRoomRequest: clean_room: Optional[CleanRoom] = None diff --git a/databricks/sdk/cleanrooms/v2/client.py b/databricks/sdk/cleanrooms/v2/client.py new file mode 100755 index 000000000..477d58621 --- /dev/null +++ b/databricks/sdk/cleanrooms/v2/client.py @@ -0,0 +1,212 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .cleanrooms import CleanRoomAssetsAPI, CleanRoomsAPI, CleanRoomTaskRunsAPI + +_LOG = logging.getLogger(__name__) + + +class CleanRoomAssetsClient(CleanRoomAssetsAPI): + """ + Clean room assets are data and code objects — Tables, volumes, and notebooks that are shared + with the clean room. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CleanRoomTaskRunsClient(CleanRoomTaskRunsAPI): + """ + Clean room task runs are the executions of notebooks in a clean room. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CleanRoomsClient(CleanRoomsAPI): + """ + A clean room uses Delta Sharing and serverless compute to provide a secure and + privacy-protecting environment where multiple parties can work together on sensitive enterprise + data without direct access to each other’s data. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/compute/__init__.py b/databricks/sdk/compute/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/compute/v2/__init__.py b/databricks/sdk/compute/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/compute/v2/client.py b/databricks/sdk/compute/v2/client.py new file mode 100755 index 000000000..7f167e6cc --- /dev/null +++ b/databricks/sdk/compute/v2/client.py @@ -0,0 +1,703 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .compute import (ClusterPoliciesAPI, CommandExecutionAPI, + GlobalInitScriptsAPI, InstancePoolsAPI, + InstanceProfilesAPI, LibrariesAPI, + PolicyComplianceForClustersAPI, PolicyFamiliesAPI) +from .mixin import ClustersExt + +_LOG = logging.getLogger(__name__) + + +class ClusterPoliciesClient(ClusterPoliciesAPI): + """ + You can use cluster policies to control users' ability to configure clusters based on a set of + rules. These rules specify which attributes or attribute values can be used during cluster + creation. Cluster policies have ACLs that limit their use to specific users and groups. + + With cluster policies, you can: - Auto-install cluster libraries on the next restart by listing + them in the policy's "libraries" field (Public Preview). - Limit users to creating clusters with + the prescribed settings. - Simplify the user interface, enabling more users to create clusters, + by fixing and hiding some fields. - Manage costs by setting limits on attributes that impact the + hourly rate. + + Cluster policy permissions limit which policies a user can select in the Policy drop-down when + the user creates a cluster: - A user who has unrestricted cluster create permission can select + the Unrestricted policy and create fully-configurable clusters. - A user who has both + unrestricted cluster create permission and access to cluster policies can select the + Unrestricted policy and policies they have access to. - A user that has access to only cluster + policies, can select the policies they have access to. + + If no policies exist in the workspace, the Policy drop-down doesn't appear. Only admin users can + create, edit, and delete policies. Admin users also have access to all policies. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ClustersClient(ClustersExt): + """ + The Clusters API allows you to create, start, edit, list, terminate, and delete clusters. + + Databricks maps cluster node instance types to compute units known as DBUs. See the instance + type pricing page for a list of the supported instance types and their corresponding DBUs. + + A Databricks cluster is a set of computation resources and configurations on which you run data + engineering, data science, and data analytics workloads, such as production ETL pipelines, + streaming analytics, ad-hoc analytics, and machine learning. + + You run these workloads as a set of commands in a notebook or as an automated job. Databricks + makes a distinction between all-purpose clusters and job clusters. You use all-purpose clusters + to analyze data collaboratively using interactive notebooks. You use job clusters to run fast + and robust automated jobs. + + You can create an all-purpose cluster using the UI, CLI, or REST API. You can manually terminate + and restart an all-purpose cluster. Multiple users can share such clusters to do collaborative + interactive analysis. + + IMPORTANT: Databricks retains cluster configuration information for terminated clusters for 30 + days. To keep an all-purpose cluster configuration even after it has been terminated for more + than 30 days, an administrator can pin a cluster to the cluster list. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CommandExecutionClient(CommandExecutionAPI): + """ + This API allows execution of Python, Scala, SQL, or R commands on running Databricks Clusters. + This API only supports (classic) all-purpose clusters. Serverless compute is not supported. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class GlobalInitScriptsClient(GlobalInitScriptsAPI): + """ + The Global Init Scripts API enables Workspace administrators to configure global initialization + scripts for their workspace. These scripts run on every node in every cluster in the workspace. + + **Important:** Existing clusters must be restarted to pick up any changes made to global init + scripts. Global init scripts are run in order. If the init script returns with a bad exit code, + the Apache Spark container fails to launch and init scripts with later position are skipped. If + enough containers fail, the entire cluster fails with a `GLOBAL_INIT_SCRIPT_FAILURE` error code. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class InstancePoolsClient(InstancePoolsAPI): + """ + Instance Pools API are used to create, edit, delete and list instance pools by using + ready-to-use cloud instances which reduces a cluster start and auto-scaling times. + + Databricks pools reduce cluster start and auto-scaling times by maintaining a set of idle, + ready-to-use instances. When a cluster is attached to a pool, cluster nodes are created using + the pool’s idle instances. If the pool has no idle instances, the pool expands by allocating a + new instance from the instance provider in order to accommodate the cluster’s request. When a + cluster releases an instance, it returns to the pool and is free for another cluster to use. + Only clusters attached to a pool can use that pool’s idle instances. + + You can specify a different pool for the driver node and worker nodes, or use the same pool for + both. + + Databricks does not charge DBUs while instances are idle in the pool. Instance provider billing + does apply. See pricing. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class InstanceProfilesClient(InstanceProfilesAPI): + """ + The Instance Profiles API allows admins to add, list, and remove instance profiles that users + can launch clusters with. Regular users can list the instance profiles available to them. See + [Secure access to S3 buckets] using instance profiles for more information. + + [Secure access to S3 buckets]: https://docs.databricks.com/administration-guide/cloud-configurations/aws/instance-profiles.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class LibrariesClient(LibrariesAPI): + """ + The Libraries API allows you to install and uninstall libraries and get the status of libraries + on a cluster. + + To make third-party or custom code available to notebooks and jobs running on your clusters, you + can install a library. Libraries can be written in Python, Java, Scala, and R. You can upload + Python, Java, Scala and R libraries and point to external packages in PyPI, Maven, and CRAN + repositories. + + Cluster libraries can be used by all notebooks running on a cluster. You can install a cluster + library directly from a public repository such as PyPI or Maven, using a previously installed + workspace library, or using an init script. + + When you uninstall a library from a cluster, the library is removed only when you restart the + cluster. Until you restart the cluster, the status of the uninstalled library appears as + Uninstall pending restart. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PolicyComplianceForClustersClient(PolicyComplianceForClustersAPI): + """ + The policy compliance APIs allow you to view and manage the policy compliance status of clusters + in your workspace. + + A cluster is compliant with its policy if its configuration satisfies all its policy rules. + Clusters could be out of compliance if their policy was updated after the cluster was last + edited. + + The get and list compliance APIs allow you to view the policy compliance status of a cluster. + The enforce compliance API allows you to update a cluster to be compliant with the current + version of its policy. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PolicyFamiliesClient(PolicyFamiliesAPI): + """ + View available policy families. A policy family contains a policy definition providing best + practices for configuring clusters for a particular use case. + + Databricks manages and provides policy families for several common cluster use cases. You cannot + create, edit, or delete policy families. + + Policy families cannot be used directly to create clusters. Instead, you create cluster policies + using a policy family. Cluster policies created using a policy family inherit the policy + family's policy definition. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/compute/v2/compute.py similarity index 95% rename from databricks/sdk/service/compute.py rename to databricks/sdk/compute/v2/compute.py index 49014519b..54b3da663 100755 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/compute/v2/compute.py @@ -3,19 +3,15 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -3519,46 +3515,6 @@ def from_dict(cls, d: Dict[str, Any]) -> Created: return cls(id=d.get("id", None)) -@dataclass -class CustomPolicyTag: - key: str - """The key of the tag. - Must be unique among all custom tags of the same policy - Cannot be - “budget-policy-name”, “budget-policy-id” or "budget-policy-resolution-result" - these - tags are preserved. - - - Follows the regex pattern defined in cluster-common/conf/src/ClusterTagConstraints.scala - (https://src.dev.databricks.com/databricks/universe@1647196627c8dc7b4152ad098a94b86484b93a6c/-/blob/cluster-common/conf/src/ClusterTagConstraints.scala?L17)""" - - value: Optional[str] = None - """The value of the tag. - - - Follows the regex pattern defined in cluster-common/conf/src/ClusterTagConstraints.scala - (https://src.dev.databricks.com/databricks/universe@1647196627c8dc7b4152ad098a94b86484b93a6c/-/blob/cluster-common/conf/src/ClusterTagConstraints.scala?L24)""" - - def as_dict(self) -> dict: - """Serializes the CustomPolicyTag into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.key is not None: - body["key"] = self.key - if self.value is not None: - body["value"] = self.value - return body - - def as_shallow_dict(self) -> dict: - """Serializes the CustomPolicyTag into a shallow dictionary of its immediate attributes.""" - body = {} - if self.key is not None: - body["key"] = self.key - if self.value is not None: - body["value"] = self.value - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> CustomPolicyTag: - """Deserializes the CustomPolicyTag from a dictionary.""" - return cls(key=d.get("key", None), value=d.get("value", None)) - - @dataclass class DataPlaneEventDetails: event_type: Optional[DataPlaneEventDetailsEventType] = None @@ -4770,48 +4726,6 @@ def from_dict(cls, d: Dict[str, Any]) -> EnforceClusterComplianceResponse: return cls(changes=_repeated_dict(d, "changes", ClusterSettingsChange), has_changes=d.get("has_changes", None)) -@dataclass -class Environment: - """The environment entity used to preserve serverless environment side panel and jobs' environment - for non-notebook task. In this minimal environment spec, only pip dependencies are supported.""" - - client: str - """Client version used by the environment The client is the user-facing environment of the runtime. - Each client comes with a specific set of pre-installed libraries. The version is a string, - consisting of the major client version.""" - - dependencies: Optional[List[str]] = None - """List of pip dependencies, as supported by the version of pip in this environment. Each - dependency is a pip requirement file line - https://pip.pypa.io/en/stable/reference/requirements-file-format/ Allowed dependency could be - , , (WSFS or Volumes in - Databricks), E.g. dependencies: ["foo==0.0.1", "-r - /Workspace/test/requirements.txt"]""" - - def as_dict(self) -> dict: - """Serializes the Environment into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.client is not None: - body["client"] = self.client - if self.dependencies: - body["dependencies"] = [v for v in self.dependencies] - return body - - def as_shallow_dict(self) -> dict: - """Serializes the Environment into a shallow dictionary of its immediate attributes.""" - body = {} - if self.client is not None: - body["client"] = self.client - if self.dependencies: - body["dependencies"] = self.dependencies - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> Environment: - """Deserializes the Environment from a dictionary.""" - return cls(client=d.get("client", None), dependencies=d.get("dependencies", None)) - - @dataclass class EventDetails: attributes: Optional[ClusterAttributes] = None @@ -8937,6 +8851,7 @@ class TerminationReasonCode(Enum): ACCESS_TOKEN_FAILURE = "ACCESS_TOKEN_FAILURE" ALLOCATION_TIMEOUT = "ALLOCATION_TIMEOUT" ALLOCATION_TIMEOUT_NODE_DAEMON_NOT_READY = "ALLOCATION_TIMEOUT_NODE_DAEMON_NOT_READY" + ALLOCATION_TIMEOUT_NO_HEALTHY_AND_WARMED_UP_CLUSTERS = "ALLOCATION_TIMEOUT_NO_HEALTHY_AND_WARMED_UP_CLUSTERS" ALLOCATION_TIMEOUT_NO_HEALTHY_CLUSTERS = "ALLOCATION_TIMEOUT_NO_HEALTHY_CLUSTERS" ALLOCATION_TIMEOUT_NO_MATCHED_CLUSTERS = "ALLOCATION_TIMEOUT_NO_MATCHED_CLUSTERS" ALLOCATION_TIMEOUT_NO_READY_CLUSTERS = "ALLOCATION_TIMEOUT_NO_READY_CLUSTERS" @@ -8989,7 +8904,10 @@ class TerminationReasonCode(Enum): DATA_ACCESS_CONFIG_CHANGED = "DATA_ACCESS_CONFIG_CHANGED" DBFS_COMPONENT_UNHEALTHY = "DBFS_COMPONENT_UNHEALTHY" DISASTER_RECOVERY_REPLICATION = "DISASTER_RECOVERY_REPLICATION" + DOCKER_CONTAINER_CREATION_EXCEPTION = "DOCKER_CONTAINER_CREATION_EXCEPTION" DOCKER_IMAGE_PULL_FAILURE = "DOCKER_IMAGE_PULL_FAILURE" + DOCKER_IMAGE_TOO_LARGE_FOR_INSTANCE_EXCEPTION = "DOCKER_IMAGE_TOO_LARGE_FOR_INSTANCE_EXCEPTION" + DOCKER_INVALID_OS_EXCEPTION = "DOCKER_INVALID_OS_EXCEPTION" DRIVER_EVICTION = "DRIVER_EVICTION" DRIVER_LAUNCH_TIMEOUT = "DRIVER_LAUNCH_TIMEOUT" DRIVER_NODE_UNREACHABLE = "DRIVER_NODE_UNREACHABLE" @@ -10074,73 +9992,6 @@ class ClustersAPI: def __init__(self, api_client): self._api = api_client - def wait_get_cluster_running( - self, - cluster_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[ClusterDetails], None]] = None, - ) -> ClusterDetails: - deadline = time.time() + timeout.total_seconds() - target_states = (State.RUNNING,) - failure_states = ( - State.ERROR, - State.TERMINATED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(cluster_id=cluster_id) - status = poll.state - status_message = poll.state_message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach RUNNING, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"cluster_id={cluster_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_get_cluster_terminated( - self, - cluster_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[ClusterDetails], None]] = None, - ) -> ClusterDetails: - deadline = time.time() + timeout.total_seconds() - target_states = (State.TERMINATED,) - failure_states = (State.ERROR,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(cluster_id=cluster_id) - status = poll.state - status_message = poll.state_message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach TERMINATED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"cluster_id={cluster_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def change_owner(self, cluster_id: str, owner_username: str): """Change cluster owner. @@ -10200,7 +10051,7 @@ def create( ssh_public_keys: Optional[List[str]] = None, use_ml_runtime: Optional[bool] = None, workload_type: Optional[WorkloadType] = None, - ) -> Wait[ClusterDetails]: + ) -> CreateClusterResponse: """Create new cluster. Creates a new Spark cluster. This method will acquire new instances from the cloud provider if @@ -10382,7 +10233,7 @@ def create( :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if apply_policy_default_values is not None: @@ -10452,84 +10303,10 @@ def create( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/create", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, - response=CreateClusterResponse.from_dict(op_response), - cluster_id=op_response["cluster_id"], - ) + res = self._api.do("POST", "/api/2.1/clusters/create", body=body, headers=headers) + return CreateClusterResponse.from_dict(res) - def create_and_wait( - self, - spark_version: str, - *, - apply_policy_default_values: Optional[bool] = None, - autoscale: Optional[AutoScale] = None, - autotermination_minutes: Optional[int] = None, - aws_attributes: Optional[AwsAttributes] = None, - azure_attributes: Optional[AzureAttributes] = None, - clone_from: Optional[CloneCluster] = None, - cluster_log_conf: Optional[ClusterLogConf] = None, - cluster_name: Optional[str] = None, - custom_tags: Optional[Dict[str, str]] = None, - data_security_mode: Optional[DataSecurityMode] = None, - docker_image: Optional[DockerImage] = None, - driver_instance_pool_id: Optional[str] = None, - driver_node_type_id: Optional[str] = None, - enable_elastic_disk: Optional[bool] = None, - enable_local_disk_encryption: Optional[bool] = None, - gcp_attributes: Optional[GcpAttributes] = None, - init_scripts: Optional[List[InitScriptInfo]] = None, - instance_pool_id: Optional[str] = None, - is_single_node: Optional[bool] = None, - kind: Optional[Kind] = None, - node_type_id: Optional[str] = None, - num_workers: Optional[int] = None, - policy_id: Optional[str] = None, - runtime_engine: Optional[RuntimeEngine] = None, - single_user_name: Optional[str] = None, - spark_conf: Optional[Dict[str, str]] = None, - spark_env_vars: Optional[Dict[str, str]] = None, - ssh_public_keys: Optional[List[str]] = None, - use_ml_runtime: Optional[bool] = None, - workload_type: Optional[WorkloadType] = None, - timeout=timedelta(minutes=20), - ) -> ClusterDetails: - return self.create( - apply_policy_default_values=apply_policy_default_values, - autoscale=autoscale, - autotermination_minutes=autotermination_minutes, - aws_attributes=aws_attributes, - azure_attributes=azure_attributes, - clone_from=clone_from, - cluster_log_conf=cluster_log_conf, - cluster_name=cluster_name, - custom_tags=custom_tags, - data_security_mode=data_security_mode, - docker_image=docker_image, - driver_instance_pool_id=driver_instance_pool_id, - driver_node_type_id=driver_node_type_id, - enable_elastic_disk=enable_elastic_disk, - enable_local_disk_encryption=enable_local_disk_encryption, - gcp_attributes=gcp_attributes, - init_scripts=init_scripts, - instance_pool_id=instance_pool_id, - is_single_node=is_single_node, - kind=kind, - node_type_id=node_type_id, - num_workers=num_workers, - policy_id=policy_id, - runtime_engine=runtime_engine, - single_user_name=single_user_name, - spark_conf=spark_conf, - spark_env_vars=spark_env_vars, - spark_version=spark_version, - ssh_public_keys=ssh_public_keys, - use_ml_runtime=use_ml_runtime, - workload_type=workload_type, - ).result(timeout=timeout) - - def delete(self, cluster_id: str) -> Wait[ClusterDetails]: + def delete(self, cluster_id: str): """Terminate cluster. Terminates the Spark cluster with the specified ID. The cluster is removed asynchronously. Once the @@ -10541,7 +10318,7 @@ def delete(self, cluster_id: str) -> Wait[ClusterDetails]: :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_terminated for more details. + See :method:WaitGetClusterTerminated for more details. """ body = {} if cluster_id is not None: @@ -10551,15 +10328,7 @@ def delete(self, cluster_id: str) -> Wait[ClusterDetails]: "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/delete", body=body, headers=headers) - return Wait( - self.wait_get_cluster_terminated, - response=DeleteClusterResponse.from_dict(op_response), - cluster_id=cluster_id, - ) - - def delete_and_wait(self, cluster_id: str, timeout=timedelta(minutes=20)) -> ClusterDetails: - return self.delete(cluster_id=cluster_id).result(timeout=timeout) + self._api.do("POST", "/api/2.1/clusters/delete", body=body, headers=headers) def edit( self, @@ -10595,7 +10364,7 @@ def edit( ssh_public_keys: Optional[List[str]] = None, use_ml_runtime: Optional[bool] = None, workload_type: Optional[WorkloadType] = None, - ) -> Wait[ClusterDetails]: + ): """Update cluster configuration. Updates the configuration of a cluster to match the provided attributes and size. A cluster can be @@ -10774,7 +10543,7 @@ def edit( :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if apply_policy_default_values is not None: @@ -10844,80 +10613,7 @@ def edit( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/edit", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, response=EditClusterResponse.from_dict(op_response), cluster_id=cluster_id - ) - - def edit_and_wait( - self, - cluster_id: str, - spark_version: str, - *, - apply_policy_default_values: Optional[bool] = None, - autoscale: Optional[AutoScale] = None, - autotermination_minutes: Optional[int] = None, - aws_attributes: Optional[AwsAttributes] = None, - azure_attributes: Optional[AzureAttributes] = None, - cluster_log_conf: Optional[ClusterLogConf] = None, - cluster_name: Optional[str] = None, - custom_tags: Optional[Dict[str, str]] = None, - data_security_mode: Optional[DataSecurityMode] = None, - docker_image: Optional[DockerImage] = None, - driver_instance_pool_id: Optional[str] = None, - driver_node_type_id: Optional[str] = None, - enable_elastic_disk: Optional[bool] = None, - enable_local_disk_encryption: Optional[bool] = None, - gcp_attributes: Optional[GcpAttributes] = None, - init_scripts: Optional[List[InitScriptInfo]] = None, - instance_pool_id: Optional[str] = None, - is_single_node: Optional[bool] = None, - kind: Optional[Kind] = None, - node_type_id: Optional[str] = None, - num_workers: Optional[int] = None, - policy_id: Optional[str] = None, - runtime_engine: Optional[RuntimeEngine] = None, - single_user_name: Optional[str] = None, - spark_conf: Optional[Dict[str, str]] = None, - spark_env_vars: Optional[Dict[str, str]] = None, - ssh_public_keys: Optional[List[str]] = None, - use_ml_runtime: Optional[bool] = None, - workload_type: Optional[WorkloadType] = None, - timeout=timedelta(minutes=20), - ) -> ClusterDetails: - return self.edit( - apply_policy_default_values=apply_policy_default_values, - autoscale=autoscale, - autotermination_minutes=autotermination_minutes, - aws_attributes=aws_attributes, - azure_attributes=azure_attributes, - cluster_id=cluster_id, - cluster_log_conf=cluster_log_conf, - cluster_name=cluster_name, - custom_tags=custom_tags, - data_security_mode=data_security_mode, - docker_image=docker_image, - driver_instance_pool_id=driver_instance_pool_id, - driver_node_type_id=driver_node_type_id, - enable_elastic_disk=enable_elastic_disk, - enable_local_disk_encryption=enable_local_disk_encryption, - gcp_attributes=gcp_attributes, - init_scripts=init_scripts, - instance_pool_id=instance_pool_id, - is_single_node=is_single_node, - kind=kind, - node_type_id=node_type_id, - num_workers=num_workers, - policy_id=policy_id, - runtime_engine=runtime_engine, - single_user_name=single_user_name, - spark_conf=spark_conf, - spark_env_vars=spark_env_vars, - spark_version=spark_version, - ssh_public_keys=ssh_public_keys, - use_ml_runtime=use_ml_runtime, - workload_type=workload_type, - ).result(timeout=timeout) + self._api.do("POST", "/api/2.1/clusters/edit", body=body, headers=headers) def events( self, @@ -11165,9 +10861,7 @@ def pin(self, cluster_id: str): self._api.do("POST", "/api/2.1/clusters/pin", body=body, headers=headers) - def resize( - self, cluster_id: str, *, autoscale: Optional[AutoScale] = None, num_workers: Optional[int] = None - ) -> Wait[ClusterDetails]: + def resize(self, cluster_id: str, *, autoscale: Optional[AutoScale] = None, num_workers: Optional[int] = None): """Resize cluster. Resizes a cluster to have a desired number of workers. This will fail unless the cluster is in a @@ -11190,7 +10884,7 @@ def resize( :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if autoscale is not None: @@ -11204,22 +10898,9 @@ def resize( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/resize", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, response=ResizeClusterResponse.from_dict(op_response), cluster_id=cluster_id - ) + self._api.do("POST", "/api/2.1/clusters/resize", body=body, headers=headers) - def resize_and_wait( - self, - cluster_id: str, - *, - autoscale: Optional[AutoScale] = None, - num_workers: Optional[int] = None, - timeout=timedelta(minutes=20), - ) -> ClusterDetails: - return self.resize(autoscale=autoscale, cluster_id=cluster_id, num_workers=num_workers).result(timeout=timeout) - - def restart(self, cluster_id: str, *, restart_user: Optional[str] = None) -> Wait[ClusterDetails]: + def restart(self, cluster_id: str, *, restart_user: Optional[str] = None): """Restart cluster. Restarts a Spark cluster with the supplied ID. If the cluster is not currently in a `RUNNING` state, @@ -11231,7 +10912,7 @@ def restart(self, cluster_id: str, *, restart_user: Optional[str] = None) -> Wai :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if cluster_id is not None: @@ -11243,15 +10924,7 @@ def restart(self, cluster_id: str, *, restart_user: Optional[str] = None) -> Wai "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/restart", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, response=RestartClusterResponse.from_dict(op_response), cluster_id=cluster_id - ) - - def restart_and_wait( - self, cluster_id: str, *, restart_user: Optional[str] = None, timeout=timedelta(minutes=20) - ) -> ClusterDetails: - return self.restart(cluster_id=cluster_id, restart_user=restart_user).result(timeout=timeout) + self._api.do("POST", "/api/2.1/clusters/restart", body=body, headers=headers) def set_permissions( self, cluster_id: str, *, access_control_list: Optional[List[ClusterAccessControlRequest]] = None @@ -11293,7 +10966,7 @@ def spark_versions(self) -> GetSparkVersionsResponse: res = self._api.do("GET", "/api/2.1/clusters/spark-versions", headers=headers) return GetSparkVersionsResponse.from_dict(res) - def start(self, cluster_id: str) -> Wait[ClusterDetails]: + def start(self, cluster_id: str): """Start terminated cluster. Starts a terminated Spark cluster with the supplied ID. This works similar to `createCluster` except: @@ -11307,7 +10980,7 @@ def start(self, cluster_id: str) -> Wait[ClusterDetails]: :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if cluster_id is not None: @@ -11317,13 +10990,7 @@ def start(self, cluster_id: str) -> Wait[ClusterDetails]: "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/start", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, response=StartClusterResponse.from_dict(op_response), cluster_id=cluster_id - ) - - def start_and_wait(self, cluster_id: str, timeout=timedelta(minutes=20)) -> ClusterDetails: - return self.start(cluster_id=cluster_id).result(timeout=timeout) + self._api.do("POST", "/api/2.1/clusters/start", body=body, headers=headers) def unpin(self, cluster_id: str): """Unpin cluster. @@ -11346,9 +11013,7 @@ def unpin(self, cluster_id: str): self._api.do("POST", "/api/2.1/clusters/unpin", body=body, headers=headers) - def update( - self, cluster_id: str, update_mask: str, *, cluster: Optional[UpdateClusterResource] = None - ) -> Wait[ClusterDetails]: + def update(self, cluster_id: str, update_mask: str, *, cluster: Optional[UpdateClusterResource] = None): """Update cluster configuration (partial). Updates the configuration of a cluster to match the partial set of attributes and size. Denote which @@ -11380,7 +11045,7 @@ def update( :returns: Long-running operation waiter for :class:`ClusterDetails`. - See :method:wait_get_cluster_running for more details. + See :method:WaitGetClusterRunning for more details. """ body = {} if cluster is not None: @@ -11394,20 +11059,7 @@ def update( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.1/clusters/update", body=body, headers=headers) - return Wait( - self.wait_get_cluster_running, response=UpdateClusterResponse.from_dict(op_response), cluster_id=cluster_id - ) - - def update_and_wait( - self, - cluster_id: str, - update_mask: str, - *, - cluster: Optional[UpdateClusterResource] = None, - timeout=timedelta(minutes=20), - ) -> ClusterDetails: - return self.update(cluster=cluster, cluster_id=cluster_id, update_mask=update_mask).result(timeout=timeout) + self._api.do("POST", "/api/2.1/clusters/update", body=body, headers=headers) def update_permissions( self, cluster_id: str, *, access_control_list: Optional[List[ClusterAccessControlRequest]] = None @@ -11441,118 +11093,9 @@ class CommandExecutionAPI: def __init__(self, api_client): self._api = api_client - def wait_command_status_command_execution_cancelled( - self, - cluster_id: str, - command_id: str, - context_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[CommandStatusResponse], None]] = None, - ) -> CommandStatusResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (CommandStatus.CANCELLED,) - failure_states = (CommandStatus.ERROR,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.command_status(cluster_id=cluster_id, command_id=command_id, context_id=context_id) - status = poll.status - status_message = f"current status: {status}" - if poll.results: - status_message = poll.results.cause - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach Cancelled, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"cluster_id={cluster_id}, command_id={command_id}, context_id={context_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_context_status_command_execution_running( - self, - cluster_id: str, - context_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[ContextStatusResponse], None]] = None, - ) -> ContextStatusResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (ContextStatus.RUNNING,) - failure_states = (ContextStatus.ERROR,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.context_status(cluster_id=cluster_id, context_id=context_id) - status = poll.status - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach Running, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"cluster_id={cluster_id}, context_id={context_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_command_status_command_execution_finished_or_error( - self, - cluster_id: str, - command_id: str, - context_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[CommandStatusResponse], None]] = None, - ) -> CommandStatusResponse: - deadline = time.time() + timeout.total_seconds() - target_states = ( - CommandStatus.FINISHED, - CommandStatus.ERROR, - ) - failure_states = ( - CommandStatus.CANCELLED, - CommandStatus.CANCELLING, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.command_status(cluster_id=cluster_id, command_id=command_id, context_id=context_id) - status = poll.status - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach Finished or Error, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"cluster_id={cluster_id}, command_id={command_id}, context_id={context_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def cancel( self, *, cluster_id: Optional[str] = None, command_id: Optional[str] = None, context_id: Optional[str] = None - ) -> Wait[CommandStatusResponse]: + ): """Cancel a command. Cancels a currently running command within an execution context. @@ -11565,7 +11108,7 @@ def cancel( :returns: Long-running operation waiter for :class:`CommandStatusResponse`. - See :method:wait_command_status_command_execution_cancelled for more details. + See :method:WaitCommandStatusCommandExecutionCancelled for more details. """ body = {} if cluster_id is not None: @@ -11579,24 +11122,7 @@ def cancel( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/1.2/commands/cancel", body=body, headers=headers) - return Wait( - self.wait_command_status_command_execution_cancelled, - response=CancelResponse.from_dict(op_response), - cluster_id=cluster_id, - command_id=command_id, - context_id=context_id, - ) - - def cancel_and_wait( - self, - *, - cluster_id: Optional[str] = None, - command_id: Optional[str] = None, - context_id: Optional[str] = None, - timeout=timedelta(minutes=20), - ) -> CommandStatusResponse: - return self.cancel(cluster_id=cluster_id, command_id=command_id, context_id=context_id).result(timeout=timeout) + self._api.do("POST", "/api/1.2/commands/cancel", body=body, headers=headers) def command_status(self, cluster_id: str, context_id: str, command_id: str) -> CommandStatusResponse: """Get command info. @@ -11649,9 +11175,7 @@ def context_status(self, cluster_id: str, context_id: str) -> ContextStatusRespo res = self._api.do("GET", "/api/1.2/contexts/status", query=query, headers=headers) return ContextStatusResponse.from_dict(res) - def create( - self, *, cluster_id: Optional[str] = None, language: Optional[Language] = None - ) -> Wait[ContextStatusResponse]: + def create(self, *, cluster_id: Optional[str] = None, language: Optional[Language] = None) -> Created: """Create an execution context. Creates an execution context for running cluster commands. @@ -11664,7 +11188,7 @@ def create( :returns: Long-running operation waiter for :class:`ContextStatusResponse`. - See :method:wait_context_status_command_execution_running for more details. + See :method:WaitContextStatusCommandExecutionRunning for more details. """ body = {} if cluster_id is not None: @@ -11676,18 +11200,8 @@ def create( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/1.2/contexts/create", body=body, headers=headers) - return Wait( - self.wait_context_status_command_execution_running, - response=Created.from_dict(op_response), - cluster_id=cluster_id, - context_id=op_response["id"], - ) - - def create_and_wait( - self, *, cluster_id: Optional[str] = None, language: Optional[Language] = None, timeout=timedelta(minutes=20) - ) -> ContextStatusResponse: - return self.create(cluster_id=cluster_id, language=language).result(timeout=timeout) + res = self._api.do("POST", "/api/1.2/contexts/create", body=body, headers=headers) + return Created.from_dict(res) def destroy(self, cluster_id: str, context_id: str): """Delete an execution context. @@ -11718,7 +11232,7 @@ def execute( command: Optional[str] = None, context_id: Optional[str] = None, language: Optional[Language] = None, - ) -> Wait[CommandStatusResponse]: + ) -> Created: """Run a command. Runs a cluster command in the given execution context, using the provided language. @@ -11735,7 +11249,7 @@ def execute( :returns: Long-running operation waiter for :class:`CommandStatusResponse`. - See :method:wait_command_status_command_execution_finished_or_error for more details. + See :method:WaitCommandStatusCommandExecutionFinishedOrError for more details. """ body = {} if cluster_id is not None: @@ -11751,27 +11265,8 @@ def execute( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/1.2/commands/execute", body=body, headers=headers) - return Wait( - self.wait_command_status_command_execution_finished_or_error, - response=Created.from_dict(op_response), - cluster_id=cluster_id, - command_id=op_response["id"], - context_id=context_id, - ) - - def execute_and_wait( - self, - *, - cluster_id: Optional[str] = None, - command: Optional[str] = None, - context_id: Optional[str] = None, - language: Optional[Language] = None, - timeout=timedelta(minutes=20), - ) -> CommandStatusResponse: - return self.execute(cluster_id=cluster_id, command=command, context_id=context_id, language=language).result( - timeout=timeout - ) + res = self._api.do("POST", "/api/1.2/commands/execute", body=body, headers=headers) + return Created.from_dict(res) class GlobalInitScriptsAPI: diff --git a/databricks/sdk/mixins/compute.py b/databricks/sdk/compute/v2/mixin.py similarity index 99% rename from databricks/sdk/mixins/compute.py rename to databricks/sdk/compute/v2/mixin.py index 3fafb416e..00f6ad575 100644 --- a/databricks/sdk/mixins/compute.py +++ b/databricks/sdk/compute/v2/mixin.py @@ -7,7 +7,8 @@ from databricks.sdk.databricks.core import DatabricksError from databricks.sdk.databricks.errors import OperationFailed -from databricks.sdk.service import compute + +from . import compute _LOG = logging.getLogger("databricks.sdk") diff --git a/databricks/sdk/dashboards/__init__.py b/databricks/sdk/dashboards/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/dashboards/v2/__init__.py b/databricks/sdk/dashboards/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/dashboards/v2/client.py b/databricks/sdk/dashboards/v2/client.py new file mode 100755 index 000000000..8d083a841 --- /dev/null +++ b/databricks/sdk/dashboards/v2/client.py @@ -0,0 +1,280 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .dashboards import (GenieAPI, LakeviewAPI, LakeviewEmbeddedAPI, + QueryExecutionAPI) + +_LOG = logging.getLogger(__name__) + + +class GenieClient(GenieAPI): + """ + Genie provides a no-code experience for business users, powered by AI/BI. Analysts set up spaces + that business users can use to ask questions using natural language. Genie uses data registered + to Unity Catalog and requires at least CAN USE permission on a Pro or Serverless SQL warehouse. + Also, Databricks Assistant must be enabled. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class LakeviewClient(LakeviewAPI): + """ + These APIs provide specific management operations for Lakeview dashboards. Generic resource + management can be done with Workspace API (import, export, get-status, list, delete). + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class LakeviewEmbeddedClient(LakeviewEmbeddedAPI): + """ + Token-based Lakeview APIs for embedding dashboards in external applications. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueryExecutionClient(QueryExecutionAPI): + """ + Query execution APIs for AI / BI Dashboards + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/dashboards.py b/databricks/sdk/dashboards/v2/dashboards.py similarity index 72% rename from databricks/sdk/service/dashboards.py rename to databricks/sdk/dashboards/v2/dashboards.py index de1042e7f..c242458e1 100755 --- a/databricks/sdk/service/dashboards.py +++ b/databricks/sdk/dashboards/v2/dashboards.py @@ -3,22 +3,151 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") +# all definitions in this file are in alphabetical order -from databricks.sdk.service import sql -# all definitions in this file are in alphabetical order +@dataclass +class AuthorizationDetails: + grant_rules: Optional[List[AuthorizationDetailsGrantRule]] = None + """Represents downscoped permission rules with specific access rights. This field is specific to + `workspace_rule_set` constraint.""" + + resource_legacy_acl_path: Optional[str] = None + """The acl path of the tree store resource resource.""" + + resource_name: Optional[str] = None + """The resource name to which the authorization rule applies. This field is specific to + `workspace_rule_set` constraint. Format: `workspaces/{workspace_id}/dashboards/{dashboard_id}`""" + + type: Optional[str] = None + """The type of authorization downscoping policy. Ex: `workspace_rule_set` defines access rules for + a specific workspace resource""" + + def as_dict(self) -> dict: + """Serializes the AuthorizationDetails into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.grant_rules: + body["grant_rules"] = [v.as_dict() for v in self.grant_rules] + if self.resource_legacy_acl_path is not None: + body["resource_legacy_acl_path"] = self.resource_legacy_acl_path + if self.resource_name is not None: + body["resource_name"] = self.resource_name + if self.type is not None: + body["type"] = self.type + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AuthorizationDetails into a shallow dictionary of its immediate attributes.""" + body = {} + if self.grant_rules: + body["grant_rules"] = self.grant_rules + if self.resource_legacy_acl_path is not None: + body["resource_legacy_acl_path"] = self.resource_legacy_acl_path + if self.resource_name is not None: + body["resource_name"] = self.resource_name + if self.type is not None: + body["type"] = self.type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AuthorizationDetails: + """Deserializes the AuthorizationDetails from a dictionary.""" + return cls( + grant_rules=_repeated_dict(d, "grant_rules", AuthorizationDetailsGrantRule), + resource_legacy_acl_path=d.get("resource_legacy_acl_path", None), + resource_name=d.get("resource_name", None), + type=d.get("type", None), + ) + + +@dataclass +class AuthorizationDetailsGrantRule: + permission_set: Optional[str] = None + """Permission sets for dashboard are defined in + iam-common/rbac-common/permission-sets/definitions/TreeStoreBasePermissionSets Ex: + `permissionSets/dashboard.runner`""" + + def as_dict(self) -> dict: + """Serializes the AuthorizationDetailsGrantRule into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.permission_set is not None: + body["permission_set"] = self.permission_set + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AuthorizationDetailsGrantRule into a shallow dictionary of its immediate attributes.""" + body = {} + if self.permission_set is not None: + body["permission_set"] = self.permission_set + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AuthorizationDetailsGrantRule: + """Deserializes the AuthorizationDetailsGrantRule from a dictionary.""" + return cls(permission_set=d.get("permission_set", None)) + + +@dataclass +class BaseChunkInfo: + """Describes metadata for a particular chunk, within a result set; this structure is used both + within a manifest, and when fetching individual chunk data or links.""" + + byte_count: Optional[int] = None + """The number of bytes in the result chunk. This field is not available when using `INLINE` + disposition.""" + + chunk_index: Optional[int] = None + """The position within the sequence of result set chunks.""" + + row_count: Optional[int] = None + """The number of rows within the result chunk.""" + + row_offset: Optional[int] = None + """The starting row offset within the result set.""" + + def as_dict(self) -> dict: + """Serializes the BaseChunkInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + def as_shallow_dict(self) -> dict: + """Serializes the BaseChunkInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> BaseChunkInfo: + """Deserializes the BaseChunkInfo from a dictionary.""" + return cls( + byte_count=d.get("byte_count", None), + chunk_index=d.get("chunk_index", None), + row_count=d.get("row_count", None), + row_offset=d.get("row_offset", None), + ) @dataclass @@ -91,6 +220,108 @@ def from_dict(cls, d: Dict[str, Any]) -> CancelQueryExecutionResponseStatus: ) +@dataclass +class ColumnInfo: + name: Optional[str] = None + """The name of the column.""" + + position: Optional[int] = None + """The ordinal position of the column (starting at position 0).""" + + type_interval_type: Optional[str] = None + """The format of the interval type.""" + + type_name: Optional[ColumnInfoTypeName] = None + """The name of the base data type. This doesn't include details for complex types such as STRUCT, + MAP or ARRAY.""" + + type_precision: Optional[int] = None + """Specifies the number of digits in a number. This applies to the DECIMAL type.""" + + type_scale: Optional[int] = None + """Specifies the number of digits to the right of the decimal point in a number. This applies to + the DECIMAL type.""" + + type_text: Optional[str] = None + """The full SQL type specification.""" + + def as_dict(self) -> dict: + """Serializes the ColumnInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.position is not None: + body["position"] = self.position + if self.type_interval_type is not None: + body["type_interval_type"] = self.type_interval_type + if self.type_name is not None: + body["type_name"] = self.type_name.value + if self.type_precision is not None: + body["type_precision"] = self.type_precision + if self.type_scale is not None: + body["type_scale"] = self.type_scale + if self.type_text is not None: + body["type_text"] = self.type_text + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ColumnInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.position is not None: + body["position"] = self.position + if self.type_interval_type is not None: + body["type_interval_type"] = self.type_interval_type + if self.type_name is not None: + body["type_name"] = self.type_name + if self.type_precision is not None: + body["type_precision"] = self.type_precision + if self.type_scale is not None: + body["type_scale"] = self.type_scale + if self.type_text is not None: + body["type_text"] = self.type_text + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ColumnInfo: + """Deserializes the ColumnInfo from a dictionary.""" + return cls( + name=d.get("name", None), + position=d.get("position", None), + type_interval_type=d.get("type_interval_type", None), + type_name=_enum(d, "type_name", ColumnInfoTypeName), + type_precision=d.get("type_precision", None), + type_scale=d.get("type_scale", None), + type_text=d.get("type_text", None), + ) + + +class ColumnInfoTypeName(Enum): + """The name of the base data type. This doesn't include details for complex types such as STRUCT, + MAP or ARRAY.""" + + ARRAY = "ARRAY" + BINARY = "BINARY" + BOOLEAN = "BOOLEAN" + BYTE = "BYTE" + CHAR = "CHAR" + DATE = "DATE" + DECIMAL = "DECIMAL" + DOUBLE = "DOUBLE" + FLOAT = "FLOAT" + INT = "INT" + INTERVAL = "INTERVAL" + LONG = "LONG" + MAP = "MAP" + NULL = "NULL" + SHORT = "SHORT" + STRING = "STRING" + STRUCT = "STRUCT" + TIMESTAMP = "TIMESTAMP" + USER_DEFINED_TYPE = "USER_DEFINED_TYPE" + + @dataclass class CronSchedule: quartz_cron_expression: str @@ -366,6 +597,112 @@ def from_dict(cls, d: Dict[str, Any]) -> ExecuteQueryResponse: return cls() +@dataclass +class ExternalLink: + byte_count: Optional[int] = None + """The number of bytes in the result chunk. This field is not available when using `INLINE` + disposition.""" + + chunk_index: Optional[int] = None + """The position within the sequence of result set chunks.""" + + expiration: Optional[str] = None + """Indicates the date-time that the given external link will expire and becomes invalid, after + which point a new `external_link` must be requested.""" + + external_link: Optional[str] = None + + http_headers: Optional[Dict[str, str]] = None + """HTTP headers that must be included with a GET request to the `external_link`. Each header is + provided as a key-value pair. Headers are typically used to pass a decryption key to the + external service. The values of these headers should be considered sensitive and the client + should not expose these values in a log.""" + + next_chunk_index: Optional[int] = None + """When fetching, provides the `chunk_index` for the _next_ chunk. If absent, indicates there are + no more chunks. The next chunk can be fetched with a + :method:statementexecution/getStatementResultChunkN request.""" + + next_chunk_internal_link: Optional[str] = None + """When fetching, provides a link to fetch the _next_ chunk. If absent, indicates there are no more + chunks. This link is an absolute `path` to be joined with your `$DATABRICKS_HOST`, and should be + treated as an opaque link. This is an alternative to using `next_chunk_index`.""" + + row_count: Optional[int] = None + """The number of rows within the result chunk.""" + + row_offset: Optional[int] = None + """The starting row offset within the result set.""" + + def as_dict(self) -> dict: + """Serializes the ExternalLink into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.expiration is not None: + body["expiration"] = self.expiration + if self.external_link is not None: + body["external_link"] = self.external_link + if self.http_headers: + body["http_headers"] = self.http_headers + if self.next_chunk_index is not None: + body["next_chunk_index"] = self.next_chunk_index + if self.next_chunk_internal_link is not None: + body["next_chunk_internal_link"] = self.next_chunk_internal_link + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ExternalLink into a shallow dictionary of its immediate attributes.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.expiration is not None: + body["expiration"] = self.expiration + if self.external_link is not None: + body["external_link"] = self.external_link + if self.http_headers: + body["http_headers"] = self.http_headers + if self.next_chunk_index is not None: + body["next_chunk_index"] = self.next_chunk_index + if self.next_chunk_internal_link is not None: + body["next_chunk_internal_link"] = self.next_chunk_internal_link + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ExternalLink: + """Deserializes the ExternalLink from a dictionary.""" + return cls( + byte_count=d.get("byte_count", None), + chunk_index=d.get("chunk_index", None), + expiration=d.get("expiration", None), + external_link=d.get("external_link", None), + http_headers=d.get("http_headers", None), + next_chunk_index=d.get("next_chunk_index", None), + next_chunk_internal_link=d.get("next_chunk_internal_link", None), + row_count=d.get("row_count", None), + row_offset=d.get("row_offset", None), + ) + + +class Format(Enum): + + ARROW_STREAM = "ARROW_STREAM" + CSV = "CSV" + JSON_ARRAY = "JSON_ARRAY" + + @dataclass class GenieAttachment: """Genie AI Response""" @@ -529,9 +866,88 @@ def from_dict(cls, d: Dict[str, Any]) -> GenieCreateConversationMessageRequest: ) +@dataclass +class GenieGenerateDownloadFullQueryResultResponse: + error: Optional[str] = None + """Error message if Genie failed to download the result""" + + status: Optional[MessageStatus] = None + """Download result status""" + + transient_statement_id: Optional[str] = None + """Transient Statement ID. Use this ID to track the download request in subsequent polling calls""" + + def as_dict(self) -> dict: + """Serializes the GenieGenerateDownloadFullQueryResultResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.error is not None: + body["error"] = self.error + if self.status is not None: + body["status"] = self.status.value + if self.transient_statement_id is not None: + body["transient_statement_id"] = self.transient_statement_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieGenerateDownloadFullQueryResultResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.error is not None: + body["error"] = self.error + if self.status is not None: + body["status"] = self.status + if self.transient_statement_id is not None: + body["transient_statement_id"] = self.transient_statement_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieGenerateDownloadFullQueryResultResponse: + """Deserializes the GenieGenerateDownloadFullQueryResultResponse from a dictionary.""" + return cls( + error=d.get("error", None), + status=_enum(d, "status", MessageStatus), + transient_statement_id=d.get("transient_statement_id", None), + ) + + +@dataclass +class GenieGetDownloadFullQueryResultResponse: + statement_response: Optional[StatementResponse] = None + """SQL Statement Execution response. See [Get status, manifest, and result first + chunk](:method:statementexecution/getstatement) for more details.""" + + transient_statement_id: Optional[str] = None + """Transient Statement ID""" + + def as_dict(self) -> dict: + """Serializes the GenieGetDownloadFullQueryResultResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.statement_response: + body["statement_response"] = self.statement_response.as_dict() + if self.transient_statement_id is not None: + body["transient_statement_id"] = self.transient_statement_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GenieGetDownloadFullQueryResultResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.statement_response: + body["statement_response"] = self.statement_response + if self.transient_statement_id is not None: + body["transient_statement_id"] = self.transient_statement_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GenieGetDownloadFullQueryResultResponse: + """Deserializes the GenieGetDownloadFullQueryResultResponse from a dictionary.""" + return cls( + statement_response=_from_dict(d, "statement_response", StatementResponse), + transient_statement_id=d.get("transient_statement_id", None), + ) + + @dataclass class GenieGetMessageQueryResultResponse: - statement_response: Optional[sql.StatementResponse] = None + statement_response: Optional[StatementResponse] = None """SQL Statement Execution response. See [Get status, manifest, and result first chunk](:method:statementexecution/getstatement) for more details.""" @@ -552,7 +968,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> GenieGetMessageQueryResultResponse: """Deserializes the GenieGetMessageQueryResultResponse from a dictionary.""" - return cls(statement_response=_from_dict(d, "statement_response", sql.StatementResponse)) + return cls(statement_response=_from_dict(d, "statement_response", StatementResponse)) @dataclass @@ -929,6 +1345,52 @@ def from_dict(cls, d: Dict[str, Any]) -> GetPublishedDashboardEmbeddedResponse: return cls() +@dataclass +class GetPublishedDashboardTokenInfoResponse: + authorization_details: Optional[List[AuthorizationDetails]] = None + """Authorization constraints for accessing the published dashboard. Currently includes + `workspace_rule_set` and could be enriched with `unity_catalog_privileges` before oAuth token + generation.""" + + custom_claim: Optional[str] = None + """Custom claim generated from external_value and external_viewer_id. Format: + `urn:aibi:external_data:::`""" + + scope: Optional[str] = None + """Scope defining access permissions.""" + + def as_dict(self) -> dict: + """Serializes the GetPublishedDashboardTokenInfoResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.authorization_details: + body["authorization_details"] = [v.as_dict() for v in self.authorization_details] + if self.custom_claim is not None: + body["custom_claim"] = self.custom_claim + if self.scope is not None: + body["scope"] = self.scope + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GetPublishedDashboardTokenInfoResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.authorization_details: + body["authorization_details"] = self.authorization_details + if self.custom_claim is not None: + body["custom_claim"] = self.custom_claim + if self.scope is not None: + body["scope"] = self.scope + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GetPublishedDashboardTokenInfoResponse: + """Deserializes the GetPublishedDashboardTokenInfoResponse from a dictionary.""" + return cls( + authorization_details=_repeated_dict(d, "authorization_details", AuthorizationDetails), + custom_claim=d.get("custom_claim", None), + scope=d.get("scope", None), + ) + + class LifecycleState(Enum): ACTIVE = "ACTIVE" @@ -1470,6 +1932,203 @@ def from_dict(cls, d: Dict[str, Any]) -> Result: ) +@dataclass +class ResultData: + byte_count: Optional[int] = None + """The number of bytes in the result chunk. This field is not available when using `INLINE` + disposition.""" + + chunk_index: Optional[int] = None + """The position within the sequence of result set chunks.""" + + data_array: Optional[List[List[str]]] = None + """The `JSON_ARRAY` format is an array of arrays of values, where each non-null value is formatted + as a string. Null values are encoded as JSON `null`.""" + + external_links: Optional[List[ExternalLink]] = None + + next_chunk_index: Optional[int] = None + """When fetching, provides the `chunk_index` for the _next_ chunk. If absent, indicates there are + no more chunks. The next chunk can be fetched with a + :method:statementexecution/getStatementResultChunkN request.""" + + next_chunk_internal_link: Optional[str] = None + """When fetching, provides a link to fetch the _next_ chunk. If absent, indicates there are no more + chunks. This link is an absolute `path` to be joined with your `$DATABRICKS_HOST`, and should be + treated as an opaque link. This is an alternative to using `next_chunk_index`.""" + + row_count: Optional[int] = None + """The number of rows within the result chunk.""" + + row_offset: Optional[int] = None + """The starting row offset within the result set.""" + + def as_dict(self) -> dict: + """Serializes the ResultData into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.data_array: + body["data_array"] = [v for v in self.data_array] + if self.external_links: + body["external_links"] = [v.as_dict() for v in self.external_links] + if self.next_chunk_index is not None: + body["next_chunk_index"] = self.next_chunk_index + if self.next_chunk_internal_link is not None: + body["next_chunk_internal_link"] = self.next_chunk_internal_link + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ResultData into a shallow dictionary of its immediate attributes.""" + body = {} + if self.byte_count is not None: + body["byte_count"] = self.byte_count + if self.chunk_index is not None: + body["chunk_index"] = self.chunk_index + if self.data_array: + body["data_array"] = self.data_array + if self.external_links: + body["external_links"] = self.external_links + if self.next_chunk_index is not None: + body["next_chunk_index"] = self.next_chunk_index + if self.next_chunk_internal_link is not None: + body["next_chunk_internal_link"] = self.next_chunk_internal_link + if self.row_count is not None: + body["row_count"] = self.row_count + if self.row_offset is not None: + body["row_offset"] = self.row_offset + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ResultData: + """Deserializes the ResultData from a dictionary.""" + return cls( + byte_count=d.get("byte_count", None), + chunk_index=d.get("chunk_index", None), + data_array=d.get("data_array", None), + external_links=_repeated_dict(d, "external_links", ExternalLink), + next_chunk_index=d.get("next_chunk_index", None), + next_chunk_internal_link=d.get("next_chunk_internal_link", None), + row_count=d.get("row_count", None), + row_offset=d.get("row_offset", None), + ) + + +@dataclass +class ResultManifest: + """The result manifest provides schema and metadata for the result set.""" + + chunks: Optional[List[BaseChunkInfo]] = None + """Array of result set chunk metadata.""" + + format: Optional[Format] = None + + schema: Optional[ResultSchema] = None + """The schema is an ordered list of column descriptions.""" + + total_byte_count: Optional[int] = None + """The total number of bytes in the result set. This field is not available when using `INLINE` + disposition.""" + + total_chunk_count: Optional[int] = None + """The total number of chunks that the result set has been divided into.""" + + total_row_count: Optional[int] = None + """The total number of rows in the result set.""" + + truncated: Optional[bool] = None + """Indicates whether the result is truncated due to `row_limit` or `byte_limit`.""" + + def as_dict(self) -> dict: + """Serializes the ResultManifest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.chunks: + body["chunks"] = [v.as_dict() for v in self.chunks] + if self.format is not None: + body["format"] = self.format.value + if self.schema: + body["schema"] = self.schema.as_dict() + if self.total_byte_count is not None: + body["total_byte_count"] = self.total_byte_count + if self.total_chunk_count is not None: + body["total_chunk_count"] = self.total_chunk_count + if self.total_row_count is not None: + body["total_row_count"] = self.total_row_count + if self.truncated is not None: + body["truncated"] = self.truncated + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ResultManifest into a shallow dictionary of its immediate attributes.""" + body = {} + if self.chunks: + body["chunks"] = self.chunks + if self.format is not None: + body["format"] = self.format + if self.schema: + body["schema"] = self.schema + if self.total_byte_count is not None: + body["total_byte_count"] = self.total_byte_count + if self.total_chunk_count is not None: + body["total_chunk_count"] = self.total_chunk_count + if self.total_row_count is not None: + body["total_row_count"] = self.total_row_count + if self.truncated is not None: + body["truncated"] = self.truncated + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ResultManifest: + """Deserializes the ResultManifest from a dictionary.""" + return cls( + chunks=_repeated_dict(d, "chunks", BaseChunkInfo), + format=_enum(d, "format", Format), + schema=_from_dict(d, "schema", ResultSchema), + total_byte_count=d.get("total_byte_count", None), + total_chunk_count=d.get("total_chunk_count", None), + total_row_count=d.get("total_row_count", None), + truncated=d.get("truncated", None), + ) + + +@dataclass +class ResultSchema: + """The schema is an ordered list of column descriptions.""" + + column_count: Optional[int] = None + + columns: Optional[List[ColumnInfo]] = None + + def as_dict(self) -> dict: + """Serializes the ResultSchema into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.column_count is not None: + body["column_count"] = self.column_count + if self.columns: + body["columns"] = [v.as_dict() for v in self.columns] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ResultSchema into a shallow dictionary of its immediate attributes.""" + body = {} + if self.column_count is not None: + body["column_count"] = self.column_count + if self.columns: + body["columns"] = self.columns + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ResultSchema: + """Deserializes the ResultSchema from a dictionary.""" + return cls(column_count=d.get("column_count", None), columns=_repeated_dict(d, "columns", ColumnInfo)) + + @dataclass class Schedule: cron_schedule: CronSchedule @@ -1569,6 +2228,158 @@ class SchedulePauseStatus(Enum): UNPAUSED = "UNPAUSED" +@dataclass +class ServiceError: + error_code: Optional[ServiceErrorCode] = None + + message: Optional[str] = None + """A brief summary of the error condition.""" + + def as_dict(self) -> dict: + """Serializes the ServiceError into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.error_code is not None: + body["error_code"] = self.error_code.value + if self.message is not None: + body["message"] = self.message + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ServiceError into a shallow dictionary of its immediate attributes.""" + body = {} + if self.error_code is not None: + body["error_code"] = self.error_code + if self.message is not None: + body["message"] = self.message + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ServiceError: + """Deserializes the ServiceError from a dictionary.""" + return cls(error_code=_enum(d, "error_code", ServiceErrorCode), message=d.get("message", None)) + + +class ServiceErrorCode(Enum): + + ABORTED = "ABORTED" + ALREADY_EXISTS = "ALREADY_EXISTS" + BAD_REQUEST = "BAD_REQUEST" + CANCELLED = "CANCELLED" + DEADLINE_EXCEEDED = "DEADLINE_EXCEEDED" + INTERNAL_ERROR = "INTERNAL_ERROR" + IO_ERROR = "IO_ERROR" + NOT_FOUND = "NOT_FOUND" + RESOURCE_EXHAUSTED = "RESOURCE_EXHAUSTED" + SERVICE_UNDER_MAINTENANCE = "SERVICE_UNDER_MAINTENANCE" + TEMPORARILY_UNAVAILABLE = "TEMPORARILY_UNAVAILABLE" + UNAUTHENTICATED = "UNAUTHENTICATED" + UNKNOWN = "UNKNOWN" + WORKSPACE_TEMPORARILY_UNAVAILABLE = "WORKSPACE_TEMPORARILY_UNAVAILABLE" + + +@dataclass +class StatementResponse: + manifest: Optional[ResultManifest] = None + """The result manifest provides schema and metadata for the result set.""" + + result: Optional[ResultData] = None + + statement_id: Optional[str] = None + """The statement ID is returned upon successfully submitting a SQL statement, and is a required + reference for all subsequent calls.""" + + status: Optional[StatementStatus] = None + """The status response includes execution state and if relevant, error information.""" + + def as_dict(self) -> dict: + """Serializes the StatementResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.manifest: + body["manifest"] = self.manifest.as_dict() + if self.result: + body["result"] = self.result.as_dict() + if self.statement_id is not None: + body["statement_id"] = self.statement_id + if self.status: + body["status"] = self.status.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the StatementResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.manifest: + body["manifest"] = self.manifest + if self.result: + body["result"] = self.result + if self.statement_id is not None: + body["statement_id"] = self.statement_id + if self.status: + body["status"] = self.status + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> StatementResponse: + """Deserializes the StatementResponse from a dictionary.""" + return cls( + manifest=_from_dict(d, "manifest", ResultManifest), + result=_from_dict(d, "result", ResultData), + statement_id=d.get("statement_id", None), + status=_from_dict(d, "status", StatementStatus), + ) + + +class StatementState(Enum): + """Statement execution state: - `PENDING`: waiting for warehouse - `RUNNING`: running - + `SUCCEEDED`: execution was successful, result data available for fetch - `FAILED`: execution + failed; reason for failure described in accomanying error message - `CANCELED`: user canceled; + can come from explicit cancel call, or timeout with `on_wait_timeout=CANCEL` - `CLOSED`: + execution successful, and statement closed; result no longer available for fetch""" + + CANCELED = "CANCELED" + CLOSED = "CLOSED" + FAILED = "FAILED" + PENDING = "PENDING" + RUNNING = "RUNNING" + SUCCEEDED = "SUCCEEDED" + + +@dataclass +class StatementStatus: + """The status response includes execution state and if relevant, error information.""" + + error: Optional[ServiceError] = None + + state: Optional[StatementState] = None + """Statement execution state: - `PENDING`: waiting for warehouse - `RUNNING`: running - + `SUCCEEDED`: execution was successful, result data available for fetch - `FAILED`: execution + failed; reason for failure described in accomanying error message - `CANCELED`: user canceled; + can come from explicit cancel call, or timeout with `on_wait_timeout=CANCEL` - `CLOSED`: + execution successful, and statement closed; result no longer available for fetch""" + + def as_dict(self) -> dict: + """Serializes the StatementStatus into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.error: + body["error"] = self.error.as_dict() + if self.state is not None: + body["state"] = self.state.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the StatementStatus into a shallow dictionary of its immediate attributes.""" + body = {} + if self.error: + body["error"] = self.error + if self.state is not None: + body["state"] = self.state + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> StatementStatus: + """Deserializes the StatementStatus from a dictionary.""" + return cls(error=_from_dict(d, "error", ServiceError), state=_enum(d, "state", StatementState)) + + @dataclass class Subscriber: destination_subscriber: Optional[SubscriptionSubscriberDestination] = None @@ -1850,41 +2661,7 @@ class GenieAPI: def __init__(self, api_client): self._api = api_client - def wait_get_message_genie_completed( - self, - conversation_id: str, - message_id: str, - space_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[GenieMessage], None]] = None, - ) -> GenieMessage: - deadline = time.time() + timeout.total_seconds() - target_states = (MessageStatus.COMPLETED,) - failure_states = (MessageStatus.FAILED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get_message(conversation_id=conversation_id, message_id=message_id, space_id=space_id) - status = poll.status - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach COMPLETED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"conversation_id={conversation_id}, message_id={message_id}, space_id={space_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def create_message(self, space_id: str, conversation_id: str, content: str) -> Wait[GenieMessage]: + def create_message(self, space_id: str, conversation_id: str, content: str) -> GenieMessage: """Create conversation message. Create new message in a [conversation](:method:genie/startconversation). The AI response uses all @@ -1899,7 +2676,7 @@ def create_message(self, space_id: str, conversation_id: str, content: str) -> W :returns: Long-running operation waiter for :class:`GenieMessage`. - See :method:wait_get_message_genie_completed for more details. + See :method:WaitGetMessageGenieCompleted for more details. """ body = {} if content is not None: @@ -1909,26 +2686,13 @@ def create_message(self, space_id: str, conversation_id: str, content: str) -> W "Content-Type": "application/json", } - op_response = self._api.do( + res = self._api.do( "POST", f"/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages", body=body, headers=headers, ) - return Wait( - self.wait_get_message_genie_completed, - response=GenieMessage.from_dict(op_response), - conversation_id=conversation_id, - message_id=op_response["id"], - space_id=space_id, - ) - - def create_message_and_wait( - self, space_id: str, conversation_id: str, content: str, timeout=timedelta(minutes=20) - ) -> GenieMessage: - return self.create_message(content=content, conversation_id=conversation_id, space_id=space_id).result( - timeout=timeout - ) + return GenieMessage.from_dict(res) def execute_message_attachment_query( self, space_id: str, conversation_id: str, message_id: str, attachment_id: str @@ -1989,6 +2753,85 @@ def execute_message_query( ) return GenieGetMessageQueryResultResponse.from_dict(res) + def generate_download_full_query_result( + self, space_id: str, conversation_id: str, message_id: str, attachment_id: str + ) -> GenieGenerateDownloadFullQueryResultResponse: + """Generate full query result download. + + Initiate full SQL query result download and obtain a transient ID for tracking the download progress. + This call initiates a new SQL execution to generate the query result. The result is stored in an + external link can be retrieved using the [Get Download Full Query + Result](:method:genie/getdownloadfullqueryresult) API. Warning: Databricks strongly recommends that + you protect the URLs that are returned by the `EXTERNAL_LINKS` disposition. See [Execute + Statement](:method:statementexecution/executestatement) for more details. + + :param space_id: str + Space ID + :param conversation_id: str + Conversation ID + :param message_id: str + Message ID + :param attachment_id: str + Attachment ID + + :returns: :class:`GenieGenerateDownloadFullQueryResultResponse` + """ + + headers = { + "Accept": "application/json", + } + + res = self._api.do( + "POST", + f"/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}/attachments/{attachment_id}/generate-download", + headers=headers, + ) + return GenieGenerateDownloadFullQueryResultResponse.from_dict(res) + + def get_download_full_query_result( + self, space_id: str, conversation_id: str, message_id: str, attachment_id: str, transient_statement_id: str + ) -> GenieGetDownloadFullQueryResultResponse: + """Get download full query result status. + + Poll download progress and retrieve the SQL query result external link(s) upon completion. Warning: + Databricks strongly recommends that you protect the URLs that are returned by the `EXTERNAL_LINKS` + disposition. When you use the `EXTERNAL_LINKS` disposition, a short-lived, presigned URL is generated, + which can be used to download the results directly from Amazon S3. As a short-lived access credential + is embedded in this presigned URL, you should protect the URL. Because presigned URLs are already + generated with embedded temporary access credentials, you must not set an Authorization header in the + download requests. See [Execute Statement](:method:statementexecution/executestatement) for more + details. + + :param space_id: str + Space ID + :param conversation_id: str + Conversation ID + :param message_id: str + Message ID + :param attachment_id: str + Attachment ID + :param transient_statement_id: str + Transient Statement ID. This ID is provided by the [Start Download + endpoint](:method:genie/startdownloadfullqueryresult) + + :returns: :class:`GenieGetDownloadFullQueryResultResponse` + """ + + query = {} + if transient_statement_id is not None: + query["transient_statement_id"] = transient_statement_id + headers = { + "Accept": "application/json", + } + + res = self._api.do( + "GET", + f"/api/2.0/genie/spaces/{space_id}/conversations/{conversation_id}/messages/{message_id}/attachments/{attachment_id}/get-download", + query=query, + headers=headers, + ) + return GenieGetDownloadFullQueryResultResponse.from_dict(res) + def get_message(self, space_id: str, conversation_id: str, message_id: str) -> GenieMessage: """Get conversation message. @@ -2124,7 +2967,7 @@ def get_space(self, space_id: str) -> GenieSpace: res = self._api.do("GET", f"/api/2.0/genie/spaces/{space_id}", headers=headers) return GenieSpace.from_dict(res) - def start_conversation(self, space_id: str, content: str) -> Wait[GenieMessage]: + def start_conversation(self, space_id: str, content: str) -> GenieStartConversationResponse: """Start conversation. Start a new conversation. @@ -2136,7 +2979,7 @@ def start_conversation(self, space_id: str, content: str) -> Wait[GenieMessage]: :returns: Long-running operation waiter for :class:`GenieMessage`. - See :method:wait_get_message_genie_completed for more details. + See :method:WaitGetMessageGenieCompleted for more details. """ body = {} if content is not None: @@ -2146,19 +2989,8 @@ def start_conversation(self, space_id: str, content: str) -> Wait[GenieMessage]: "Content-Type": "application/json", } - op_response = self._api.do( - "POST", f"/api/2.0/genie/spaces/{space_id}/start-conversation", body=body, headers=headers - ) - return Wait( - self.wait_get_message_genie_completed, - response=GenieStartConversationResponse.from_dict(op_response), - conversation_id=op_response["conversation_id"], - message_id=op_response["message_id"], - space_id=space_id, - ) - - def start_conversation_and_wait(self, space_id: str, content: str, timeout=timedelta(minutes=20)) -> GenieMessage: - return self.start_conversation(content=content, space_id=space_id).result(timeout=timeout) + res = self._api.do("POST", f"/api/2.0/genie/spaces/{space_id}/start-conversation", body=body, headers=headers) + return GenieStartConversationResponse.from_dict(res) class LakeviewAPI: @@ -2665,6 +3497,42 @@ def get_published_dashboard_embedded(self, dashboard_id: str): self._api.do("GET", f"/api/2.0/lakeview/dashboards/{dashboard_id}/published/embedded", headers=headers) + def get_published_dashboard_token_info( + self, dashboard_id: str, *, external_value: Optional[str] = None, external_viewer_id: Optional[str] = None + ) -> GetPublishedDashboardTokenInfoResponse: + """Read an information of a published dashboard to mint an OAuth token. + + Get a required authorization details and scopes of a published dashboard to mint an OAuth token. The + `authorization_details` can be enriched to apply additional restriction. + + Example: Adding the following `authorization_details` object to downscope the viewer permission to + specific table ``` { type: "unity_catalog_privileges", privileges: ["SELECT"], object_type: "TABLE", + object_full_path: "main.default.testdata" } ``` + + :param dashboard_id: str + UUID identifying the published dashboard. + :param external_value: str (optional) + Provided external value to be included in the custom claim. + :param external_viewer_id: str (optional) + Provided external viewer id to be included in the custom claim. + + :returns: :class:`GetPublishedDashboardTokenInfoResponse` + """ + + query = {} + if external_value is not None: + query["external_value"] = external_value + if external_viewer_id is not None: + query["external_viewer_id"] = external_viewer_id + headers = { + "Accept": "application/json", + } + + res = self._api.do( + "GET", f"/api/2.0/lakeview/dashboards/{dashboard_id}/published/tokeninfo", query=query, headers=headers + ) + return GetPublishedDashboardTokenInfoResponse.from_dict(res) + class QueryExecutionAPI: """Query execution APIs for AI / BI Dashboards""" diff --git a/databricks/sdk/databricks/azure.py b/databricks/sdk/databricks/azure.py index 84a638969..abfda955a 100644 --- a/databricks/sdk/databricks/azure.py +++ b/databricks/sdk/databricks/azure.py @@ -1,6 +1,7 @@ from typing import Dict -from ..service.provisioning import Workspace +from databricks.sdk.provisioning.v2.provisioning import Workspace + from .oauth import TokenSource diff --git a/databricks/sdk/databricks/dbutils.py b/databricks/sdk/databricks/dbutils.py index 64b410e36..1c09f028c 100644 --- a/databricks/sdk/databricks/dbutils.py +++ b/databricks/sdk/databricks/dbutils.py @@ -7,9 +7,10 @@ from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional -from ..mixins import compute as compute_ext -from ..mixins import files as dbfs_ext -from ..service import compute, workspace +from ..compute.v2 import compute +from ..compute.v2.mixin import ClustersExt as compute_ext +from ..files.v2.mixin import DbfsExt as dbfs_ext +from ..workspace.v2 import workspace from .core import ApiClient, Config, DatabricksError _LOG = logging.getLogger("databricks.sdk") @@ -38,7 +39,7 @@ class _FsUtil: def __init__( self, - dbfs_ext: dbfs_ext.DbfsExt, + dbfs_ext: dbfs_ext, proxy_factory: Callable[[str], "_ProxyUtil"], ): self._dbfs = dbfs_ext @@ -212,12 +213,12 @@ class RemoteDbUtils: def __init__(self, config: "Config" = None): self._config = Config() if not config else config self._client = ApiClient(self._config) - self._clusters = compute_ext.ClustersExt(self._client) + self._clusters = compute_ext(self._client) self._commands = compute.CommandExecutionAPI(self._client) self._lock = threading.Lock() self._ctx = None - self.fs = _FsUtil(dbfs_ext.DbfsExt(self._client), self.__getattr__) + self.fs = _FsUtil(dbfs_ext(self._client), self.__getattr__) self.secrets = _SecretsUtil(workspace.SecretsAPI(self._client)) self.jobs = _JobsUtil() self._widgets = None diff --git a/databricks/sdk/files/__init__.py b/databricks/sdk/files/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/files/v2/__init__.py b/databricks/sdk/files/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/files/v2/client.py b/databricks/sdk/files/v2/client.py new file mode 100755 index 000000000..7b6fd8a37 --- /dev/null +++ b/databricks/sdk/files/v2/client.py @@ -0,0 +1,161 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .mixin import DbfsExt, FilesExt + +_LOG = logging.getLogger(__name__) + + +class DbfsClient(DbfsExt): + """ + DBFS API makes it simple to interact with various data sources without having to include a users + credentials every time to read a file. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class FilesClient(FilesExt): + """ + The Files API is a standard HTTP API that allows you to read, write, list, and delete files and + directories by referring to their URI. The API makes working with file content as raw bytes + easier and more efficient. + + The API supports [Unity Catalog volumes], where files and directories to operate on are + specified using their volume URI path, which follows the format + /Volumes/<catalog_name>/<schema_name>/<volume_name>/<path_to_file>. + + The Files API has two distinct endpoints, one for working with files (`/fs/files`) and another + one for working with directories (`/fs/directories`). Both endpoints use the standard HTTP + methods GET, HEAD, PUT, and DELETE to manage files and directories specified using their URI + path. The path is always absolute. + + Some Files API client features are currently experimental. To enable them, set + `enable_experimental_files_api_client = True` in your configuration profile or use the + environment variable `DATABRICKS_ENABLE_EXPERIMENTAL_FILES_API_CLIENT=True`. + + [Unity Catalog volumes]: https://docs.databricks.com/en/connect/unity-catalog/volumes.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config), config=config) diff --git a/databricks/sdk/service/files.py b/databricks/sdk/files/v2/files.py similarity index 99% rename from databricks/sdk/service/files.py rename to databricks/sdk/files/v2/files.py index 8d60b842f..c2cc92775 100755 --- a/databricks/sdk/service/files.py +++ b/databricks/sdk/files/v2/files.py @@ -6,11 +6,11 @@ from dataclasses import dataclass from typing import Any, BinaryIO, Dict, Iterator, List, Optional -from ._internal import _escape_multi_segment_path_parameter, _repeated_dict +from ...service._internal import (_escape_multi_segment_path_parameter, + _repeated_dict) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order diff --git a/databricks/sdk/mixins/files.py b/databricks/sdk/files/v2/mixin.py similarity index 98% rename from databricks/sdk/mixins/files.py rename to databricks/sdk/files/v2/mixin.py index 08c4d4806..c45014765 100644 --- a/databricks/sdk/mixins/files.py +++ b/databricks/sdk/files/v2/mixin.py @@ -24,17 +24,17 @@ import requests.adapters from requests import RequestException -from ..databricks._base_client import (_BaseClient, _RawResponse, - _StreamingResponse) -from ..databricks._property import _cached_property -from ..databricks.config import Config -from ..databricks.errors import AlreadyExists, NotFound -from ..databricks.errors.customizer import _RetryAfterCustomizer -from ..databricks.errors.mapper import _error_mapper -from ..databricks.retries import retried -from ..service import files -from ..service._internal import _escape_multi_segment_path_parameter -from ..service.files import DownloadResponse +from ...databricks._base_client import (_BaseClient, _RawResponse, + _StreamingResponse) +from ...databricks._property import _cached_property +from ...databricks.config import Config +from ...databricks.errors import AlreadyExists, NotFound +from ...databricks.errors.customizer import _RetryAfterCustomizer +from ...databricks.errors.mapper import _error_mapper +from ...databricks.retries import retried +from ...service._internal import _escape_multi_segment_path_parameter +from . import files +from .files import DownloadResponse if TYPE_CHECKING: from _typeshed import Self @@ -718,6 +718,10 @@ def download(self, file_path: str) -> DownloadResponse: :returns: :class:`DownloadResponse` """ + if not self._config.enable_experimental_files_api_client: + # Use the new Files API client for downloads + return super().download(file_path) + initial_response: DownloadResponse = self._open_download_stream( file_path=file_path, start_byte_offset=0, @@ -743,6 +747,10 @@ def upload(self, file_path: str, contents: BinaryIO, *, overwrite: Optional[bool If true, an existing file will be overwritten. When not specified, assumed True. """ + if not self._config.enable_experimental_files_api_client: + # Use the new Files API client for downloads + return super().upload(file_path, contents, overwrite=overwrite) + # Upload empty and small files with one-shot upload. pre_read_buffer = contents.read(self._config.multipart_upload_min_stream_size) if len(pre_read_buffer) < self._config.multipart_upload_min_stream_size: diff --git a/databricks/sdk/iam/__init__.py b/databricks/sdk/iam/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/iam/v2/__init__.py b/databricks/sdk/iam/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/iam/v2/client.py b/databricks/sdk/iam/v2/client.py new file mode 100755 index 000000000..4c5c27cc0 --- /dev/null +++ b/databricks/sdk/iam/v2/client.py @@ -0,0 +1,963 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .iam import (AccessControlAPI, AccountAccessControlAPI, + AccountAccessControlProxyAPI, AccountGroupsAPI, + AccountServicePrincipalsAPI, AccountUsersAPI, CurrentUserAPI, + GroupsAPI, PermissionMigrationAPI, PermissionsAPI, + ServicePrincipalsAPI, UsersAPI, WorkspaceAssignmentAPI) + +_LOG = logging.getLogger(__name__) + + +class AccessControlClient(AccessControlAPI): + """ + Rule based Access Control for Databricks Resources. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountAccessControlClient(AccountAccessControlAPI): + """ + These APIs manage access rules on resources in an account. Currently, only grant rules are + supported. A grant rule specifies a role assigned to a set of principals. A list of rules + attached to a resource is called a rule set. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountAccessControlProxyClient(AccountAccessControlProxyAPI): + """ + These APIs manage access rules on resources in an account. Currently, only grant rules are + supported. A grant rule specifies a role assigned to a set of principals. A list of rules + attached to a resource is called a rule set. A workspace must belong to an account for these + APIs to work. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountGroupsClient(AccountGroupsAPI): + """ + Groups simplify identity management, making it easier to assign access to Databricks account, + data, and other securable objects. + + It is best practice to assign access to workspaces and access-control policies in Unity Catalog + to groups, instead of to users individually. All Databricks account identities can be assigned + as members of groups, and members inherit permissions that are assigned to their group. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountServicePrincipalsClient(AccountServicePrincipalsAPI): + """ + Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD + platforms. Databricks recommends creating service principals to run production jobs or modify + production data. If all processes that act on production data run with service principals, + interactive users do not need any write, delete, or modify privileges in production. This + eliminates the risk of a user overwriting production data by accident. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountUsersClient(AccountUsersAPI): + """ + User identities recognized by Databricks and represented by email addresses. + + Databricks recommends using SCIM provisioning to sync users and groups automatically from your + identity provider to your Databricks account. SCIM streamlines onboarding a new employee or team + by using your identity provider to create users and groups in Databricks account and give them + the proper level of access. When a user leaves your organization or no longer needs access to + Databricks account, admins can terminate the user in your identity provider and that user’s + account will also be removed from Databricks account. This ensures a consistent offboarding + process and prevents unauthorized users from accessing sensitive data. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CurrentUserClient(CurrentUserAPI): + """ + This API allows retrieving information about currently authenticated user or service principal. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class GroupsClient(GroupsAPI): + """ + Groups simplify identity management, making it easier to assign access to Databricks workspace, + data, and other securable objects. + + It is best practice to assign access to workspaces and access-control policies in Unity Catalog + to groups, instead of to users individually. All Databricks workspace identities can be assigned + as members of groups, and members inherit permissions that are assigned to their group. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PermissionMigrationClient(PermissionMigrationAPI): + """ + APIs for migrating acl permissions, used only by the ucx tool: + https://github.com/databrickslabs/ucx + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PermissionsClient(PermissionsAPI): + """ + Permissions API are used to create read, write, edit, update and manage access for various users + on different objects and endpoints. + + * **[Apps permissions](:service:apps)** — Manage which users can manage or use apps. + + * **[Cluster permissions](:service:clusters)** — Manage which users can manage, restart, or + attach to clusters. + + * **[Cluster policy permissions](:service:clusterpolicies)** — Manage which users can use + cluster policies. + + * **[Delta Live Tables pipeline permissions](:service:pipelines)** — Manage which users can + view, manage, run, cancel, or own a Delta Live Tables pipeline. + + * **[Job permissions](:service:jobs)** — Manage which users can view, manage, trigger, cancel, + or own a job. + + * **[MLflow experiment permissions](:service:experiments)** — Manage which users can read, + edit, or manage MLflow experiments. + + * **[MLflow registered model permissions](:service:modelregistry)** — Manage which users can + read, edit, or manage MLflow registered models. + + * **[Password permissions](:service:users)** — Manage which users can use password login when + SSO is enabled. + + * **[Instance Pool permissions](:service:instancepools)** — Manage which users can manage or + attach to pools. + + * **[Repo permissions](repos)** — Manage which users can read, run, edit, or manage a repo. + + * **[Serving endpoint permissions](:service:servingendpoints)** — Manage which users can view, + query, or manage a serving endpoint. + + * **[SQL warehouse permissions](:service:warehouses)** — Manage which users can use or manage + SQL warehouses. + + * **[Token permissions](:service:tokenmanagement)** — Manage which users can create or use + tokens. + + * **[Workspace object permissions](:service:workspace)** — Manage which users can read, run, + edit, or manage alerts, dbsql-dashboards, directories, files, notebooks and queries. + + For the mapping of the required permissions for specific actions or abilities and other + important information, see [Access Control]. + + Note that to manage access control on service principals, use **[Account Access Control + Proxy](:service:accountaccesscontrolproxy)**. + + [Access Control]: https://docs.databricks.com/security/auth-authz/access-control/index.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ServicePrincipalsClient(ServicePrincipalsAPI): + """ + Identities for use with jobs, automated tools, and systems such as scripts, apps, and CI/CD + platforms. Databricks recommends creating service principals to run production jobs or modify + production data. If all processes that act on production data run with service principals, + interactive users do not need any write, delete, or modify privileges in production. This + eliminates the risk of a user overwriting production data by accident. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class UsersClient(UsersAPI): + """ + User identities recognized by Databricks and represented by email addresses. + + Databricks recommends using SCIM provisioning to sync users and groups automatically from your + identity provider to your Databricks workspace. SCIM streamlines onboarding a new employee or + team by using your identity provider to create users and groups in Databricks workspace and give + them the proper level of access. When a user leaves your organization or no longer needs access + to Databricks workspace, admins can terminate the user in your identity provider and that + user’s account will also be removed from Databricks workspace. This ensures a consistent + offboarding process and prevents unauthorized users from accessing sensitive data. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WorkspaceAssignmentClient(WorkspaceAssignmentAPI): + """ + The Workspace Permission Assignment API allows you to manage workspace permissions for + principals in your account. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/iam.py b/databricks/sdk/iam/v2/iam.py similarity index 99% rename from databricks/sdk/service/iam.py rename to databricks/sdk/iam/v2/iam.py index d5fe5645e..132580aa9 100755 --- a/databricks/sdk/service/iam.py +++ b/databricks/sdk/iam/v2/iam.py @@ -7,11 +7,11 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order diff --git a/databricks/sdk/jobs/__init__.py b/databricks/sdk/jobs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/jobs/v2/__init__.py b/databricks/sdk/jobs/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/jobs/v2/client.py b/databricks/sdk/jobs/v2/client.py new file mode 100755 index 000000000..8a7ef9d23 --- /dev/null +++ b/databricks/sdk/jobs/v2/client.py @@ -0,0 +1,170 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .jobs import PolicyComplianceForJobsAPI +from .mixin import JobsExt + +_LOG = logging.getLogger(__name__) + + +class JobsClient(JobsExt): + """ + The Jobs API allows you to create, edit, and delete jobs. + + You can use a Databricks job to run a data processing or data analysis task in a Databricks + cluster with scalable resources. Your job can consist of a single task or can be a large, + multi-task workflow with complex dependencies. Databricks manages the task orchestration, + cluster management, monitoring, and error reporting for all of your jobs. You can run your jobs + immediately or periodically through an easy-to-use scheduling system. You can implement job + tasks using notebooks, JARS, Delta Live Tables pipelines, or Python, Scala, Spark submit, and + Java applications. + + You should never hard code secrets or store them in plain text. Use the [Secrets CLI] to manage + secrets in the [Databricks CLI]. Use the [Secrets utility] to reference secrets in notebooks and + jobs. + + [Databricks CLI]: https://docs.databricks.com/dev-tools/cli/index.html + [Secrets CLI]: https://docs.databricks.com/dev-tools/cli/secrets-cli.html + [Secrets utility]: https://docs.databricks.com/dev-tools/databricks-utils.html#dbutils-secrets + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PolicyComplianceForJobsClient(PolicyComplianceForJobsAPI): + """ + The compliance APIs allow you to view and manage the policy compliance status of jobs in your + workspace. This API currently only supports compliance controls for cluster policies. + + A job is in compliance if its cluster configurations satisfy the rules of all their respective + cluster policies. A job could be out of compliance if a cluster policy it uses was updated after + the job was last edited. The job is considered out of compliance if any of its clusters no + longer comply with their updated policies. + + The get and list compliance APIs allow you to view the policy compliance status of a job. The + enforce compliance API allows you to update a job so that it becomes compliant with all of its + policies. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/jobs/v2/jobs.py similarity index 82% rename from databricks/sdk/service/jobs.py rename to databricks/sdk/jobs/v2/jobs.py index 5885ace01..76aeaac34 100755 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/jobs/v2/jobs.py @@ -3,22 +3,312 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") +# all definitions in this file are in alphabetical order -from databricks.sdk.service import compute -# all definitions in this file are in alphabetical order +@dataclass +class Adlsgen2Info: + """A storage location in Adls Gen2""" + + destination: str + """abfss destination, e.g. + `abfss://@.dfs.core.windows.net/`.""" + + def as_dict(self) -> dict: + """Serializes the Adlsgen2Info into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Adlsgen2Info into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Adlsgen2Info: + """Deserializes the Adlsgen2Info from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +class AuthenticationMethod(Enum): + + OAUTH = "OAUTH" + PAT = "PAT" + + +@dataclass +class AutoScale: + max_workers: Optional[int] = None + """The maximum number of workers to which the cluster can scale up when overloaded. Note that + `max_workers` must be strictly greater than `min_workers`.""" + + min_workers: Optional[int] = None + """The minimum number of workers to which the cluster can scale down when underutilized. It is also + the initial number of workers the cluster will have after creation.""" + + def as_dict(self) -> dict: + """Serializes the AutoScale into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.max_workers is not None: + body["max_workers"] = self.max_workers + if self.min_workers is not None: + body["min_workers"] = self.min_workers + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AutoScale into a shallow dictionary of its immediate attributes.""" + body = {} + if self.max_workers is not None: + body["max_workers"] = self.max_workers + if self.min_workers is not None: + body["min_workers"] = self.min_workers + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AutoScale: + """Deserializes the AutoScale from a dictionary.""" + return cls(max_workers=d.get("max_workers", None), min_workers=d.get("min_workers", None)) + + +@dataclass +class AwsAttributes: + """Attributes set during cluster creation which are related to Amazon Web Services.""" + + availability: Optional[AwsAvailability] = None + """Availability type used for all subsequent nodes past the `first_on_demand` ones. + + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ebs_volume_count: Optional[int] = None + """The number of volumes launched for each instance. Users can choose up to 10 volumes. This + feature is only enabled for supported node types. Legacy node types cannot specify custom EBS + volumes. For node types with no instance store, at least one EBS volume needs to be specified; + otherwise, cluster creation will fail. + + These EBS volumes will be mounted at `/ebs0`, `/ebs1`, and etc. Instance store volumes will be + mounted at `/local_disk0`, `/local_disk1`, and etc. + + If EBS volumes are attached, Databricks will configure Spark to use only the EBS volumes for + scratch storage because heterogenously sized scratch devices can lead to inefficient disk + utilization. If no EBS volumes are attached, Databricks will configure Spark to use instance + store volumes. + + Please note that if EBS volumes are specified, then the Spark configuration `spark.local.dir` + will be overridden.""" + + ebs_volume_iops: Optional[int] = None + """If using gp3 volumes, what IOPS to use for the disk. If this is not set, the maximum performance + of a gp2 volume with the same volume size will be used.""" + + ebs_volume_size: Optional[int] = None + """The size of each EBS volume (in GiB) launched for each instance. For general purpose SSD, this + value must be within the range 100 - 4096. For throughput optimized HDD, this value must be + within the range 500 - 4096.""" + + ebs_volume_throughput: Optional[int] = None + """If using gp3 volumes, what throughput to use for the disk. If this is not set, the maximum + performance of a gp2 volume with the same volume size will be used.""" + + ebs_volume_type: Optional[EbsVolumeType] = None + """The type of EBS volumes that will be launched with this cluster.""" + + first_on_demand: Optional[int] = None + """The first `first_on_demand` nodes of the cluster will be placed on on-demand instances. If this + value is greater than 0, the cluster driver node in particular will be placed on an on-demand + instance. If this value is greater than or equal to the current cluster size, all nodes will be + placed on on-demand instances. If this value is less than the current cluster size, + `first_on_demand` nodes will be placed on on-demand instances and the remainder will be placed + on `availability` instances. Note that this value does not affect cluster size and cannot + currently be mutated over the lifetime of a cluster.""" + + instance_profile_arn: Optional[str] = None + """Nodes for this cluster will only be placed on AWS instances with this instance profile. If + ommitted, nodes will be placed on instances without an IAM instance profile. The instance + profile must have previously been added to the Databricks environment by an account + administrator. + + This feature may only be available to certain customer plans.""" + + spot_bid_price_percent: Optional[int] = None + """The bid price for AWS spot instances, as a percentage of the corresponding instance type's + on-demand price. For example, if this field is set to 50, and the cluster needs a new + `r3.xlarge` spot instance, then the bid price is half of the price of on-demand `r3.xlarge` + instances. Similarly, if this field is set to 200, the bid price is twice the price of on-demand + `r3.xlarge` instances. If not specified, the default value is 100. When spot instances are + requested for this cluster, only spot instances whose bid price percentage matches this field + will be considered. Note that, for safety, we enforce this field to be no more than 10000.""" + + zone_id: Optional[str] = None + """Identifier for the availability zone/datacenter in which the cluster resides. This string will + be of a form like "us-west-2a". The provided availability zone must be in the same region as the + Databricks deployment. For example, "us-west-2a" is not a valid zone id if the Databricks + deployment resides in the "us-east-1" region. This is an optional field at cluster creation, and + if not specified, a default zone will be used. If the zone specified is "auto", will try to + place cluster in a zone with high availability, and will retry placement in a different AZ if + there is not enough capacity. + + The list of available zones as well as the default value can be found by using the `List Zones` + method.""" + + def as_dict(self) -> dict: + """Serializes the AwsAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.ebs_volume_count is not None: + body["ebs_volume_count"] = self.ebs_volume_count + if self.ebs_volume_iops is not None: + body["ebs_volume_iops"] = self.ebs_volume_iops + if self.ebs_volume_size is not None: + body["ebs_volume_size"] = self.ebs_volume_size + if self.ebs_volume_throughput is not None: + body["ebs_volume_throughput"] = self.ebs_volume_throughput + if self.ebs_volume_type is not None: + body["ebs_volume_type"] = self.ebs_volume_type.value + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.instance_profile_arn is not None: + body["instance_profile_arn"] = self.instance_profile_arn + if self.spot_bid_price_percent is not None: + body["spot_bid_price_percent"] = self.spot_bid_price_percent + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AwsAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.ebs_volume_count is not None: + body["ebs_volume_count"] = self.ebs_volume_count + if self.ebs_volume_iops is not None: + body["ebs_volume_iops"] = self.ebs_volume_iops + if self.ebs_volume_size is not None: + body["ebs_volume_size"] = self.ebs_volume_size + if self.ebs_volume_throughput is not None: + body["ebs_volume_throughput"] = self.ebs_volume_throughput + if self.ebs_volume_type is not None: + body["ebs_volume_type"] = self.ebs_volume_type + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.instance_profile_arn is not None: + body["instance_profile_arn"] = self.instance_profile_arn + if self.spot_bid_price_percent is not None: + body["spot_bid_price_percent"] = self.spot_bid_price_percent + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AwsAttributes: + """Deserializes the AwsAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", AwsAvailability), + ebs_volume_count=d.get("ebs_volume_count", None), + ebs_volume_iops=d.get("ebs_volume_iops", None), + ebs_volume_size=d.get("ebs_volume_size", None), + ebs_volume_throughput=d.get("ebs_volume_throughput", None), + ebs_volume_type=_enum(d, "ebs_volume_type", EbsVolumeType), + first_on_demand=d.get("first_on_demand", None), + instance_profile_arn=d.get("instance_profile_arn", None), + spot_bid_price_percent=d.get("spot_bid_price_percent", None), + zone_id=d.get("zone_id", None), + ) + + +class AwsAvailability(Enum): + """Availability type used for all subsequent nodes past the `first_on_demand` ones. + + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ON_DEMAND = "ON_DEMAND" + SPOT = "SPOT" + SPOT_WITH_FALLBACK = "SPOT_WITH_FALLBACK" + + +@dataclass +class AzureAttributes: + """Attributes set during cluster creation which are related to Microsoft Azure.""" + + availability: Optional[AzureAvailability] = None + """Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If + `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + first_on_demand: Optional[int] = None + """The first `first_on_demand` nodes of the cluster will be placed on on-demand instances. This + value should be greater than 0, to make sure the cluster driver node is placed on an on-demand + instance. If this value is greater than or equal to the current cluster size, all nodes will be + placed on on-demand instances. If this value is less than the current cluster size, + `first_on_demand` nodes will be placed on on-demand instances and the remainder will be placed + on `availability` instances. Note that this value does not affect cluster size and cannot + currently be mutated over the lifetime of a cluster.""" + + log_analytics_info: Optional[LogAnalyticsInfo] = None + """Defines values necessary to configure and run Azure Log Analytics agent""" + + spot_bid_max_price: Optional[float] = None + """The max bid price to be used for Azure spot instances. The Max price for the bid cannot be + higher than the on-demand price of the instance. If not specified, the default value is -1, + which specifies that the instance cannot be evicted on the basis of price, and only on the basis + of availability. Further, the value should > 0 or -1.""" + + def as_dict(self) -> dict: + """Serializes the AzureAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.log_analytics_info: + body["log_analytics_info"] = self.log_analytics_info.as_dict() + if self.spot_bid_max_price is not None: + body["spot_bid_max_price"] = self.spot_bid_max_price + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AzureAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.log_analytics_info: + body["log_analytics_info"] = self.log_analytics_info + if self.spot_bid_max_price is not None: + body["spot_bid_max_price"] = self.spot_bid_max_price + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AzureAttributes: + """Deserializes the AzureAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", AzureAvailability), + first_on_demand=d.get("first_on_demand", None), + log_analytics_info=_from_dict(d, "log_analytics_info", LogAnalyticsInfo), + spot_bid_max_price=d.get("spot_bid_max_price", None), + ) + + +class AzureAvailability(Enum): + """Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If + `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ON_DEMAND_AZURE = "ON_DEMAND_AZURE" + SPOT_AZURE = "SPOT_AZURE" + SPOT_WITH_FALLBACK_AZURE = "SPOT_WITH_FALLBACK_AZURE" @dataclass @@ -37,9 +327,9 @@ class BaseJob: on accessible budget policies of the run_as identity on job creation or modification.""" has_more: Optional[bool] = None - """Indicates if the job has more sub-resources (`tasks`, `job_clusters`) that are not shown. They - can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 :method:jobs/list - requests with `expand_tasks=true`.""" + """Indicates if the job has more array properties (`tasks`, `job_clusters`) that are not shown. + They can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 + :method:jobs/list requests with `expand_tasks=true`.""" job_id: Optional[int] = None """The canonical identifier for this job.""" @@ -125,9 +415,9 @@ class BaseRun: """Description of the run""" effective_performance_target: Optional[PerformanceTarget] = None - """effective_performance_target is the actual performance target used by the run during execution. - effective_performance_target can differ from the client-set performance_target depending on if - the job was eligible to be cost-optimized.""" + """The actual performance target used by the serverless run during execution. This can differ from + the client-set performance target on the request depending on whether the performance mode is + supported by the job type.""" end_time: Optional[int] = None """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This @@ -151,8 +441,8 @@ class BaseRun: are used, `git_source` must be defined on the job.""" has_more: Optional[bool] = None - """Indicates if the run has more sub-resources (`tasks`, `job_clusters`) that are not shown. They - can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 + """Indicates if the run has more array properties (`tasks`, `job_clusters`) that are not shown. + They can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 :method:jobs/listruns requests with `expand_tasks=true`.""" job_clusters: Optional[List[JobCluster]] = None @@ -693,6 +983,38 @@ def from_dict(cls, d: Dict[str, Any]) -> CleanRoomsNotebookTaskCleanRoomsNoteboo ) +@dataclass +class ClientsTypes: + jobs: Optional[bool] = None + """With jobs set, the cluster can be used for jobs""" + + notebooks: Optional[bool] = None + """With notebooks set, this cluster can be used for notebooks""" + + def as_dict(self) -> dict: + """Serializes the ClientsTypes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.jobs is not None: + body["jobs"] = self.jobs + if self.notebooks is not None: + body["notebooks"] = self.notebooks + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ClientsTypes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.jobs is not None: + body["jobs"] = self.jobs + if self.notebooks is not None: + body["notebooks"] = self.notebooks + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ClientsTypes: + """Deserializes the ClientsTypes from a dictionary.""" + return cls(jobs=d.get("jobs", None), notebooks=d.get("notebooks", None)) + + @dataclass class ClusterInstance: cluster_id: Optional[str] = None @@ -736,6 +1058,56 @@ def from_dict(cls, d: Dict[str, Any]) -> ClusterInstance: return cls(cluster_id=d.get("cluster_id", None), spark_context_id=d.get("spark_context_id", None)) +@dataclass +class ClusterLogConf: + """Cluster log delivery config""" + + dbfs: Optional[DbfsStorageInfo] = None + """destination needs to be provided. e.g. `{ "dbfs" : { "destination" : "dbfs:/home/cluster_log" } + }`""" + + s3: Optional[S3StorageInfo] = None + """destination and either the region or endpoint need to be provided. e.g. `{ "s3": { "destination" + : "s3://cluster_log_bucket/prefix", "region" : "us-west-2" } }` Cluster iam role is used to + access s3, please make sure the cluster iam role in `instance_profile_arn` has permission to + write data to the s3 destination.""" + + volumes: Optional[VolumesStorageInfo] = None + """destination needs to be provided, e.g. `{ "volumes": { "destination": + "/Volumes/catalog/schema/volume/cluster_log" } }`""" + + def as_dict(self) -> dict: + """Serializes the ClusterLogConf into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dbfs: + body["dbfs"] = self.dbfs.as_dict() + if self.s3: + body["s3"] = self.s3.as_dict() + if self.volumes: + body["volumes"] = self.volumes.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ClusterLogConf into a shallow dictionary of its immediate attributes.""" + body = {} + if self.dbfs: + body["dbfs"] = self.dbfs + if self.s3: + body["s3"] = self.s3 + if self.volumes: + body["volumes"] = self.volumes + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ClusterLogConf: + """Deserializes the ClusterLogConf from a dictionary.""" + return cls( + dbfs=_from_dict(d, "dbfs", DbfsStorageInfo), + s3=_from_dict(d, "s3", S3StorageInfo), + volumes=_from_dict(d, "volumes", VolumesStorageInfo), + ) + + @dataclass class ClusterSpec: existing_cluster_id: Optional[str] = None @@ -747,11 +1119,11 @@ class ClusterSpec: """If job_cluster_key, this task is executed reusing the cluster specified in `job.settings.job_clusters`.""" - libraries: Optional[List[compute.Library]] = None + libraries: Optional[List[Library]] = None """An optional list of libraries to be installed on the cluster. The default value is an empty list.""" - new_cluster: Optional[compute.ClusterSpec] = None + new_cluster: Optional[ClusterSpec] = None """If new_cluster, a description of a new cluster that is created for each run.""" def as_dict(self) -> dict: @@ -786,15 +1158,13 @@ def from_dict(cls, d: Dict[str, Any]) -> ClusterSpec: return cls( existing_cluster_id=d.get("existing_cluster_id", None), job_cluster_key=d.get("job_cluster_key", None), - libraries=_repeated_dict(d, "libraries", compute.Library), - new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + libraries=_repeated_dict(d, "libraries", Library), + new_cluster=_from_dict(d, "new_cluster", ClusterSpec), ) @dataclass class ComputeConfig: - """Next field: 4""" - num_gpus: int """Number of GPUs.""" @@ -1016,8 +1386,8 @@ class CreateJob: """Job-level parameter definitions""" performance_target: Optional[PerformanceTarget] = None - """PerformanceTarget defines how performant or cost efficient the execution of run on serverless - should be.""" + """The performance mode on a serverless job. The performance target determines the level of compute + performance or cost-efficiency for the run.""" queue: Optional[QueueSettings] = None """The queue settings of the job.""" @@ -1271,6 +1641,172 @@ def from_dict(cls, d: Dict[str, Any]) -> CronSchedule: ) +@dataclass +class DashboardPageSnapshot: + page_display_name: Optional[str] = None + + widget_error_details: Optional[List[WidgetErrorDetail]] = None + + def as_dict(self) -> dict: + """Serializes the DashboardPageSnapshot into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.page_display_name is not None: + body["page_display_name"] = self.page_display_name + if self.widget_error_details: + body["widget_error_details"] = [v.as_dict() for v in self.widget_error_details] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DashboardPageSnapshot into a shallow dictionary of its immediate attributes.""" + body = {} + if self.page_display_name is not None: + body["page_display_name"] = self.page_display_name + if self.widget_error_details: + body["widget_error_details"] = self.widget_error_details + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DashboardPageSnapshot: + """Deserializes the DashboardPageSnapshot from a dictionary.""" + return cls( + page_display_name=d.get("page_display_name", None), + widget_error_details=_repeated_dict(d, "widget_error_details", WidgetErrorDetail), + ) + + +@dataclass +class DashboardTask: + """Configures the Lakeview Dashboard job task type.""" + + dashboard_id: Optional[str] = None + + subscription: Optional[Subscription] = None + + warehouse_id: Optional[str] = None + """The warehouse id to execute the dashboard with for the schedule""" + + def as_dict(self) -> dict: + """Serializes the DashboardTask into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dashboard_id is not None: + body["dashboard_id"] = self.dashboard_id + if self.subscription: + body["subscription"] = self.subscription.as_dict() + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DashboardTask into a shallow dictionary of its immediate attributes.""" + body = {} + if self.dashboard_id is not None: + body["dashboard_id"] = self.dashboard_id + if self.subscription: + body["subscription"] = self.subscription + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DashboardTask: + """Deserializes the DashboardTask from a dictionary.""" + return cls( + dashboard_id=d.get("dashboard_id", None), + subscription=_from_dict(d, "subscription", Subscription), + warehouse_id=d.get("warehouse_id", None), + ) + + +@dataclass +class DashboardTaskOutput: + page_snapshots: Optional[List[DashboardPageSnapshot]] = None + """Should only be populated for manual PDF download jobs.""" + + def as_dict(self) -> dict: + """Serializes the DashboardTaskOutput into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.page_snapshots: + body["page_snapshots"] = [v.as_dict() for v in self.page_snapshots] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DashboardTaskOutput into a shallow dictionary of its immediate attributes.""" + body = {} + if self.page_snapshots: + body["page_snapshots"] = self.page_snapshots + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DashboardTaskOutput: + """Deserializes the DashboardTaskOutput from a dictionary.""" + return cls(page_snapshots=_repeated_dict(d, "page_snapshots", DashboardPageSnapshot)) + + +class DataSecurityMode(Enum): + """Data security mode decides what data governance model to use when accessing data from a cluster. + + The following modes can only be used when `kind = CLASSIC_PREVIEW`. * `DATA_SECURITY_MODE_AUTO`: + Databricks will choose the most appropriate access mode depending on your compute configuration. + * `DATA_SECURITY_MODE_STANDARD`: Alias for `USER_ISOLATION`. * `DATA_SECURITY_MODE_DEDICATED`: + Alias for `SINGLE_USER`. + + The following modes can be used regardless of `kind`. * `NONE`: No security isolation for + multiple users sharing the cluster. Data governance features are not available in this mode. * + `SINGLE_USER`: A secure cluster that can only be exclusively used by a single user specified in + `single_user_name`. Most programming languages, cluster features and data governance features + are available in this mode. * `USER_ISOLATION`: A secure cluster that can be shared by multiple + users. Cluster users are fully isolated so that they cannot see each other's data and + credentials. Most data governance features are supported in this mode. But programming languages + and cluster features might be limited. + + The following modes are deprecated starting with Databricks Runtime 15.0 and will be removed for + future Databricks Runtime versions: + + * `LEGACY_TABLE_ACL`: This mode is for users migrating from legacy Table ACL clusters. * + `LEGACY_PASSTHROUGH`: This mode is for users migrating from legacy Passthrough on high + concurrency clusters. * `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy + Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This mode provides a way that + doesn’t have UC nor passthrough enabled.""" + + DATA_SECURITY_MODE_AUTO = "DATA_SECURITY_MODE_AUTO" + DATA_SECURITY_MODE_DEDICATED = "DATA_SECURITY_MODE_DEDICATED" + DATA_SECURITY_MODE_STANDARD = "DATA_SECURITY_MODE_STANDARD" + LEGACY_PASSTHROUGH = "LEGACY_PASSTHROUGH" + LEGACY_SINGLE_USER = "LEGACY_SINGLE_USER" + LEGACY_SINGLE_USER_STANDARD = "LEGACY_SINGLE_USER_STANDARD" + LEGACY_TABLE_ACL = "LEGACY_TABLE_ACL" + NONE = "NONE" + SINGLE_USER = "SINGLE_USER" + USER_ISOLATION = "USER_ISOLATION" + + +@dataclass +class DbfsStorageInfo: + """A storage location in DBFS""" + + destination: str + """dbfs destination, e.g. `dbfs:/my/path`""" + + def as_dict(self) -> dict: + """Serializes the DbfsStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DbfsStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DbfsStorageInfo: + """Deserializes the DbfsStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + @dataclass class DbtOutput: artifacts_headers: Optional[Dict[str, str]] = None @@ -1479,6 +2015,78 @@ def from_dict(cls, d: Dict[str, Any]) -> DeleteRunResponse: return cls() +@dataclass +class DockerBasicAuth: + password: Optional[str] = None + """Password of the user""" + + username: Optional[str] = None + """Name of the user""" + + def as_dict(self) -> dict: + """Serializes the DockerBasicAuth into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.password is not None: + body["password"] = self.password + if self.username is not None: + body["username"] = self.username + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DockerBasicAuth into a shallow dictionary of its immediate attributes.""" + body = {} + if self.password is not None: + body["password"] = self.password + if self.username is not None: + body["username"] = self.username + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DockerBasicAuth: + """Deserializes the DockerBasicAuth from a dictionary.""" + return cls(password=d.get("password", None), username=d.get("username", None)) + + +@dataclass +class DockerImage: + basic_auth: Optional[DockerBasicAuth] = None + """Basic auth with username and password""" + + url: Optional[str] = None + """URL of the docker image.""" + + def as_dict(self) -> dict: + """Serializes the DockerImage into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.basic_auth: + body["basic_auth"] = self.basic_auth.as_dict() + if self.url is not None: + body["url"] = self.url + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DockerImage into a shallow dictionary of its immediate attributes.""" + body = {} + if self.basic_auth: + body["basic_auth"] = self.basic_auth + if self.url is not None: + body["url"] = self.url + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DockerImage: + """Deserializes the DockerImage from a dictionary.""" + return cls(basic_auth=_from_dict(d, "basic_auth", DockerBasicAuth), url=d.get("url", None)) + + +class EbsVolumeType(Enum): + """All EBS volume types that Databricks supports. See https://aws.amazon.com/ebs/details/ for + details.""" + + GENERAL_PURPOSE_SSD = "GENERAL_PURPOSE_SSD" + THROUGHPUT_OPTIMIZED_HDD = "THROUGHPUT_OPTIMIZED_HDD" + + @dataclass class EnforcePolicyComplianceForJobResponseJobClusterSettingsChange: """Represents a change to the job cluster's settings that would be required for the job clusters to @@ -1609,6 +2217,60 @@ def from_dict(cls, d: Dict[str, Any]) -> EnforcePolicyComplianceResponse: ) +@dataclass +class Environment: + """The environment entity used to preserve serverless environment side panel and jobs' environment + for non-notebook task. In this minimal environment spec, only pip dependencies are supported.""" + + client: str + """Client version used by the environment The client is the user-facing environment of the runtime. + Each client comes with a specific set of pre-installed libraries. The version is a string, + consisting of the major client version.""" + + dependencies: Optional[List[str]] = None + """List of pip dependencies, as supported by the version of pip in this environment. Each + dependency is a pip requirement file line + https://pip.pypa.io/en/stable/reference/requirements-file-format/ Allowed dependency could be + , , (WSFS or Volumes in + Databricks), E.g. dependencies: ["foo==0.0.1", "-r + /Workspace/test/requirements.txt"]""" + + jar_dependencies: Optional[List[str]] = None + """List of jar dependencies, should be string representing volume paths. For example: + `/Volumes/path/to/test.jar`.""" + + def as_dict(self) -> dict: + """Serializes the Environment into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.client is not None: + body["client"] = self.client + if self.dependencies: + body["dependencies"] = [v for v in self.dependencies] + if self.jar_dependencies: + body["jar_dependencies"] = [v for v in self.jar_dependencies] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Environment into a shallow dictionary of its immediate attributes.""" + body = {} + if self.client is not None: + body["client"] = self.client + if self.dependencies: + body["dependencies"] = self.dependencies + if self.jar_dependencies: + body["jar_dependencies"] = self.jar_dependencies + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Environment: + """Deserializes the Environment from a dictionary.""" + return cls( + client=d.get("client", None), + dependencies=d.get("dependencies", None), + jar_dependencies=d.get("jar_dependencies", None), + ) + + @dataclass class ExportRunOutput: """Run was exported successfully.""" @@ -1880,40 +2542,155 @@ class Format(Enum): @dataclass -class GenAiComputeTask: - """Next field: 9""" +class GcpAttributes: + """Attributes set during cluster creation which are related to GCP.""" - dl_runtime_image: str - """Runtime image""" + availability: Optional[GcpAvailability] = None + """This field determines whether the spark executors will be scheduled to run on preemptible VMs, + on-demand VMs, or preemptible VMs with a fallback to on-demand VMs if the former is unavailable.""" - command: Optional[str] = None - """Command launcher to run the actual script, e.g. bash, python etc.""" + boot_disk_size: Optional[int] = None + """Boot disk size in GB""" - compute: Optional[ComputeConfig] = None - """Next field: 4""" + google_service_account: Optional[str] = None + """If provided, the cluster will impersonate the google service account when accessing gcloud + services (like GCS). The google service account must have previously been added to the + Databricks environment by an account administrator.""" - mlflow_experiment_name: Optional[str] = None - """Optional string containing the name of the MLflow experiment to log the run to. If name is not - found, backend will create the mlflow experiment using the name.""" + local_ssd_count: Optional[int] = None + """If provided, each node (workers and driver) in the cluster will have this number of local SSDs + attached. Each local SSD is 375GB in size. Refer to [GCP documentation] for the supported number + of local SSDs for each instance type. + + [GCP documentation]: https://cloud.google.com/compute/docs/disks/local-ssd#choose_number_local_ssds""" - source: Optional[Source] = None - """Optional location type of the training script. When set to `WORKSPACE`, the script will be - retrieved from the local Databricks workspace. When set to `GIT`, the script will be retrieved - from a Git repository defined in `git_source`. If the value is empty, the task will use `GIT` if - `git_source` is defined and `WORKSPACE` otherwise. * `WORKSPACE`: Script is located in - Databricks workspace. * `GIT`: Script is located in cloud Git provider.""" + use_preemptible_executors: Optional[bool] = None + """This field determines whether the spark executors will be scheduled to run on preemptible VMs + (when set to true) versus standard compute engine VMs (when set to false; default). Note: Soon + to be deprecated, use the 'availability' field instead.""" - training_script_path: Optional[str] = None - """The training script file path to be executed. Cloud file URIs (such as dbfs:/, s3:/, adls:/, - gcs:/) and workspace paths are supported. For python files stored in the Databricks workspace, - the path must be absolute and begin with `/`. For files stored in a remote repository, the path - must be relative. This field is required.""" + zone_id: Optional[str] = None + """Identifier for the availability zone in which the cluster resides. This can be one of the + following: - "HA" => High availability, spread nodes across availability zones for a Databricks + deployment region [default]. - "AUTO" => Databricks picks an availability zone to schedule the + cluster on. - A GCP availability zone => Pick One of the available zones for (machine type + + region) from https://cloud.google.com/compute/docs/regions-zones.""" - yaml_parameters: Optional[str] = None - """Optional string containing model parameters passed to the training script in yaml format. If - present, then the content in yaml_parameters_file_path will be ignored.""" + def as_dict(self) -> dict: + """Serializes the GcpAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.boot_disk_size is not None: + body["boot_disk_size"] = self.boot_disk_size + if self.google_service_account is not None: + body["google_service_account"] = self.google_service_account + if self.local_ssd_count is not None: + body["local_ssd_count"] = self.local_ssd_count + if self.use_preemptible_executors is not None: + body["use_preemptible_executors"] = self.use_preemptible_executors + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body - yaml_parameters_file_path: Optional[str] = None + def as_shallow_dict(self) -> dict: + """Serializes the GcpAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.boot_disk_size is not None: + body["boot_disk_size"] = self.boot_disk_size + if self.google_service_account is not None: + body["google_service_account"] = self.google_service_account + if self.local_ssd_count is not None: + body["local_ssd_count"] = self.local_ssd_count + if self.use_preemptible_executors is not None: + body["use_preemptible_executors"] = self.use_preemptible_executors + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GcpAttributes: + """Deserializes the GcpAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", GcpAvailability), + boot_disk_size=d.get("boot_disk_size", None), + google_service_account=d.get("google_service_account", None), + local_ssd_count=d.get("local_ssd_count", None), + use_preemptible_executors=d.get("use_preemptible_executors", None), + zone_id=d.get("zone_id", None), + ) + + +class GcpAvailability(Enum): + """This field determines whether the instance pool will contain preemptible VMs, on-demand VMs, or + preemptible VMs with a fallback to on-demand VMs if the former is unavailable.""" + + ON_DEMAND_GCP = "ON_DEMAND_GCP" + PREEMPTIBLE_GCP = "PREEMPTIBLE_GCP" + PREEMPTIBLE_WITH_FALLBACK_GCP = "PREEMPTIBLE_WITH_FALLBACK_GCP" + + +@dataclass +class GcsStorageInfo: + """A storage location in Google Cloud Platform's GCS""" + + destination: str + """GCS destination/URI, e.g. `gs://my-bucket/some-prefix`""" + + def as_dict(self) -> dict: + """Serializes the GcsStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GcsStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GcsStorageInfo: + """Deserializes the GcsStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +@dataclass +class GenAiComputeTask: + dl_runtime_image: str + """Runtime image""" + + command: Optional[str] = None + """Command launcher to run the actual script, e.g. bash, python etc.""" + + compute: Optional[ComputeConfig] = None + + mlflow_experiment_name: Optional[str] = None + """Optional string containing the name of the MLflow experiment to log the run to. If name is not + found, backend will create the mlflow experiment using the name.""" + + source: Optional[Source] = None + """Optional location type of the training script. When set to `WORKSPACE`, the script will be + retrieved from the local Databricks workspace. When set to `GIT`, the script will be retrieved + from a Git repository defined in `git_source`. If the value is empty, the task will use `GIT` if + `git_source` is defined and `WORKSPACE` otherwise. * `WORKSPACE`: Script is located in + Databricks workspace. * `GIT`: Script is located in cloud Git provider.""" + + training_script_path: Optional[str] = None + """The training script file path to be executed. Cloud file URIs (such as dbfs:/, s3:/, adls:/, + gcs:/) and workspace paths are supported. For python files stored in the Databricks workspace, + the path must be absolute and begin with `/`. For files stored in a remote repository, the path + must be relative. This field is required.""" + + yaml_parameters: Optional[str] = None + """Optional string containing model parameters passed to the training script in yaml format. If + present, then the content in yaml_parameters_file_path will be ignored.""" + + yaml_parameters_file_path: Optional[str] = None """Optional path to a YAML file containing model parameters passed to the training script.""" def as_dict(self) -> dict: @@ -2165,6 +2942,90 @@ def from_dict(cls, d: Dict[str, Any]) -> GitSource: ) +@dataclass +class InitScriptInfo: + """Config for an individual init script Next ID: 11""" + + abfss: Optional[Adlsgen2Info] = None + """destination needs to be provided, e.g. + `abfss://@.dfs.core.windows.net/`""" + + dbfs: Optional[DbfsStorageInfo] = None + """destination needs to be provided. e.g. `{ "dbfs": { "destination" : "dbfs:/home/cluster_log" } + }`""" + + file: Optional[LocalFileInfo] = None + """destination needs to be provided, e.g. `{ "file": { "destination": "file:/my/local/file.sh" } }`""" + + gcs: Optional[GcsStorageInfo] = None + """destination needs to be provided, e.g. `{ "gcs": { "destination": "gs://my-bucket/file.sh" } }`""" + + s3: Optional[S3StorageInfo] = None + """destination and either the region or endpoint need to be provided. e.g. `{ \"s3\": { + \"destination\": \"s3://cluster_log_bucket/prefix\", \"region\": \"us-west-2\" } }` Cluster iam + role is used to access s3, please make sure the cluster iam role in `instance_profile_arn` has + permission to write data to the s3 destination.""" + + volumes: Optional[VolumesStorageInfo] = None + """destination needs to be provided. e.g. `{ \"volumes\" : { \"destination\" : + \"/Volumes/my-init.sh\" } }`""" + + workspace: Optional[WorkspaceStorageInfo] = None + """destination needs to be provided, e.g. `{ "workspace": { "destination": + "/cluster-init-scripts/setup-datadog.sh" } }`""" + + def as_dict(self) -> dict: + """Serializes the InitScriptInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.abfss: + body["abfss"] = self.abfss.as_dict() + if self.dbfs: + body["dbfs"] = self.dbfs.as_dict() + if self.file: + body["file"] = self.file.as_dict() + if self.gcs: + body["gcs"] = self.gcs.as_dict() + if self.s3: + body["s3"] = self.s3.as_dict() + if self.volumes: + body["volumes"] = self.volumes.as_dict() + if self.workspace: + body["workspace"] = self.workspace.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the InitScriptInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.abfss: + body["abfss"] = self.abfss + if self.dbfs: + body["dbfs"] = self.dbfs + if self.file: + body["file"] = self.file + if self.gcs: + body["gcs"] = self.gcs + if self.s3: + body["s3"] = self.s3 + if self.volumes: + body["volumes"] = self.volumes + if self.workspace: + body["workspace"] = self.workspace + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> InitScriptInfo: + """Deserializes the InitScriptInfo from a dictionary.""" + return cls( + abfss=_from_dict(d, "abfss", Adlsgen2Info), + dbfs=_from_dict(d, "dbfs", DbfsStorageInfo), + file=_from_dict(d, "file", LocalFileInfo), + gcs=_from_dict(d, "gcs", GcsStorageInfo), + s3=_from_dict(d, "s3", S3StorageInfo), + volumes=_from_dict(d, "volumes", VolumesStorageInfo), + workspace=_from_dict(d, "workspace", WorkspaceStorageInfo), + ) + + @dataclass class Job: """Job was retrieved successfully.""" @@ -2183,15 +3044,15 @@ class Job: on accessible budget policies of the run_as identity on job creation or modification.""" has_more: Optional[bool] = None - """Indicates if the job has more sub-resources (`tasks`, `job_clusters`) that are not shown. They - can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 :method:jobs/list - requests with `expand_tasks=true`.""" + """Indicates if the job has more array properties (`tasks`, `job_clusters`) that are not shown. + They can be accessed via :method:jobs/get endpoint. It is only relevant for API 2.2 + :method:jobs/list requests with `expand_tasks=true`.""" job_id: Optional[int] = None """The canonical identifier for this job.""" next_page_token: Optional[str] = None - """A token that can be used to list the next page of sub-resources.""" + """A token that can be used to list the next page of array properties.""" run_as_user_name: Optional[str] = None """The email of an active workspace user or the application ID of a service principal that the job @@ -2380,7 +3241,7 @@ class JobCluster: `JobTaskSettings` may refer to this field to determine which cluster to launch for the task execution.""" - new_cluster: compute.ClusterSpec + new_cluster: JobsClusterSpec """If new_cluster, a description of a cluster that is created for each task.""" def as_dict(self) -> dict: @@ -2405,8 +3266,7 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> JobCluster: """Deserializes the JobCluster from a dictionary.""" return cls( - job_cluster_key=d.get("job_cluster_key", None), - new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + job_cluster_key=d.get("job_cluster_key", None), new_cluster=_from_dict(d, "new_cluster", JobsClusterSpec) ) @@ -2591,7 +3451,7 @@ class JobEnvironment: environment_key: str """The key of an environment. It has to be unique within a job.""" - spec: Optional[compute.Environment] = None + spec: Optional[Environment] = None """The environment entity used to preserve serverless environment side panel and jobs' environment for non-notebook task. In this minimal environment spec, only pip dependencies are supported.""" @@ -2616,7 +3476,7 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> JobEnvironment: """Deserializes the JobEnvironment from a dictionary.""" - return cls(environment_key=d.get("environment_key", None), spec=_from_dict(d, "spec", compute.Environment)) + return cls(environment_key=d.get("environment_key", None), spec=_from_dict(d, "spec", Environment)) @dataclass @@ -3001,8 +3861,8 @@ class JobSettings: """Job-level parameter definitions""" performance_target: Optional[PerformanceTarget] = None - """PerformanceTarget defines how performant or cost efficient the execution of run on serverless - should be.""" + """The performance mode on a serverless job. The performance target determines the level of compute + performance or cost-efficiency for the run.""" queue: Optional[QueueSettings] = None """The queue settings of the job.""" @@ -3241,6 +4101,369 @@ class JobSourceDirtyState(Enum): NOT_SYNCED = "NOT_SYNCED" +@dataclass +class JobsClusterSpec: + """Contains a snapshot of the latest user specified settings that were used to create/edit the + cluster.""" + + apply_policy_default_values: Optional[bool] = None + """When set to true, fixed and default values from the policy will be used for fields that are + omitted. When set to false, only fixed values from the policy will be applied.""" + + autoscale: Optional[AutoScale] = None + """Parameters needed in order to automatically scale clusters up and down based on load. Note: + autoscaling works best with DB runtime versions 3.0 or later.""" + + autotermination_minutes: Optional[int] = None + """Automatically terminates the cluster after it is inactive for this time in minutes. If not set, + this cluster will not be automatically terminated. If specified, the threshold must be between + 10 and 10000 minutes. Users can also set this value to 0 to explicitly disable automatic + termination.""" + + aws_attributes: Optional[AwsAttributes] = None + """Attributes related to clusters running on Amazon Web Services. If not specified at cluster + creation, a set of default values will be used.""" + + azure_attributes: Optional[AzureAttributes] = None + """Attributes related to clusters running on Microsoft Azure. If not specified at cluster creation, + a set of default values will be used.""" + + cluster_log_conf: Optional[ClusterLogConf] = None + """The configuration for delivering spark logs to a long-term storage destination. Three kinds of + destinations (DBFS, S3 and Unity Catalog volumes) are supported. Only one destination can be + specified for one cluster. If the conf is given, the logs will be delivered to the destination + every `5 mins`. The destination of driver logs is `$destination/$clusterId/driver`, while the + destination of executor logs is `$destination/$clusterId/executor`.""" + + cluster_name: Optional[str] = None + """Cluster name requested by the user. This doesn't have to be unique. If not specified at + creation, the cluster name will be an empty string.""" + + custom_tags: Optional[Dict[str, str]] = None + """Additional tags for cluster resources. Databricks will tag all cluster resources (e.g., AWS + instances and EBS volumes) with these tags in addition to `default_tags`. Notes: + + - Currently, Databricks allows at most 45 custom tags + + - Clusters can only reuse cloud resources if the resources' tags are a subset of the cluster + tags""" + + data_security_mode: Optional[DataSecurityMode] = None + """Data security mode decides what data governance model to use when accessing data from a cluster. + + The following modes can only be used when `kind = CLASSIC_PREVIEW`. * `DATA_SECURITY_MODE_AUTO`: + Databricks will choose the most appropriate access mode depending on your compute configuration. + * `DATA_SECURITY_MODE_STANDARD`: Alias for `USER_ISOLATION`. * `DATA_SECURITY_MODE_DEDICATED`: + Alias for `SINGLE_USER`. + + The following modes can be used regardless of `kind`. * `NONE`: No security isolation for + multiple users sharing the cluster. Data governance features are not available in this mode. * + `SINGLE_USER`: A secure cluster that can only be exclusively used by a single user specified in + `single_user_name`. Most programming languages, cluster features and data governance features + are available in this mode. * `USER_ISOLATION`: A secure cluster that can be shared by multiple + users. Cluster users are fully isolated so that they cannot see each other's data and + credentials. Most data governance features are supported in this mode. But programming languages + and cluster features might be limited. + + The following modes are deprecated starting with Databricks Runtime 15.0 and will be removed for + future Databricks Runtime versions: + + * `LEGACY_TABLE_ACL`: This mode is for users migrating from legacy Table ACL clusters. * + `LEGACY_PASSTHROUGH`: This mode is for users migrating from legacy Passthrough on high + concurrency clusters. * `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy + Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This mode provides a way that + doesn’t have UC nor passthrough enabled.""" + + docker_image: Optional[DockerImage] = None + """Custom docker image BYOC""" + + driver_instance_pool_id: Optional[str] = None + """The optional ID of the instance pool for the driver of the cluster belongs. The pool cluster + uses the instance pool with id (instance_pool_id) if the driver pool is not assigned.""" + + driver_node_type_id: Optional[str] = None + """The node type of the Spark driver. Note that this field is optional; if unset, the driver node + type will be set as the same value as `node_type_id` defined above. + + This field, along with node_type_id, should not be set if virtual_cluster_size is set. If both + driver_node_type_id, node_type_id, and virtual_cluster_size are specified, driver_node_type_id + and node_type_id take precedence.""" + + enable_elastic_disk: Optional[bool] = None + """Autoscaling Local Storage: when enabled, this cluster will dynamically acquire additional disk + space when its Spark workers are running low on disk space. This feature requires specific AWS + permissions to function correctly - refer to the User Guide for more details.""" + + enable_local_disk_encryption: Optional[bool] = None + """Whether to enable LUKS on cluster VMs' local disks""" + + gcp_attributes: Optional[GcpAttributes] = None + """Attributes related to clusters running on Google Cloud Platform. If not specified at cluster + creation, a set of default values will be used.""" + + init_scripts: Optional[List[InitScriptInfo]] = None + """The configuration for storing init scripts. Any number of destinations can be specified. The + scripts are executed sequentially in the order provided. If `cluster_log_conf` is specified, + init script logs are sent to `//init_scripts`.""" + + instance_pool_id: Optional[str] = None + """The optional ID of the instance pool to which the cluster belongs.""" + + is_single_node: Optional[bool] = None + """This field can only be used when `kind = CLASSIC_PREVIEW`. + + When set to true, Databricks will automatically set single node related `custom_tags`, + `spark_conf`, and `num_workers`""" + + kind: Optional[Kind] = None + """The kind of compute described by this compute specification. + + Depending on `kind`, different validations and default values will be applied. + + Clusters with `kind = CLASSIC_PREVIEW` support the following fields, whereas clusters with no + specified `kind` do not. * [is_single_node](/api/workspace/clusters/create#is_single_node) * + [use_ml_runtime](/api/workspace/clusters/create#use_ml_runtime) * + [data_security_mode](/api/workspace/clusters/create#data_security_mode) set to + `DATA_SECURITY_MODE_AUTO`, `DATA_SECURITY_MODE_DEDICATED`, or `DATA_SECURITY_MODE_STANDARD` + + By using the [simple form], your clusters are automatically using `kind = CLASSIC_PREVIEW`. + + [simple form]: https://docs.databricks.com/compute/simple-form.html""" + + node_type_id: Optional[str] = None + """This field encodes, through a single value, the resources available to each of the Spark nodes + in this cluster. For example, the Spark nodes can be provisioned and optimized for memory or + compute intensive workloads. A list of available node types can be retrieved by using the + :method:clusters/listNodeTypes API call.""" + + num_workers: Optional[int] = None + """Number of worker nodes that this cluster should have. A cluster has one Spark Driver and + `num_workers` Executors for a total of `num_workers` + 1 Spark nodes. + + Note: When reading the properties of a cluster, this field reflects the desired number of + workers rather than the actual current number of workers. For instance, if a cluster is resized + from 5 to 10 workers, this field will immediately be updated to reflect the target size of 10 + workers, whereas the workers listed in `spark_info` will gradually increase from 5 to 10 as the + new nodes are provisioned.""" + + policy_id: Optional[str] = None + """The ID of the cluster policy used to create the cluster if applicable.""" + + runtime_engine: Optional[RuntimeEngine] = None + """Determines the cluster's runtime engine, either standard or Photon. + + This field is not compatible with legacy `spark_version` values that contain `-photon-`. Remove + `-photon-` from the `spark_version` and set `runtime_engine` to `PHOTON`. + + If left unspecified, the runtime engine defaults to standard unless the spark_version contains + -photon-, in which case Photon will be used.""" + + single_user_name: Optional[str] = None + """Single user name if data_security_mode is `SINGLE_USER`""" + + spark_conf: Optional[Dict[str, str]] = None + """An object containing a set of optional, user-specified Spark configuration key-value pairs. + Users can also pass in a string of extra JVM options to the driver and the executors via + `spark.driver.extraJavaOptions` and `spark.executor.extraJavaOptions` respectively.""" + + spark_env_vars: Optional[Dict[str, str]] = None + """An object containing a set of optional, user-specified environment variable key-value pairs. + Please note that key-value pair of the form (X,Y) will be exported as is (i.e., `export X='Y'`) + while launching the driver and workers. + + In order to specify an additional set of `SPARK_DAEMON_JAVA_OPTS`, we recommend appending them + to `$SPARK_DAEMON_JAVA_OPTS` as shown in the example below. This ensures that all default + databricks managed environmental variables are included as well. + + Example Spark environment variables: `{"SPARK_WORKER_MEMORY": "28000m", "SPARK_LOCAL_DIRS": + "/local_disk0"}` or `{"SPARK_DAEMON_JAVA_OPTS": "$SPARK_DAEMON_JAVA_OPTS + -Dspark.shuffle.service.enabled=true"}`""" + + spark_version: Optional[str] = None + """The Spark version of the cluster, e.g. `3.3.x-scala2.11`. A list of available Spark versions can + be retrieved by using the :method:clusters/sparkVersions API call.""" + + ssh_public_keys: Optional[List[str]] = None + """SSH public key contents that will be added to each Spark node in this cluster. The corresponding + private keys can be used to login with the user name `ubuntu` on port `2200`. Up to 10 keys can + be specified.""" + + use_ml_runtime: Optional[bool] = None + """This field can only be used when `kind = CLASSIC_PREVIEW`. + + `effective_spark_version` is determined by `spark_version` (DBR release), this field + `use_ml_runtime`, and whether `node_type_id` is gpu node or not.""" + + workload_type: Optional[WorkloadType] = None + """Cluster Attributes showing for clusters workload types.""" + + def as_dict(self) -> dict: + """Serializes the JobsClusterSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.apply_policy_default_values is not None: + body["apply_policy_default_values"] = self.apply_policy_default_values + if self.autoscale: + body["autoscale"] = self.autoscale.as_dict() + if self.autotermination_minutes is not None: + body["autotermination_minutes"] = self.autotermination_minutes + if self.aws_attributes: + body["aws_attributes"] = self.aws_attributes.as_dict() + if self.azure_attributes: + body["azure_attributes"] = self.azure_attributes.as_dict() + if self.cluster_log_conf: + body["cluster_log_conf"] = self.cluster_log_conf.as_dict() + if self.cluster_name is not None: + body["cluster_name"] = self.cluster_name + if self.custom_tags: + body["custom_tags"] = self.custom_tags + if self.data_security_mode is not None: + body["data_security_mode"] = self.data_security_mode.value + if self.docker_image: + body["docker_image"] = self.docker_image.as_dict() + if self.driver_instance_pool_id is not None: + body["driver_instance_pool_id"] = self.driver_instance_pool_id + if self.driver_node_type_id is not None: + body["driver_node_type_id"] = self.driver_node_type_id + if self.enable_elastic_disk is not None: + body["enable_elastic_disk"] = self.enable_elastic_disk + if self.enable_local_disk_encryption is not None: + body["enable_local_disk_encryption"] = self.enable_local_disk_encryption + if self.gcp_attributes: + body["gcp_attributes"] = self.gcp_attributes.as_dict() + if self.init_scripts: + body["init_scripts"] = [v.as_dict() for v in self.init_scripts] + if self.instance_pool_id is not None: + body["instance_pool_id"] = self.instance_pool_id + if self.is_single_node is not None: + body["is_single_node"] = self.is_single_node + if self.kind is not None: + body["kind"] = self.kind.value + if self.node_type_id is not None: + body["node_type_id"] = self.node_type_id + if self.num_workers is not None: + body["num_workers"] = self.num_workers + if self.policy_id is not None: + body["policy_id"] = self.policy_id + if self.runtime_engine is not None: + body["runtime_engine"] = self.runtime_engine.value + if self.single_user_name is not None: + body["single_user_name"] = self.single_user_name + if self.spark_conf: + body["spark_conf"] = self.spark_conf + if self.spark_env_vars: + body["spark_env_vars"] = self.spark_env_vars + if self.spark_version is not None: + body["spark_version"] = self.spark_version + if self.ssh_public_keys: + body["ssh_public_keys"] = [v for v in self.ssh_public_keys] + if self.use_ml_runtime is not None: + body["use_ml_runtime"] = self.use_ml_runtime + if self.workload_type: + body["workload_type"] = self.workload_type.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the JobsClusterSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.apply_policy_default_values is not None: + body["apply_policy_default_values"] = self.apply_policy_default_values + if self.autoscale: + body["autoscale"] = self.autoscale + if self.autotermination_minutes is not None: + body["autotermination_minutes"] = self.autotermination_minutes + if self.aws_attributes: + body["aws_attributes"] = self.aws_attributes + if self.azure_attributes: + body["azure_attributes"] = self.azure_attributes + if self.cluster_log_conf: + body["cluster_log_conf"] = self.cluster_log_conf + if self.cluster_name is not None: + body["cluster_name"] = self.cluster_name + if self.custom_tags: + body["custom_tags"] = self.custom_tags + if self.data_security_mode is not None: + body["data_security_mode"] = self.data_security_mode + if self.docker_image: + body["docker_image"] = self.docker_image + if self.driver_instance_pool_id is not None: + body["driver_instance_pool_id"] = self.driver_instance_pool_id + if self.driver_node_type_id is not None: + body["driver_node_type_id"] = self.driver_node_type_id + if self.enable_elastic_disk is not None: + body["enable_elastic_disk"] = self.enable_elastic_disk + if self.enable_local_disk_encryption is not None: + body["enable_local_disk_encryption"] = self.enable_local_disk_encryption + if self.gcp_attributes: + body["gcp_attributes"] = self.gcp_attributes + if self.init_scripts: + body["init_scripts"] = self.init_scripts + if self.instance_pool_id is not None: + body["instance_pool_id"] = self.instance_pool_id + if self.is_single_node is not None: + body["is_single_node"] = self.is_single_node + if self.kind is not None: + body["kind"] = self.kind + if self.node_type_id is not None: + body["node_type_id"] = self.node_type_id + if self.num_workers is not None: + body["num_workers"] = self.num_workers + if self.policy_id is not None: + body["policy_id"] = self.policy_id + if self.runtime_engine is not None: + body["runtime_engine"] = self.runtime_engine + if self.single_user_name is not None: + body["single_user_name"] = self.single_user_name + if self.spark_conf: + body["spark_conf"] = self.spark_conf + if self.spark_env_vars: + body["spark_env_vars"] = self.spark_env_vars + if self.spark_version is not None: + body["spark_version"] = self.spark_version + if self.ssh_public_keys: + body["ssh_public_keys"] = self.ssh_public_keys + if self.use_ml_runtime is not None: + body["use_ml_runtime"] = self.use_ml_runtime + if self.workload_type: + body["workload_type"] = self.workload_type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> JobsClusterSpec: + """Deserializes the JobsClusterSpec from a dictionary.""" + return cls( + apply_policy_default_values=d.get("apply_policy_default_values", None), + autoscale=_from_dict(d, "autoscale", AutoScale), + autotermination_minutes=d.get("autotermination_minutes", None), + aws_attributes=_from_dict(d, "aws_attributes", AwsAttributes), + azure_attributes=_from_dict(d, "azure_attributes", AzureAttributes), + cluster_log_conf=_from_dict(d, "cluster_log_conf", ClusterLogConf), + cluster_name=d.get("cluster_name", None), + custom_tags=d.get("custom_tags", None), + data_security_mode=_enum(d, "data_security_mode", DataSecurityMode), + docker_image=_from_dict(d, "docker_image", DockerImage), + driver_instance_pool_id=d.get("driver_instance_pool_id", None), + driver_node_type_id=d.get("driver_node_type_id", None), + enable_elastic_disk=d.get("enable_elastic_disk", None), + enable_local_disk_encryption=d.get("enable_local_disk_encryption", None), + gcp_attributes=_from_dict(d, "gcp_attributes", GcpAttributes), + init_scripts=_repeated_dict(d, "init_scripts", InitScriptInfo), + instance_pool_id=d.get("instance_pool_id", None), + is_single_node=d.get("is_single_node", None), + kind=_enum(d, "kind", Kind), + node_type_id=d.get("node_type_id", None), + num_workers=d.get("num_workers", None), + policy_id=d.get("policy_id", None), + runtime_engine=_enum(d, "runtime_engine", RuntimeEngine), + single_user_name=d.get("single_user_name", None), + spark_conf=d.get("spark_conf", None), + spark_env_vars=d.get("spark_env_vars", None), + spark_version=d.get("spark_version", None), + ssh_public_keys=d.get("ssh_public_keys", None), + use_ml_runtime=d.get("use_ml_runtime", None), + workload_type=_from_dict(d, "workload_type", WorkloadType), + ) + + class JobsHealthMetric(Enum): """Specifies the health metric that is being evaluated for a particular health rule. @@ -3342,6 +4565,111 @@ def from_dict(cls, d: Dict[str, Any]) -> JobsHealthRules: return cls(rules=_repeated_dict(d, "rules", JobsHealthRule)) +class Kind(Enum): + """The kind of compute described by this compute specification. + + Depending on `kind`, different validations and default values will be applied. + + Clusters with `kind = CLASSIC_PREVIEW` support the following fields, whereas clusters with no + specified `kind` do not. * [is_single_node](/api/workspace/clusters/create#is_single_node) * + [use_ml_runtime](/api/workspace/clusters/create#use_ml_runtime) * + [data_security_mode](/api/workspace/clusters/create#data_security_mode) set to + `DATA_SECURITY_MODE_AUTO`, `DATA_SECURITY_MODE_DEDICATED`, or `DATA_SECURITY_MODE_STANDARD` + + By using the [simple form], your clusters are automatically using `kind = CLASSIC_PREVIEW`. + + [simple form]: https://docs.databricks.com/compute/simple-form.html""" + + CLASSIC_PREVIEW = "CLASSIC_PREVIEW" + + +@dataclass +class Library: + cran: Optional[RCranLibrary] = None + """Specification of a CRAN library to be installed as part of the library""" + + egg: Optional[str] = None + """Deprecated. URI of the egg library to install. Installing Python egg files is deprecated and is + not supported in Databricks Runtime 14.0 and above.""" + + jar: Optional[str] = None + """URI of the JAR library to install. Supported URIs include Workspace paths, Unity Catalog Volumes + paths, and S3 URIs. For example: `{ "jar": "/Workspace/path/to/library.jar" }`, `{ "jar" : + "/Volumes/path/to/library.jar" }` or `{ "jar": "s3://my-bucket/library.jar" }`. If S3 is used, + please make sure the cluster has read access on the library. You may need to launch the cluster + with an IAM role to access the S3 URI.""" + + maven: Optional[MavenLibrary] = None + """Specification of a maven library to be installed. For example: `{ "coordinates": + "org.jsoup:jsoup:1.7.2" }`""" + + pypi: Optional[PythonPyPiLibrary] = None + """Specification of a PyPi library to be installed. For example: `{ "package": "simplejson" }`""" + + requirements: Optional[str] = None + """URI of the requirements.txt file to install. Only Workspace paths and Unity Catalog Volumes + paths are supported. For example: `{ "requirements": "/Workspace/path/to/requirements.txt" }` or + `{ "requirements" : "/Volumes/path/to/requirements.txt" }`""" + + whl: Optional[str] = None + """URI of the wheel library to install. Supported URIs include Workspace paths, Unity Catalog + Volumes paths, and S3 URIs. For example: `{ "whl": "/Workspace/path/to/library.whl" }`, `{ "whl" + : "/Volumes/path/to/library.whl" }` or `{ "whl": "s3://my-bucket/library.whl" }`. If S3 is used, + please make sure the cluster has read access on the library. You may need to launch the cluster + with an IAM role to access the S3 URI.""" + + def as_dict(self) -> dict: + """Serializes the Library into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.cran: + body["cran"] = self.cran.as_dict() + if self.egg is not None: + body["egg"] = self.egg + if self.jar is not None: + body["jar"] = self.jar + if self.maven: + body["maven"] = self.maven.as_dict() + if self.pypi: + body["pypi"] = self.pypi.as_dict() + if self.requirements is not None: + body["requirements"] = self.requirements + if self.whl is not None: + body["whl"] = self.whl + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Library into a shallow dictionary of its immediate attributes.""" + body = {} + if self.cran: + body["cran"] = self.cran + if self.egg is not None: + body["egg"] = self.egg + if self.jar is not None: + body["jar"] = self.jar + if self.maven: + body["maven"] = self.maven + if self.pypi: + body["pypi"] = self.pypi + if self.requirements is not None: + body["requirements"] = self.requirements + if self.whl is not None: + body["whl"] = self.whl + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Library: + """Deserializes the Library from a dictionary.""" + return cls( + cran=_from_dict(d, "cran", RCranLibrary), + egg=d.get("egg", None), + jar=d.get("jar", None), + maven=_from_dict(d, "maven", MavenLibrary), + pypi=_from_dict(d, "pypi", PythonPyPiLibrary), + requirements=d.get("requirements", None), + whl=d.get("whl", None), + ) + + @dataclass class ListJobComplianceForPolicyResponse: jobs: Optional[List[JobCompliance]] = None @@ -3494,6 +4822,109 @@ def from_dict(cls, d: Dict[str, Any]) -> ListRunsResponse: ) +@dataclass +class LocalFileInfo: + destination: str + """local file destination, e.g. `file:/my/local/file.sh`""" + + def as_dict(self) -> dict: + """Serializes the LocalFileInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the LocalFileInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> LocalFileInfo: + """Deserializes the LocalFileInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +@dataclass +class LogAnalyticsInfo: + log_analytics_primary_key: Optional[str] = None + + log_analytics_workspace_id: Optional[str] = None + + def as_dict(self) -> dict: + """Serializes the LogAnalyticsInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.log_analytics_primary_key is not None: + body["log_analytics_primary_key"] = self.log_analytics_primary_key + if self.log_analytics_workspace_id is not None: + body["log_analytics_workspace_id"] = self.log_analytics_workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the LogAnalyticsInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.log_analytics_primary_key is not None: + body["log_analytics_primary_key"] = self.log_analytics_primary_key + if self.log_analytics_workspace_id is not None: + body["log_analytics_workspace_id"] = self.log_analytics_workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> LogAnalyticsInfo: + """Deserializes the LogAnalyticsInfo from a dictionary.""" + return cls( + log_analytics_primary_key=d.get("log_analytics_primary_key", None), + log_analytics_workspace_id=d.get("log_analytics_workspace_id", None), + ) + + +@dataclass +class MavenLibrary: + coordinates: str + """Gradle-style maven coordinates. For example: "org.jsoup:jsoup:1.7.2".""" + + exclusions: Optional[List[str]] = None + """List of dependences to exclude. For example: `["slf4j:slf4j", "*:hadoop-client"]`. + + Maven dependency exclusions: + https://maven.apache.org/guides/introduction/introduction-to-optional-and-excludes-dependencies.html.""" + + repo: Optional[str] = None + """Maven repo to install the Maven package from. If omitted, both Maven Central Repository and + Spark Packages are searched.""" + + def as_dict(self) -> dict: + """Serializes the MavenLibrary into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.coordinates is not None: + body["coordinates"] = self.coordinates + if self.exclusions: + body["exclusions"] = [v for v in self.exclusions] + if self.repo is not None: + body["repo"] = self.repo + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MavenLibrary into a shallow dictionary of its immediate attributes.""" + body = {} + if self.coordinates is not None: + body["coordinates"] = self.coordinates + if self.exclusions: + body["exclusions"] = self.exclusions + if self.repo is not None: + body["repo"] = self.repo + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MavenLibrary: + """Deserializes the MavenLibrary from a dictionary.""" + return cls( + coordinates=d.get("coordinates", None), exclusions=d.get("exclusions", None), repo=d.get("repo", None) + ) + + @dataclass class NotebookOutput: result: Optional[str] = None @@ -3659,9 +5090,9 @@ class PerformanceTarget(Enum): on serverless compute should be. The performance mode on the job or pipeline should map to a performance setting that is passed to Cluster Manager (see cluster-common PerformanceTarget).""" - BALANCED = "BALANCED" COST_OPTIMIZED = "COST_OPTIMIZED" PERFORMANCE_OPTIMIZED = "PERFORMANCE_OPTIMIZED" + STANDARD = "STANDARD" @dataclass @@ -3703,61 +5134,263 @@ class PeriodicTriggerConfigurationTimeUnit(Enum): WEEKS = "WEEKS" -@dataclass -class PipelineParams: - full_refresh: Optional[bool] = None - """If true, triggers a full refresh on the delta live table.""" +@dataclass +class PipelineParams: + full_refresh: Optional[bool] = None + """If true, triggers a full refresh on the delta live table.""" + + def as_dict(self) -> dict: + """Serializes the PipelineParams into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.full_refresh is not None: + body["full_refresh"] = self.full_refresh + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PipelineParams into a shallow dictionary of its immediate attributes.""" + body = {} + if self.full_refresh is not None: + body["full_refresh"] = self.full_refresh + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PipelineParams: + """Deserializes the PipelineParams from a dictionary.""" + return cls(full_refresh=d.get("full_refresh", None)) + + +@dataclass +class PipelineTask: + pipeline_id: str + """The full name of the pipeline task to execute.""" + + full_refresh: Optional[bool] = None + """If true, triggers a full refresh on the delta live table.""" + + def as_dict(self) -> dict: + """Serializes the PipelineTask into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.full_refresh is not None: + body["full_refresh"] = self.full_refresh + if self.pipeline_id is not None: + body["pipeline_id"] = self.pipeline_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PipelineTask into a shallow dictionary of its immediate attributes.""" + body = {} + if self.full_refresh is not None: + body["full_refresh"] = self.full_refresh + if self.pipeline_id is not None: + body["pipeline_id"] = self.pipeline_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PipelineTask: + """Deserializes the PipelineTask from a dictionary.""" + return cls(full_refresh=d.get("full_refresh", None), pipeline_id=d.get("pipeline_id", None)) + + +@dataclass +class PowerBiModel: + authentication_method: Optional[AuthenticationMethod] = None + """How the published Power BI model authenticates to Databricks""" + + model_name: Optional[str] = None + """The name of the Power BI model""" + + overwrite_existing: Optional[bool] = None + """Whether to overwrite existing Power BI models""" + + storage_mode: Optional[StorageMode] = None + """The default storage mode of the Power BI model""" + + workspace_name: Optional[str] = None + """The name of the Power BI workspace of the model""" + + def as_dict(self) -> dict: + """Serializes the PowerBiModel into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.authentication_method is not None: + body["authentication_method"] = self.authentication_method.value + if self.model_name is not None: + body["model_name"] = self.model_name + if self.overwrite_existing is not None: + body["overwrite_existing"] = self.overwrite_existing + if self.storage_mode is not None: + body["storage_mode"] = self.storage_mode.value + if self.workspace_name is not None: + body["workspace_name"] = self.workspace_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PowerBiModel into a shallow dictionary of its immediate attributes.""" + body = {} + if self.authentication_method is not None: + body["authentication_method"] = self.authentication_method + if self.model_name is not None: + body["model_name"] = self.model_name + if self.overwrite_existing is not None: + body["overwrite_existing"] = self.overwrite_existing + if self.storage_mode is not None: + body["storage_mode"] = self.storage_mode + if self.workspace_name is not None: + body["workspace_name"] = self.workspace_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PowerBiModel: + """Deserializes the PowerBiModel from a dictionary.""" + return cls( + authentication_method=_enum(d, "authentication_method", AuthenticationMethod), + model_name=d.get("model_name", None), + overwrite_existing=d.get("overwrite_existing", None), + storage_mode=_enum(d, "storage_mode", StorageMode), + workspace_name=d.get("workspace_name", None), + ) + + +@dataclass +class PowerBiTable: + catalog: Optional[str] = None + """The catalog name in Databricks""" + + name: Optional[str] = None + """The table name in Databricks""" + + schema: Optional[str] = None + """The schema name in Databricks""" + + storage_mode: Optional[StorageMode] = None + """The Power BI storage mode of the table""" + + def as_dict(self) -> dict: + """Serializes the PowerBiTable into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + if self.storage_mode is not None: + body["storage_mode"] = self.storage_mode.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the PowerBiTable into a shallow dictionary of its immediate attributes.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + if self.storage_mode is not None: + body["storage_mode"] = self.storage_mode + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> PowerBiTable: + """Deserializes the PowerBiTable from a dictionary.""" + return cls( + catalog=d.get("catalog", None), + name=d.get("name", None), + schema=d.get("schema", None), + storage_mode=_enum(d, "storage_mode", StorageMode), + ) + + +@dataclass +class PowerBiTask: + connection_resource_name: Optional[str] = None + """The resource name of the UC connection to authenticate from Databricks to Power BI""" + + power_bi_model: Optional[PowerBiModel] = None + """The semantic model to update""" + + refresh_after_update: Optional[bool] = None + """Whether the model should be refreshed after the update""" + + tables: Optional[List[PowerBiTable]] = None + """The tables to be exported to Power BI""" + + warehouse_id: Optional[str] = None + """The SQL warehouse ID to use as the Power BI data source""" def as_dict(self) -> dict: - """Serializes the PipelineParams into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.full_refresh is not None: - body["full_refresh"] = self.full_refresh + """Serializes the PowerBiTask into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.connection_resource_name is not None: + body["connection_resource_name"] = self.connection_resource_name + if self.power_bi_model: + body["power_bi_model"] = self.power_bi_model.as_dict() + if self.refresh_after_update is not None: + body["refresh_after_update"] = self.refresh_after_update + if self.tables: + body["tables"] = [v.as_dict() for v in self.tables] + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id return body def as_shallow_dict(self) -> dict: - """Serializes the PipelineParams into a shallow dictionary of its immediate attributes.""" - body = {} - if self.full_refresh is not None: - body["full_refresh"] = self.full_refresh + """Serializes the PowerBiTask into a shallow dictionary of its immediate attributes.""" + body = {} + if self.connection_resource_name is not None: + body["connection_resource_name"] = self.connection_resource_name + if self.power_bi_model: + body["power_bi_model"] = self.power_bi_model + if self.refresh_after_update is not None: + body["refresh_after_update"] = self.refresh_after_update + if self.tables: + body["tables"] = self.tables + if self.warehouse_id is not None: + body["warehouse_id"] = self.warehouse_id return body @classmethod - def from_dict(cls, d: Dict[str, Any]) -> PipelineParams: - """Deserializes the PipelineParams from a dictionary.""" - return cls(full_refresh=d.get("full_refresh", None)) + def from_dict(cls, d: Dict[str, Any]) -> PowerBiTask: + """Deserializes the PowerBiTask from a dictionary.""" + return cls( + connection_resource_name=d.get("connection_resource_name", None), + power_bi_model=_from_dict(d, "power_bi_model", PowerBiModel), + refresh_after_update=d.get("refresh_after_update", None), + tables=_repeated_dict(d, "tables", PowerBiTable), + warehouse_id=d.get("warehouse_id", None), + ) @dataclass -class PipelineTask: - pipeline_id: str - """The full name of the pipeline task to execute.""" +class PythonPyPiLibrary: + package: str + """The name of the pypi package to install. An optional exact version specification is also + supported. Examples: "simplejson" and "simplejson==3.8.0".""" - full_refresh: Optional[bool] = None - """If true, triggers a full refresh on the delta live table.""" + repo: Optional[str] = None + """The repository where the package can be found. If not specified, the default pip index is used.""" def as_dict(self) -> dict: - """Serializes the PipelineTask into a dictionary suitable for use as a JSON request body.""" + """Serializes the PythonPyPiLibrary into a dictionary suitable for use as a JSON request body.""" body = {} - if self.full_refresh is not None: - body["full_refresh"] = self.full_refresh - if self.pipeline_id is not None: - body["pipeline_id"] = self.pipeline_id + if self.package is not None: + body["package"] = self.package + if self.repo is not None: + body["repo"] = self.repo return body def as_shallow_dict(self) -> dict: - """Serializes the PipelineTask into a shallow dictionary of its immediate attributes.""" + """Serializes the PythonPyPiLibrary into a shallow dictionary of its immediate attributes.""" body = {} - if self.full_refresh is not None: - body["full_refresh"] = self.full_refresh - if self.pipeline_id is not None: - body["pipeline_id"] = self.pipeline_id + if self.package is not None: + body["package"] = self.package + if self.repo is not None: + body["repo"] = self.repo return body @classmethod - def from_dict(cls, d: Dict[str, Any]) -> PipelineTask: - """Deserializes the PipelineTask from a dictionary.""" - return cls(full_refresh=d.get("full_refresh", None), pipeline_id=d.get("pipeline_id", None)) + def from_dict(cls, d: Dict[str, Any]) -> PythonPyPiLibrary: + """Deserializes the PythonPyPiLibrary from a dictionary.""" + return cls(package=d.get("package", None), repo=d.get("repo", None)) @dataclass @@ -3888,6 +5521,38 @@ def from_dict(cls, d: Dict[str, Any]) -> QueueSettings: return cls(enabled=d.get("enabled", None)) +@dataclass +class RCranLibrary: + package: str + """The name of the CRAN package to install.""" + + repo: Optional[str] = None + """The repository where the package can be found. If not specified, the default CRAN repo is used.""" + + def as_dict(self) -> dict: + """Serializes the RCranLibrary into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.package is not None: + body["package"] = self.package + if self.repo is not None: + body["repo"] = self.repo + return body + + def as_shallow_dict(self) -> dict: + """Serializes the RCranLibrary into a shallow dictionary of its immediate attributes.""" + body = {} + if self.package is not None: + body["package"] = self.package + if self.repo is not None: + body["repo"] = self.repo + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> RCranLibrary: + """Deserializes the RCranLibrary from a dictionary.""" + return cls(package=d.get("package", None), repo=d.get("repo", None)) + + @dataclass class RepairHistoryItem: end_time: Optional[int] = None @@ -4542,9 +6207,9 @@ class Run: """Description of the run""" effective_performance_target: Optional[PerformanceTarget] = None - """effective_performance_target is the actual performance target used by the run during execution. - effective_performance_target can differ from the client-set performance_target depending on if - the job was eligible to be cost-optimized.""" + """The actual performance target used by the serverless run during execution. This can differ from + the client-set performance target on the request depending on whether the performance mode is + supported by the job type.""" end_time: Optional[int] = None """The time at which this run ended in epoch milliseconds (milliseconds since 1/1/1970 UTC). This @@ -4568,8 +6233,8 @@ class Run: are used, `git_source` must be defined on the job.""" has_more: Optional[bool] = None - """Indicates if the run has more sub-resources (`tasks`, `job_clusters`) that are not shown. They - can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 + """Indicates if the run has more array properties (`tasks`, `job_clusters`) that are not shown. + They can be accessed via :method:jobs/getrun endpoint. It is only relevant for API 2.2 :method:jobs/listruns requests with `expand_tasks=true`.""" iterations: Optional[List[RunTask]] = None @@ -4593,7 +6258,7 @@ class Run: that the task run belongs to.""" next_page_token: Optional[str] = None - """A token that can be used to list the next page of sub-resources.""" + """A token that can be used to list the next page of array properties.""" number_in_job: Optional[int] = None """A unique identifier for this job run. This is set to the same value as `run_id`.""" @@ -5266,9 +6931,9 @@ class RunNow: job will be run.""" performance_target: Optional[PerformanceTarget] = None - """PerformanceTarget defines how performant or cost efficient the execution of run on serverless - compute should be. For RunNow, this performance target will override the target defined on the - job-level.""" + """The performance mode on a serverless job. The performance target determines the level of compute + performance or cost-efficiency for the run. This field overrides the performance target defined + on the job-level.""" pipeline_params: Optional[PipelineParams] = None """Controls whether the pipeline should perform a full refresh""" @@ -5443,6 +7108,9 @@ class RunOutput: clean_rooms_notebook_output: Optional[CleanRoomsNotebookTaskCleanRoomsNotebookTaskOutput] = None """The output of a clean rooms notebook task, if available""" + dashboard_output: Optional[DashboardTaskOutput] = None + """The output of a dashboard task, if available""" + dbt_output: Optional[DbtOutput] = None """The output of a dbt task, if available.""" @@ -5489,6 +7157,8 @@ def as_dict(self) -> dict: body = {} if self.clean_rooms_notebook_output: body["clean_rooms_notebook_output"] = self.clean_rooms_notebook_output.as_dict() + if self.dashboard_output: + body["dashboard_output"] = self.dashboard_output.as_dict() if self.dbt_output: body["dbt_output"] = self.dbt_output.as_dict() if self.error is not None: @@ -5516,6 +7186,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.clean_rooms_notebook_output: body["clean_rooms_notebook_output"] = self.clean_rooms_notebook_output + if self.dashboard_output: + body["dashboard_output"] = self.dashboard_output if self.dbt_output: body["dbt_output"] = self.dbt_output if self.error is not None: @@ -5545,6 +7217,7 @@ def from_dict(cls, d: Dict[str, Any]) -> RunOutput: clean_rooms_notebook_output=_from_dict( d, "clean_rooms_notebook_output", CleanRoomsNotebookTaskCleanRoomsNotebookTaskOutput ), + dashboard_output=_from_dict(d, "dashboard_output", DashboardTaskOutput), dbt_output=_from_dict(d, "dbt_output", DbtOutput), error=d.get("error", None), error_trace=d.get("error_trace", None), @@ -5860,6 +7533,9 @@ class RunTask: `condition_task` field is present. The condition task does not require a cluster to execute and does not support retries or notifications.""" + dashboard_task: Optional[DashboardTask] = None + """The task runs a DashboardTask when the `dashboard_task` field is present.""" + dbt_task: Optional[DbtTask] = None """The task runs one or more dbt commands when the `dbt_task` field is present. The dbt task requires both Databricks SQL and the ability to use a serverless or a pro SQL warehouse.""" @@ -5873,13 +7549,12 @@ class RunTask: """An optional description for this task.""" disabled: Optional[bool] = None - """Denotes whether or not the task was disabled by the user. Disabled tasks do not execute and are - immediately skipped as soon as they are unblocked.""" + """Deprecated, field was never used in production.""" effective_performance_target: Optional[PerformanceTarget] = None - """effective_performance_target is the actual performance target used by the run during execution. - effective_performance_target can differ from the client-set performance_target depending on if - the job was eligible to be cost-optimized.""" + """The actual performance target used by the serverless run during execution. This can differ from + the client-set performance target on the request depending on whether the performance mode is + supported by the job type.""" email_notifications: Optional[JobEmailNotifications] = None """An optional set of email addresses notified when the task run begins or completes. The default @@ -5910,7 +7585,6 @@ class RunTask: present.""" gen_ai_compute_task: Optional[GenAiComputeTask] = None - """Next field: 9""" git_source: Optional[GitSource] = None """An optional specification for a remote Git repository containing the source code used by tasks. @@ -5924,11 +7598,11 @@ class RunTask: """If job_cluster_key, this task is executed reusing the cluster specified in `job.settings.job_clusters`.""" - libraries: Optional[List[compute.Library]] = None + libraries: Optional[List[Library]] = None """An optional list of libraries to be installed on the cluster. The default value is an empty list.""" - new_cluster: Optional[compute.ClusterSpec] = None + new_cluster: Optional[ClusterSpec] = None """If new_cluster, a description of a new cluster that is created for each run.""" notebook_task: Optional[NotebookTask] = None @@ -5942,6 +7616,9 @@ class RunTask: """The task triggers a pipeline update when the `pipeline_task` field is present. Only pipelines configured to use triggered more are supported.""" + power_bi_task: Optional[PowerBiTask] = None + """The task triggers a Power BI semantic model update when the `power_bi_task` field is present.""" + python_wheel_task: Optional[PythonWheelTask] = None """The task runs a Python wheel when the `python_wheel_task` field is present.""" @@ -6033,6 +7710,8 @@ def as_dict(self) -> dict: body["cluster_instance"] = self.cluster_instance.as_dict() if self.condition_task: body["condition_task"] = self.condition_task.as_dict() + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task.as_dict() if self.dbt_task: body["dbt_task"] = self.dbt_task.as_dict() if self.depends_on: @@ -6071,6 +7750,8 @@ def as_dict(self) -> dict: body["notification_settings"] = self.notification_settings.as_dict() if self.pipeline_task: body["pipeline_task"] = self.pipeline_task.as_dict() + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task.as_dict() if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task.as_dict() if self.queue_duration is not None: @@ -6124,6 +7805,8 @@ def as_shallow_dict(self) -> dict: body["cluster_instance"] = self.cluster_instance if self.condition_task: body["condition_task"] = self.condition_task + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task if self.dbt_task: body["dbt_task"] = self.dbt_task if self.depends_on: @@ -6162,6 +7845,8 @@ def as_shallow_dict(self) -> dict: body["notification_settings"] = self.notification_settings if self.pipeline_task: body["pipeline_task"] = self.pipeline_task + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task if self.queue_duration is not None: @@ -6211,6 +7896,7 @@ def from_dict(cls, d: Dict[str, Any]) -> RunTask: cleanup_duration=d.get("cleanup_duration", None), cluster_instance=_from_dict(d, "cluster_instance", ClusterInstance), condition_task=_from_dict(d, "condition_task", RunConditionTask), + dashboard_task=_from_dict(d, "dashboard_task", DashboardTask), dbt_task=_from_dict(d, "dbt_task", DbtTask), depends_on=_repeated_dict(d, "depends_on", TaskDependency), description=d.get("description", None), @@ -6225,11 +7911,12 @@ def from_dict(cls, d: Dict[str, Any]) -> RunTask: gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), git_source=_from_dict(d, "git_source", GitSource), job_cluster_key=d.get("job_cluster_key", None), - libraries=_repeated_dict(d, "libraries", compute.Library), - new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + libraries=_repeated_dict(d, "libraries", Library), + new_cluster=_from_dict(d, "new_cluster", ClusterSpec), notebook_task=_from_dict(d, "notebook_task", NotebookTask), notification_settings=_from_dict(d, "notification_settings", TaskNotificationSettings), pipeline_task=_from_dict(d, "pipeline_task", PipelineTask), + power_bi_task=_from_dict(d, "power_bi_task", PowerBiTask), python_wheel_task=_from_dict(d, "python_wheel_task", PythonWheelTask), queue_duration=d.get("queue_duration", None), resolved_values=_from_dict(d, "resolved_values", ResolvedValues), @@ -6264,6 +7951,102 @@ class RunType(Enum): WORKFLOW_RUN = "WORKFLOW_RUN" +class RuntimeEngine(Enum): + + NULL = "NULL" + PHOTON = "PHOTON" + STANDARD = "STANDARD" + + +@dataclass +class S3StorageInfo: + """A storage location in Amazon S3""" + + destination: str + """S3 destination, e.g. `s3://my-bucket/some-prefix` Note that logs will be delivered using cluster + iam role, please make sure you set cluster iam role and the role has write access to the + destination. Please also note that you cannot use AWS keys to deliver logs.""" + + canned_acl: Optional[str] = None + """(Optional) Set canned access control list for the logs, e.g. `bucket-owner-full-control`. If + `canned_cal` is set, please make sure the cluster iam role has `s3:PutObjectAcl` permission on + the destination bucket and prefix. The full list of possible canned acl can be found at + http://docs.aws.amazon.com/AmazonS3/latest/dev/acl-overview.html#canned-acl. Please also note + that by default only the object owner gets full controls. If you are using cross account role + for writing data, you may want to set `bucket-owner-full-control` to make bucket owner able to + read the logs.""" + + enable_encryption: Optional[bool] = None + """(Optional) Flag to enable server side encryption, `false` by default.""" + + encryption_type: Optional[str] = None + """(Optional) The encryption type, it could be `sse-s3` or `sse-kms`. It will be used only when + encryption is enabled and the default type is `sse-s3`.""" + + endpoint: Optional[str] = None + """S3 endpoint, e.g. `https://s3-us-west-2.amazonaws.com`. Either region or endpoint needs to be + set. If both are set, endpoint will be used.""" + + kms_key: Optional[str] = None + """(Optional) Kms key which will be used if encryption is enabled and encryption type is set to + `sse-kms`.""" + + region: Optional[str] = None + """S3 region, e.g. `us-west-2`. Either region or endpoint needs to be set. If both are set, + endpoint will be used.""" + + def as_dict(self) -> dict: + """Serializes the S3StorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.canned_acl is not None: + body["canned_acl"] = self.canned_acl + if self.destination is not None: + body["destination"] = self.destination + if self.enable_encryption is not None: + body["enable_encryption"] = self.enable_encryption + if self.encryption_type is not None: + body["encryption_type"] = self.encryption_type + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.kms_key is not None: + body["kms_key"] = self.kms_key + if self.region is not None: + body["region"] = self.region + return body + + def as_shallow_dict(self) -> dict: + """Serializes the S3StorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.canned_acl is not None: + body["canned_acl"] = self.canned_acl + if self.destination is not None: + body["destination"] = self.destination + if self.enable_encryption is not None: + body["enable_encryption"] = self.enable_encryption + if self.encryption_type is not None: + body["encryption_type"] = self.encryption_type + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.kms_key is not None: + body["kms_key"] = self.kms_key + if self.region is not None: + body["region"] = self.region + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> S3StorageInfo: + """Deserializes the S3StorageInfo from a dictionary.""" + return cls( + canned_acl=d.get("canned_acl", None), + destination=d.get("destination", None), + enable_encryption=d.get("enable_encryption", None), + encryption_type=d.get("encryption_type", None), + endpoint=d.get("endpoint", None), + kms_key=d.get("kms_key", None), + region=d.get("region", None), + ) + + class Source(Enum): """Optional location type of the SQL file. When set to `WORKSPACE`, the SQL file will be retrieved\ from the local Databricks workspace. When set to `GIT`, the SQL file will be retrieved from a @@ -7028,6 +8811,13 @@ def from_dict(cls, d: Dict[str, Any]) -> SqlTaskSubscription: return cls(destination_id=d.get("destination_id", None), user_name=d.get("user_name", None)) +class StorageMode(Enum): + + DIRECT_QUERY = "DIRECT_QUERY" + DUAL = "DUAL" + IMPORT = "IMPORT" + + @dataclass class SubmitRun: access_control_list: Optional[List[JobAccessControlRequest]] = None @@ -7223,6 +9013,9 @@ class SubmitTask: `condition_task` field is present. The condition task does not require a cluster to execute and does not support retries or notifications.""" + dashboard_task: Optional[DashboardTask] = None + """The task runs a DashboardTask when the `dashboard_task` field is present.""" + dbt_task: Optional[DbtTask] = None """The task runs one or more dbt commands when the `dbt_task` field is present. The dbt task requires both Databricks SQL and the ability to use a serverless or a pro SQL warehouse.""" @@ -7253,16 +9046,15 @@ class SubmitTask: present.""" gen_ai_compute_task: Optional[GenAiComputeTask] = None - """Next field: 9""" health: Optional[JobsHealthRules] = None """An optional set of health rules that can be defined for this job.""" - libraries: Optional[List[compute.Library]] = None + libraries: Optional[List[Library]] = None """An optional list of libraries to be installed on the cluster. The default value is an empty list.""" - new_cluster: Optional[compute.ClusterSpec] = None + new_cluster: Optional[JobsClusterSpec] = None """If new_cluster, a description of a new cluster that is created for each run.""" notebook_task: Optional[NotebookTask] = None @@ -7276,6 +9068,9 @@ class SubmitTask: """The task triggers a pipeline update when the `pipeline_task` field is present. Only pipelines configured to use triggered more are supported.""" + power_bi_task: Optional[PowerBiTask] = None + """The task triggers a Power BI semantic model update when the `power_bi_task` field is present.""" + python_wheel_task: Optional[PythonWheelTask] = None """The task runs a Python wheel when the `python_wheel_task` field is present.""" @@ -7329,6 +9124,8 @@ def as_dict(self) -> dict: body["clean_rooms_notebook_task"] = self.clean_rooms_notebook_task.as_dict() if self.condition_task: body["condition_task"] = self.condition_task.as_dict() + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task.as_dict() if self.dbt_task: body["dbt_task"] = self.dbt_task.as_dict() if self.depends_on: @@ -7357,6 +9154,8 @@ def as_dict(self) -> dict: body["notification_settings"] = self.notification_settings.as_dict() if self.pipeline_task: body["pipeline_task"] = self.pipeline_task.as_dict() + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task.as_dict() if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task.as_dict() if self.run_if is not None: @@ -7386,6 +9185,8 @@ def as_shallow_dict(self) -> dict: body["clean_rooms_notebook_task"] = self.clean_rooms_notebook_task if self.condition_task: body["condition_task"] = self.condition_task + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task if self.dbt_task: body["dbt_task"] = self.dbt_task if self.depends_on: @@ -7414,6 +9215,8 @@ def as_shallow_dict(self) -> dict: body["notification_settings"] = self.notification_settings if self.pipeline_task: body["pipeline_task"] = self.pipeline_task + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task if self.run_if is not None: @@ -7442,6 +9245,7 @@ def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: return cls( clean_rooms_notebook_task=_from_dict(d, "clean_rooms_notebook_task", CleanRoomsNotebookTask), condition_task=_from_dict(d, "condition_task", ConditionTask), + dashboard_task=_from_dict(d, "dashboard_task", DashboardTask), dbt_task=_from_dict(d, "dbt_task", DbtTask), depends_on=_repeated_dict(d, "depends_on", TaskDependency), description=d.get("description", None), @@ -7451,11 +9255,12 @@ def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: for_each_task=_from_dict(d, "for_each_task", ForEachTask), gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), health=_from_dict(d, "health", JobsHealthRules), - libraries=_repeated_dict(d, "libraries", compute.Library), - new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + libraries=_repeated_dict(d, "libraries", Library), + new_cluster=_from_dict(d, "new_cluster", JobsClusterSpec), notebook_task=_from_dict(d, "notebook_task", NotebookTask), notification_settings=_from_dict(d, "notification_settings", TaskNotificationSettings), pipeline_task=_from_dict(d, "pipeline_task", PipelineTask), + power_bi_task=_from_dict(d, "power_bi_task", PowerBiTask), python_wheel_task=_from_dict(d, "python_wheel_task", PythonWheelTask), run_if=_enum(d, "run_if", RunIf), run_job_task=_from_dict(d, "run_job_task", RunJobTask), @@ -7469,6 +9274,78 @@ def from_dict(cls, d: Dict[str, Any]) -> SubmitTask: ) +@dataclass +class Subscription: + custom_subject: Optional[str] = None + """Optional: Allows users to specify a custom subject line on the email sent to subscribers.""" + + paused: Optional[bool] = None + """When true, the subscription will not send emails.""" + + subscribers: Optional[List[SubscriptionSubscriber]] = None + + def as_dict(self) -> dict: + """Serializes the Subscription into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.custom_subject is not None: + body["custom_subject"] = self.custom_subject + if self.paused is not None: + body["paused"] = self.paused + if self.subscribers: + body["subscribers"] = [v.as_dict() for v in self.subscribers] + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Subscription into a shallow dictionary of its immediate attributes.""" + body = {} + if self.custom_subject is not None: + body["custom_subject"] = self.custom_subject + if self.paused is not None: + body["paused"] = self.paused + if self.subscribers: + body["subscribers"] = self.subscribers + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Subscription: + """Deserializes the Subscription from a dictionary.""" + return cls( + custom_subject=d.get("custom_subject", None), + paused=d.get("paused", None), + subscribers=_repeated_dict(d, "subscribers", SubscriptionSubscriber), + ) + + +@dataclass +class SubscriptionSubscriber: + destination_id: Optional[str] = None + + user_name: Optional[str] = None + + def as_dict(self) -> dict: + """Serializes the SubscriptionSubscriber into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination_id is not None: + body["destination_id"] = self.destination_id + if self.user_name is not None: + body["user_name"] = self.user_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the SubscriptionSubscriber into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination_id is not None: + body["destination_id"] = self.destination_id + if self.user_name is not None: + body["user_name"] = self.user_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> SubscriptionSubscriber: + """Deserializes the SubscriptionSubscriber from a dictionary.""" + return cls(destination_id=d.get("destination_id", None), user_name=d.get("user_name", None)) + + @dataclass class TableUpdateTriggerConfiguration: condition: Optional[Condition] = None @@ -7541,6 +9418,9 @@ class Task: `condition_task` field is present. The condition task does not require a cluster to execute and does not support retries or notifications.""" + dashboard_task: Optional[DashboardTask] = None + """The task runs a DashboardTask when the `dashboard_task` field is present.""" + dbt_task: Optional[DbtTask] = None """The task runs one or more dbt commands when the `dbt_task` field is present. The dbt task requires both Databricks SQL and the ability to use a serverless or a pro SQL warehouse.""" @@ -7575,7 +9455,6 @@ class Task: present.""" gen_ai_compute_task: Optional[GenAiComputeTask] = None - """Next field: 9""" health: Optional[JobsHealthRules] = None """An optional set of health rules that can be defined for this job.""" @@ -7584,7 +9463,7 @@ class Task: """If job_cluster_key, this task is executed reusing the cluster specified in `job.settings.job_clusters`.""" - libraries: Optional[List[compute.Library]] = None + libraries: Optional[List[Library]] = None """An optional list of libraries to be installed on the cluster. The default value is an empty list.""" @@ -7598,7 +9477,7 @@ class Task: """An optional minimal interval in milliseconds between the start of the failed run and the subsequent retry run. The default behavior is that unsuccessful runs are immediately retried.""" - new_cluster: Optional[compute.ClusterSpec] = None + new_cluster: Optional[JobsClusterSpec] = None """If new_cluster, a description of a new cluster that is created for each run.""" notebook_task: Optional[NotebookTask] = None @@ -7612,6 +9491,9 @@ class Task: """The task triggers a pipeline update when the `pipeline_task` field is present. Only pipelines configured to use triggered more are supported.""" + power_bi_task: Optional[PowerBiTask] = None + """The task triggers a Power BI semantic model update when the `power_bi_task` field is present.""" + python_wheel_task: Optional[PythonWheelTask] = None """The task runs a Python wheel when the `python_wheel_task` field is present.""" @@ -7672,6 +9554,8 @@ def as_dict(self) -> dict: body["clean_rooms_notebook_task"] = self.clean_rooms_notebook_task.as_dict() if self.condition_task: body["condition_task"] = self.condition_task.as_dict() + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task.as_dict() if self.dbt_task: body["dbt_task"] = self.dbt_task.as_dict() if self.depends_on: @@ -7708,6 +9592,8 @@ def as_dict(self) -> dict: body["notification_settings"] = self.notification_settings.as_dict() if self.pipeline_task: body["pipeline_task"] = self.pipeline_task.as_dict() + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task.as_dict() if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task.as_dict() if self.retry_on_timeout is not None: @@ -7739,6 +9625,8 @@ def as_shallow_dict(self) -> dict: body["clean_rooms_notebook_task"] = self.clean_rooms_notebook_task if self.condition_task: body["condition_task"] = self.condition_task + if self.dashboard_task: + body["dashboard_task"] = self.dashboard_task if self.dbt_task: body["dbt_task"] = self.dbt_task if self.depends_on: @@ -7775,6 +9663,8 @@ def as_shallow_dict(self) -> dict: body["notification_settings"] = self.notification_settings if self.pipeline_task: body["pipeline_task"] = self.pipeline_task + if self.power_bi_task: + body["power_bi_task"] = self.power_bi_task if self.python_wheel_task: body["python_wheel_task"] = self.python_wheel_task if self.retry_on_timeout is not None: @@ -7805,6 +9695,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Task: return cls( clean_rooms_notebook_task=_from_dict(d, "clean_rooms_notebook_task", CleanRoomsNotebookTask), condition_task=_from_dict(d, "condition_task", ConditionTask), + dashboard_task=_from_dict(d, "dashboard_task", DashboardTask), dbt_task=_from_dict(d, "dbt_task", DbtTask), depends_on=_repeated_dict(d, "depends_on", TaskDependency), description=d.get("description", None), @@ -7816,13 +9707,14 @@ def from_dict(cls, d: Dict[str, Any]) -> Task: gen_ai_compute_task=_from_dict(d, "gen_ai_compute_task", GenAiComputeTask), health=_from_dict(d, "health", JobsHealthRules), job_cluster_key=d.get("job_cluster_key", None), - libraries=_repeated_dict(d, "libraries", compute.Library), + libraries=_repeated_dict(d, "libraries", Library), max_retries=d.get("max_retries", None), min_retry_interval_millis=d.get("min_retry_interval_millis", None), - new_cluster=_from_dict(d, "new_cluster", compute.ClusterSpec), + new_cluster=_from_dict(d, "new_cluster", JobsClusterSpec), notebook_task=_from_dict(d, "notebook_task", NotebookTask), notification_settings=_from_dict(d, "notification_settings", TaskNotificationSettings), pipeline_task=_from_dict(d, "pipeline_task", PipelineTask), + power_bi_task=_from_dict(d, "power_bi_task", PowerBiTask), python_wheel_task=_from_dict(d, "python_wheel_task", PythonWheelTask), retry_on_timeout=d.get("retry_on_timeout", None), run_if=_enum(d, "run_if", RunIf), @@ -8392,6 +10284,34 @@ class ViewsToExport(Enum): DASHBOARDS = "DASHBOARDS" +@dataclass +class VolumesStorageInfo: + """A storage location back by UC Volumes.""" + + destination: str + """UC Volumes destination, e.g. `/Volumes/catalog/schema/vol1/init-scripts/setup-datadog.sh` or + `dbfs:/Volumes/catalog/schema/vol1/init-scripts/setup-datadog.sh`""" + + def as_dict(self) -> dict: + """Serializes the VolumesStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the VolumesStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> VolumesStorageInfo: + """Deserializes the VolumesStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + @dataclass class Webhook: id: str @@ -8487,6 +10407,84 @@ def from_dict(cls, d: Dict[str, Any]) -> WebhookNotifications: ) +@dataclass +class WidgetErrorDetail: + message: Optional[str] = None + + def as_dict(self) -> dict: + """Serializes the WidgetErrorDetail into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.message is not None: + body["message"] = self.message + return body + + def as_shallow_dict(self) -> dict: + """Serializes the WidgetErrorDetail into a shallow dictionary of its immediate attributes.""" + body = {} + if self.message is not None: + body["message"] = self.message + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> WidgetErrorDetail: + """Deserializes the WidgetErrorDetail from a dictionary.""" + return cls(message=d.get("message", None)) + + +@dataclass +class WorkloadType: + """Cluster Attributes showing for clusters workload types.""" + + clients: ClientsTypes + """defined what type of clients can use the cluster. E.g. Notebooks, Jobs""" + + def as_dict(self) -> dict: + """Serializes the WorkloadType into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.clients: + body["clients"] = self.clients.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the WorkloadType into a shallow dictionary of its immediate attributes.""" + body = {} + if self.clients: + body["clients"] = self.clients + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> WorkloadType: + """Deserializes the WorkloadType from a dictionary.""" + return cls(clients=_from_dict(d, "clients", ClientsTypes)) + + +@dataclass +class WorkspaceStorageInfo: + """A storage location in Workspace Filesystem (WSFS)""" + + destination: str + """wsfs destination, e.g. `workspace:/cluster-init-scripts/setup-datadog.sh`""" + + def as_dict(self) -> dict: + """Serializes the WorkspaceStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the WorkspaceStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> WorkspaceStorageInfo: + """Deserializes the WorkspaceStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + class JobsAPI: """The Jobs API allows you to create, edit, and delete jobs. @@ -8507,40 +10505,6 @@ class JobsAPI: def __init__(self, api_client): self._api = api_client - def wait_get_run_job_terminated_or_skipped( - self, run_id: int, timeout=timedelta(minutes=20), callback: Optional[Callable[[Run], None]] = None - ) -> Run: - deadline = time.time() + timeout.total_seconds() - target_states = ( - RunLifeCycleState.TERMINATED, - RunLifeCycleState.SKIPPED, - ) - failure_states = (RunLifeCycleState.INTERNAL_ERROR,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get_run(run_id=run_id) - status = poll.state.life_cycle_state - status_message = f"current status: {status}" - if poll.state: - status_message = poll.state.state_message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach TERMINATED or SKIPPED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"run_id={run_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def cancel_all_runs(self, *, all_queued_runs: Optional[bool] = None, job_id: Optional[int] = None): """Cancel all runs of a job. @@ -8566,7 +10530,7 @@ def cancel_all_runs(self, *, all_queued_runs: Optional[bool] = None, job_id: Opt self._api.do("POST", "/api/2.2/jobs/runs/cancel-all", body=body, headers=headers) - def cancel_run(self, run_id: int) -> Wait[Run]: + def cancel_run(self, run_id: int): """Cancel a run. Cancels a job run or a task run. The run is canceled asynchronously, so it may still be running when @@ -8577,7 +10541,7 @@ def cancel_run(self, run_id: int) -> Wait[Run]: :returns: Long-running operation waiter for :class:`Run`. - See :method:wait_get_run_job_terminated_or_skipped for more details. + See :method:WaitGetRunJobTerminatedOrSkipped for more details. """ body = {} if run_id is not None: @@ -8586,15 +10550,7 @@ def cancel_run(self, run_id: int) -> Wait[Run]: "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.2/jobs/runs/cancel", body=body, headers=headers) - return Wait( - self.wait_get_run_job_terminated_or_skipped, - response=CancelRunResponse.from_dict(op_response), - run_id=run_id, - ) - - def cancel_run_and_wait(self, run_id: int, timeout=timedelta(minutes=20)) -> Run: - return self.cancel_run(run_id=run_id).result(timeout=timeout) + self._api.do("POST", "/api/2.2/jobs/runs/cancel", body=body, headers=headers) def create( self, @@ -8690,8 +10646,8 @@ def create( :param parameters: List[:class:`JobParameterDefinition`] (optional) Job-level parameter definitions :param performance_target: :class:`PerformanceTarget` (optional) - PerformanceTarget defines how performant or cost efficient the execution of run on serverless should - be. + The performance mode on a serverless job. The performance target determines the level of compute + performance or cost-efficiency for the run. :param queue: :class:`QueueSettings` (optional) The queue settings of the job. :param run_as: :class:`JobRunAs` (optional) @@ -8848,16 +10804,18 @@ def get(self, job_id: int, *, page_token: Optional[str] = None) -> Job: Retrieves the details for a single job. - In Jobs API 2.2, requests for a single job support pagination of `tasks` and `job_clusters` when - either exceeds 100 elements. Use the `next_page_token` field to check for more results and pass its - value as the `page_token` in subsequent requests. Arrays with fewer than 100 elements in a page will - be empty on later pages. + Large arrays in the results will be paginated when they exceed 100 elements. A request for a single + job will return all properties for that job, and the first 100 elements of array properties (`tasks`, + `job_clusters`, `environments` and `parameters`). Use the `next_page_token` field to check for more + results and pass its value as the `page_token` in subsequent requests. If any array properties have + more than 100 elements, additional results will be returned on subsequent requests. Arrays without + additional results will be empty on later pages. :param job_id: int The canonical identifier of the job to retrieve information about. This field is required. :param page_token: str (optional) - Use `next_page_token` returned from the previous GetJob to request the next page of the job's - sub-resources. + Use `next_page_token` returned from the previous GetJob response to request the next page of the + job's array properties. :returns: :class:`Job` """ @@ -8922,10 +10880,12 @@ def get_run( Retrieves the metadata of a run. - In Jobs API 2.2, requests for a single job run support pagination of `tasks` and `job_clusters` when - either exceeds 100 elements. Use the `next_page_token` field to check for more results and pass its - value as the `page_token` in subsequent requests. Arrays with fewer than 100 elements in a page will - be empty on later pages. + Large arrays in the results will be paginated when they exceed 100 elements. A request for a single + run will return all properties for that run, and the first 100 elements of array properties (`tasks`, + `job_clusters`, `job_parameters` and `repair_history`). Use the next_page_token field to check for + more results and pass its value as the page_token in subsequent requests. If any array properties have + more than 100 elements, additional results will be returned on subsequent requests. Arrays without + additional results will be empty on later pages. :param run_id: int The canonical identifier of the run for which to retrieve the metadata. This field is required. @@ -8934,8 +10894,8 @@ def get_run( :param include_resolved_values: bool (optional) Whether to include resolved parameter values in the response. :param page_token: str (optional) - Use `next_page_token` returned from the previous GetRun to request the next page of the run's - sub-resources. + Use `next_page_token` returned from the previous GetRun response to request the next page of the + run's array properties. :returns: :class:`Run` """ @@ -9141,7 +11101,7 @@ def repair_run( rerun_tasks: Optional[List[str]] = None, spark_submit_params: Optional[List[str]] = None, sql_params: Optional[Dict[str, str]] = None, - ) -> Wait[Run]: + ) -> RepairRunResponse: """Repair a job run. Re-run one or more tasks. Tasks are re-run as part of the original job run. They use the current job @@ -9230,7 +11190,7 @@ def repair_run( :returns: Long-running operation waiter for :class:`Run`. - See :method:wait_get_run_job_terminated_or_skipped for more details. + See :method:WaitGetRunJobTerminatedOrSkipped for more details. """ body = {} if dbt_commands is not None: @@ -9266,48 +11226,8 @@ def repair_run( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.2/jobs/runs/repair", body=body, headers=headers) - return Wait( - self.wait_get_run_job_terminated_or_skipped, - response=RepairRunResponse.from_dict(op_response), - run_id=run_id, - ) - - def repair_run_and_wait( - self, - run_id: int, - *, - dbt_commands: Optional[List[str]] = None, - jar_params: Optional[List[str]] = None, - job_parameters: Optional[Dict[str, str]] = None, - latest_repair_id: Optional[int] = None, - notebook_params: Optional[Dict[str, str]] = None, - pipeline_params: Optional[PipelineParams] = None, - python_named_params: Optional[Dict[str, str]] = None, - python_params: Optional[List[str]] = None, - rerun_all_failed_tasks: Optional[bool] = None, - rerun_dependent_tasks: Optional[bool] = None, - rerun_tasks: Optional[List[str]] = None, - spark_submit_params: Optional[List[str]] = None, - sql_params: Optional[Dict[str, str]] = None, - timeout=timedelta(minutes=20), - ) -> Run: - return self.repair_run( - dbt_commands=dbt_commands, - jar_params=jar_params, - job_parameters=job_parameters, - latest_repair_id=latest_repair_id, - notebook_params=notebook_params, - pipeline_params=pipeline_params, - python_named_params=python_named_params, - python_params=python_params, - rerun_all_failed_tasks=rerun_all_failed_tasks, - rerun_dependent_tasks=rerun_dependent_tasks, - rerun_tasks=rerun_tasks, - run_id=run_id, - spark_submit_params=spark_submit_params, - sql_params=sql_params, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.2/jobs/runs/repair", body=body, headers=headers) + return RepairRunResponse.from_dict(res) def reset(self, job_id: int, new_settings: JobSettings): """Update all job settings (reset). @@ -9353,7 +11273,7 @@ def run_now( queue: Optional[QueueSettings] = None, spark_submit_params: Optional[List[str]] = None, sql_params: Optional[Dict[str, str]] = None, - ) -> Wait[Run]: + ) -> RunNowResponse: """Trigger a new job run. Run a job and return the `run_id` of the triggered run. @@ -9408,9 +11328,9 @@ def run_now( A list of task keys to run inside of the job. If this field is not provided, all tasks in the job will be run. :param performance_target: :class:`PerformanceTarget` (optional) - PerformanceTarget defines how performant or cost efficient the execution of run on serverless - compute should be. For RunNow, this performance target will override the target defined on the - job-level. + The performance mode on a serverless job. The performance target determines the level of compute + performance or cost-efficiency for the run. This field overrides the performance target defined on + the job-level. :param pipeline_params: :class:`PipelineParams` (optional) Controls whether the pipeline should perform a full refresh :param python_named_params: Dict[str,str] (optional) @@ -9453,7 +11373,7 @@ def run_now( :returns: Long-running operation waiter for :class:`Run`. - See :method:wait_get_run_job_terminated_or_skipped for more details. + See :method:WaitGetRunJobTerminatedOrSkipped for more details. """ body = {} if dbt_commands is not None: @@ -9489,48 +11409,8 @@ def run_now( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.2/jobs/run-now", body=body, headers=headers) - return Wait( - self.wait_get_run_job_terminated_or_skipped, - response=RunNowResponse.from_dict(op_response), - run_id=op_response["run_id"], - ) - - def run_now_and_wait( - self, - job_id: int, - *, - dbt_commands: Optional[List[str]] = None, - idempotency_token: Optional[str] = None, - jar_params: Optional[List[str]] = None, - job_parameters: Optional[Dict[str, str]] = None, - notebook_params: Optional[Dict[str, str]] = None, - only: Optional[List[str]] = None, - performance_target: Optional[PerformanceTarget] = None, - pipeline_params: Optional[PipelineParams] = None, - python_named_params: Optional[Dict[str, str]] = None, - python_params: Optional[List[str]] = None, - queue: Optional[QueueSettings] = None, - spark_submit_params: Optional[List[str]] = None, - sql_params: Optional[Dict[str, str]] = None, - timeout=timedelta(minutes=20), - ) -> Run: - return self.run_now( - dbt_commands=dbt_commands, - idempotency_token=idempotency_token, - jar_params=jar_params, - job_id=job_id, - job_parameters=job_parameters, - notebook_params=notebook_params, - only=only, - performance_target=performance_target, - pipeline_params=pipeline_params, - python_named_params=python_named_params, - python_params=python_params, - queue=queue, - spark_submit_params=spark_submit_params, - sql_params=sql_params, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.2/jobs/run-now", body=body, headers=headers) + return RunNowResponse.from_dict(res) def set_permissions( self, job_id: str, *, access_control_list: Optional[List[JobAccessControlRequest]] = None @@ -9574,7 +11454,7 @@ def submit( tasks: Optional[List[SubmitTask]] = None, timeout_seconds: Optional[int] = None, webhook_notifications: Optional[WebhookNotifications] = None, - ) -> Wait[Run]: + ) -> SubmitRunResponse: """Create and trigger a one-time run. Submit a one-time run. This endpoint allows you to submit a workload directly without creating a job. @@ -9632,7 +11512,7 @@ def submit( :returns: Long-running operation waiter for :class:`Run`. - See :method:wait_get_run_job_terminated_or_skipped for more details. + See :method:WaitGetRunJobTerminatedOrSkipped for more details. """ body = {} if access_control_list is not None: @@ -9668,48 +11548,8 @@ def submit( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.2/jobs/runs/submit", body=body, headers=headers) - return Wait( - self.wait_get_run_job_terminated_or_skipped, - response=SubmitRunResponse.from_dict(op_response), - run_id=op_response["run_id"], - ) - - def submit_and_wait( - self, - *, - access_control_list: Optional[List[JobAccessControlRequest]] = None, - budget_policy_id: Optional[str] = None, - email_notifications: Optional[JobEmailNotifications] = None, - environments: Optional[List[JobEnvironment]] = None, - git_source: Optional[GitSource] = None, - health: Optional[JobsHealthRules] = None, - idempotency_token: Optional[str] = None, - notification_settings: Optional[JobNotificationSettings] = None, - queue: Optional[QueueSettings] = None, - run_as: Optional[JobRunAs] = None, - run_name: Optional[str] = None, - tasks: Optional[List[SubmitTask]] = None, - timeout_seconds: Optional[int] = None, - webhook_notifications: Optional[WebhookNotifications] = None, - timeout=timedelta(minutes=20), - ) -> Run: - return self.submit( - access_control_list=access_control_list, - budget_policy_id=budget_policy_id, - email_notifications=email_notifications, - environments=environments, - git_source=git_source, - health=health, - idempotency_token=idempotency_token, - notification_settings=notification_settings, - queue=queue, - run_as=run_as, - run_name=run_name, - tasks=tasks, - timeout_seconds=timeout_seconds, - webhook_notifications=webhook_notifications, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.2/jobs/runs/submit", body=body, headers=headers) + return SubmitRunResponse.from_dict(res) def update( self, job_id: int, *, fields_to_remove: Optional[List[str]] = None, new_settings: Optional[JobSettings] = None diff --git a/databricks/sdk/mixins/jobs.py b/databricks/sdk/jobs/v2/mixin.py similarity index 99% rename from databricks/sdk/mixins/jobs.py rename to databricks/sdk/jobs/v2/mixin.py index 1e6cf25d5..9618143d6 100644 --- a/databricks/sdk/mixins/jobs.py +++ b/databricks/sdk/jobs/v2/mixin.py @@ -1,7 +1,7 @@ from typing import Iterator, Optional -from databricks.sdk.service import jobs -from databricks.sdk.service.jobs import BaseJob, BaseRun, Job, RunType +from . import jobs +from .jobs import BaseJob, BaseRun, Job, RunType class JobsExt(jobs.JobsAPI): diff --git a/databricks/sdk/marketplace/__init__.py b/databricks/sdk/marketplace/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/marketplace/v2/__init__.py b/databricks/sdk/marketplace/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/marketplace/v2/client.py b/databricks/sdk/marketplace/v2/client.py new file mode 100755 index 000000000..0ce4fb7bd --- /dev/null +++ b/databricks/sdk/marketplace/v2/client.py @@ -0,0 +1,817 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .marketplace import (ConsumerFulfillmentsAPI, ConsumerInstallationsAPI, + ConsumerListingsAPI, + ConsumerPersonalizationRequestsAPI, + ConsumerProvidersAPI, ProviderExchangeFiltersAPI, + ProviderExchangesAPI, ProviderFilesAPI, + ProviderListingsAPI, + ProviderPersonalizationRequestsAPI, + ProviderProviderAnalyticsDashboardsAPI, + ProviderProvidersAPI) + +_LOG = logging.getLogger(__name__) + + +class ConsumerFulfillmentsClient(ConsumerFulfillmentsAPI): + """ + Fulfillments are entities that allow consumers to preview installations. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ConsumerInstallationsClient(ConsumerInstallationsAPI): + """ + Installations are entities that allow consumers to interact with Databricks Marketplace + listings. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ConsumerListingsClient(ConsumerListingsAPI): + """ + Listings are the core entities in the Marketplace. They represent the products that are + available for consumption. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ConsumerPersonalizationRequestsClient(ConsumerPersonalizationRequestsAPI): + """ + Personalization Requests allow customers to interact with the individualized Marketplace listing + flow. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ConsumerProvidersClient(ConsumerProvidersAPI): + """ + Providers are the entities that publish listings to the Marketplace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderExchangeFiltersClient(ProviderExchangeFiltersAPI): + """ + Marketplace exchanges filters curate which groups can access an exchange. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderExchangesClient(ProviderExchangesAPI): + """ + Marketplace exchanges allow providers to share their listings with a curated set of customers. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderFilesClient(ProviderFilesAPI): + """ + Marketplace offers a set of file APIs for various purposes such as preview notebooks and + provider icons. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderListingsClient(ProviderListingsAPI): + """ + Listings are the core entities in the Marketplace. They represent the products that are + available for consumption. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderPersonalizationRequestsClient(ProviderPersonalizationRequestsAPI): + """ + Personalization requests are an alternate to instantly available listings. Control the lifecycle + of personalized solutions. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderProviderAnalyticsDashboardsClient(ProviderProviderAnalyticsDashboardsAPI): + """ + Manage templated analytics solution for providers. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ProviderProvidersClient(ProviderProvidersAPI): + """ + Providers are entities that manage assets in Marketplace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/marketplace.py b/databricks/sdk/marketplace/v2/marketplace.py similarity index 99% rename from databricks/sdk/service/marketplace.py rename to databricks/sdk/marketplace/v2/marketplace.py index 41992fd69..d249c3966 100755 --- a/databricks/sdk/service/marketplace.py +++ b/databricks/sdk/marketplace/v2/marketplace.py @@ -7,11 +7,11 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -1948,7 +1948,6 @@ def from_dict(cls, d: Dict[str, Any]) -> ListProvidersResponse: @dataclass class Listing: summary: ListingSummary - """Next Number: 26""" detail: Optional[ListingDetail] = None @@ -2252,8 +2251,6 @@ class ListingStatus(Enum): @dataclass class ListingSummary: - """Next Number: 26""" - name: str listing_type: ListingType diff --git a/databricks/sdk/ml/__init__.py b/databricks/sdk/ml/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/ml/v2/__init__.py b/databricks/sdk/ml/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/ml/v2/client.py b/databricks/sdk/ml/v2/client.py new file mode 100755 index 000000000..56b465241 --- /dev/null +++ b/databricks/sdk/ml/v2/client.py @@ -0,0 +1,221 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .ml import ExperimentsAPI, ForecastingAPI, ModelRegistryAPI + +_LOG = logging.getLogger(__name__) + + +class ExperimentsClient(ExperimentsAPI): + """ + Experiments are the primary unit of organization in MLflow; all MLflow runs belong to an + experiment. Each experiment lets you visualize, search, and compare runs, as well as download + run artifacts or metadata for analysis in other tools. Experiments are maintained in a + Databricks hosted MLflow tracking server. + + Experiments are located in the workspace file tree. You manage experiments using the same tools + you use to manage other workspace objects such as folders, notebooks, and libraries. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ForecastingClient(ForecastingAPI): + """ + The Forecasting API allows you to create and get serverless forecasting experiments + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ModelRegistryClient(ModelRegistryAPI): + """ + Note: This API reference documents APIs for the Workspace Model Registry. Databricks recommends + using [Models in Unity Catalog](/api/workspace/registeredmodels) instead. Models in Unity + Catalog provides centralized model governance, cross-workspace access, lineage, and deployment. + Workspace Model Registry will be deprecated in the future. + + The Workspace Model Registry is a centralized model repository and a UI and set of APIs that + enable you to manage the full lifecycle of MLflow Models. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/ml/v2/ml.py similarity index 93% rename from databricks/sdk/service/ml.py rename to databricks/sdk/ml/v2/ml.py index 4a17541af..44f5f28c5 100755 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/ml/v2/ml.py @@ -3,19 +3,15 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -271,6 +267,109 @@ def from_dict(cls, d: Dict[str, Any]) -> ApproveTransitionRequestResponse: return cls(activity=_from_dict(d, "activity", Activity)) +@dataclass +class ArtifactCredentialInfo: + headers: Optional[List[ArtifactCredentialInfoHttpHeader]] = None + """A collection of HTTP headers that should be specified when uploading to or downloading from the + specified `signed_uri`.""" + + path: Optional[str] = None + """The path, relative to the Run's artifact root location, of the artifact that can be accessed + with the credential.""" + + run_id: Optional[str] = None + """The ID of the MLflow Run containing the artifact that can be accessed with the credential.""" + + signed_uri: Optional[str] = None + """The signed URI credential that provides access to the artifact.""" + + type: Optional[ArtifactCredentialType] = None + """The type of the signed credential URI (e.g., an AWS presigned URL or an Azure Shared Access + Signature URI).""" + + def as_dict(self) -> dict: + """Serializes the ArtifactCredentialInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.headers: + body["headers"] = [v.as_dict() for v in self.headers] + if self.path is not None: + body["path"] = self.path + if self.run_id is not None: + body["run_id"] = self.run_id + if self.signed_uri is not None: + body["signed_uri"] = self.signed_uri + if self.type is not None: + body["type"] = self.type.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ArtifactCredentialInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.headers: + body["headers"] = self.headers + if self.path is not None: + body["path"] = self.path + if self.run_id is not None: + body["run_id"] = self.run_id + if self.signed_uri is not None: + body["signed_uri"] = self.signed_uri + if self.type is not None: + body["type"] = self.type + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ArtifactCredentialInfo: + """Deserializes the ArtifactCredentialInfo from a dictionary.""" + return cls( + headers=_repeated_dict(d, "headers", ArtifactCredentialInfoHttpHeader), + path=d.get("path", None), + run_id=d.get("run_id", None), + signed_uri=d.get("signed_uri", None), + type=_enum(d, "type", ArtifactCredentialType), + ) + + +@dataclass +class ArtifactCredentialInfoHttpHeader: + name: Optional[str] = None + """The HTTP header name.""" + + value: Optional[str] = None + """The HTTP header value.""" + + def as_dict(self) -> dict: + """Serializes the ArtifactCredentialInfoHttpHeader into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.value is not None: + body["value"] = self.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ArtifactCredentialInfoHttpHeader into a shallow dictionary of its immediate attributes.""" + body = {} + if self.name is not None: + body["name"] = self.name + if self.value is not None: + body["value"] = self.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ArtifactCredentialInfoHttpHeader: + """Deserializes the ArtifactCredentialInfoHttpHeader from a dictionary.""" + return cls(name=d.get("name", None), value=d.get("value", None)) + + +class ArtifactCredentialType(Enum): + """The type of a given artifact access credential""" + + AWS_PRESIGNED_URL = "AWS_PRESIGNED_URL" + AZURE_ADLS_GEN2_SAS_URI = "AZURE_ADLS_GEN2_SAS_URI" + AZURE_SAS_URI = "AZURE_SAS_URI" + GCP_SIGNED_URL = "GCP_SIGNED_URL" + + class CommentActivityAction(Enum): """An action that a user (with sufficient permissions) could take on a comment. Valid values are: * `EDIT_COMMENT`: Edit the comment @@ -489,63 +588,67 @@ def from_dict(cls, d: Dict[str, Any]) -> CreateExperimentResponse: @dataclass class CreateForecastingExperimentRequest: train_data_path: str - """The three-level (fully qualified) name of a unity catalog table. This table serves as the - training data for the forecasting model.""" + """The fully qualified name of a Unity Catalog table, formatted as + catalog_name.schema_name.table_name, used as training data for the forecasting model.""" target_column: str - """Name of the column in the input training table that serves as the prediction target. The values - in this column will be used as the ground truth for model training.""" + """The column in the input training table used as the prediction target for model training. The + values in this column are used as the ground truth for model training.""" time_column: str - """Name of the column in the input training table that represents the timestamp of each row.""" + """The column in the input training table that represents each row's timestamp.""" forecast_granularity: str - """The granularity of the forecast. This defines the time interval between consecutive rows in the - time series data. Possible values: '1 second', '1 minute', '5 minutes', '10 minutes', '15 - minutes', '30 minutes', 'Hourly', 'Daily', 'Weekly', 'Monthly', 'Quarterly', 'Yearly'.""" + """The time interval between consecutive rows in the time series data. Possible values include: '1 + second', '1 minute', '5 minutes', '10 minutes', '15 minutes', '30 minutes', 'Hourly', 'Daily', + 'Weekly', 'Monthly', 'Quarterly', 'Yearly'.""" forecast_horizon: int - """The number of time steps into the future for which predictions should be made. This value - represents a multiple of forecast_granularity determining how far ahead the model will forecast.""" + """The number of time steps into the future to make predictions, calculated as a multiple of + forecast_granularity. This value represents how far ahead the model should forecast.""" custom_weights_column: Optional[str] = None - """Name of the column in the input training table used to customize the weight for each time series - to calculate weighted metrics.""" + """The column in the training table used to customize weights for each time series.""" experiment_path: Optional[str] = None - """The path to the created experiment. This is the path where the experiment will be stored in the - workspace.""" + """The path in the workspace to store the created experiment.""" holiday_regions: Optional[List[str]] = None - """Region code(s) to consider when automatically adding holiday features. When empty, no holiday - features are added. Only supports 1 holiday region for now.""" + """The region code(s) to automatically add holiday features. Currently supports only one region.""" + + include_features: Optional[List[str]] = None + """Specifies the list of feature columns to include in model training. These columns must exist in + the training data and be of type string, numerical, or boolean. If not specified, no additional + features will be included. Note: Certain columns are automatically handled: - Automatically + excluded: split_column, target_column, custom_weights_column. - Automatically included: + time_column.""" max_runtime: Optional[int] = None - """The maximum duration in minutes for which the experiment is allowed to run. If the experiment - exceeds this time limit it will be stopped automatically.""" + """The maximum duration for the experiment in minutes. The experiment stops automatically if it + exceeds this limit.""" prediction_data_path: Optional[str] = None - """The three-level (fully qualified) path to a unity catalog table. This table path serves to store - the predictions.""" + """The fully qualified path of a Unity Catalog table, formatted as + catalog_name.schema_name.table_name, used to store predictions.""" primary_metric: Optional[str] = None """The evaluation metric used to optimize the forecasting model.""" register_to: Optional[str] = None - """The three-level (fully qualified) path to a unity catalog model. This model path serves to store - the best model.""" + """The fully qualified path of a Unity Catalog model, formatted as + catalog_name.schema_name.model_name, used to store the best model.""" split_column: Optional[str] = None - """Name of the column in the input training table used for custom data splits. The values in this - column must be "train", "validate", or "test" to indicate which split each row belongs to.""" + """// The column in the training table used for custom data splits. Values must be 'train', + 'validate', or 'test'.""" timeseries_identifier_columns: Optional[List[str]] = None - """Name of the column in the input training table used to group the dataset to predict individual - time series""" + """The column in the training table used to group the dataset for predicting individual time + series.""" training_frameworks: Optional[List[str]] = None - """The list of frameworks to include for model tuning. Possible values: 'Prophet', 'ARIMA', - 'DeepAR'. An empty list will include all supported frameworks.""" + """List of frameworks to include for model tuning. Possible values are 'Prophet', 'ARIMA', + 'DeepAR'. An empty list includes all supported frameworks.""" def as_dict(self) -> dict: """Serializes the CreateForecastingExperimentRequest into a dictionary suitable for use as a JSON request body.""" @@ -560,6 +663,8 @@ def as_dict(self) -> dict: body["forecast_horizon"] = self.forecast_horizon if self.holiday_regions: body["holiday_regions"] = [v for v in self.holiday_regions] + if self.include_features: + body["include_features"] = [v for v in self.include_features] if self.max_runtime is not None: body["max_runtime"] = self.max_runtime if self.prediction_data_path is not None: @@ -595,6 +700,8 @@ def as_shallow_dict(self) -> dict: body["forecast_horizon"] = self.forecast_horizon if self.holiday_regions: body["holiday_regions"] = self.holiday_regions + if self.include_features: + body["include_features"] = self.include_features if self.max_runtime is not None: body["max_runtime"] = self.max_runtime if self.prediction_data_path is not None: @@ -626,6 +733,7 @@ def from_dict(cls, d: Dict[str, Any]) -> CreateForecastingExperimentRequest: forecast_granularity=d.get("forecast_granularity", None), forecast_horizon=d.get("forecast_horizon", None), holiday_regions=d.get("holiday_regions", None), + include_features=d.get("include_features", None), max_runtime=d.get("max_runtime", None), prediction_data_path=d.get("prediction_data_path", None), primary_metric=d.get("primary_metric", None), @@ -1946,13 +2054,13 @@ class FileInfo: """Metadata of a single artifact file or directory.""" file_size: Optional[int] = None - """Size in bytes. Unset for directories.""" + """The size in bytes of the file. Unset for directories.""" is_dir: Optional[bool] = None """Whether the path is a directory.""" path: Optional[str] = None - """Path relative to the root artifact directory run.""" + """The path relative to the root artifact directory run.""" def as_dict(self) -> dict: """Serializes the FileInfo into a dictionary suitable for use as a JSON request body.""" @@ -2036,6 +2144,56 @@ class ForecastingExperimentState(Enum): SUCCEEDED = "SUCCEEDED" +@dataclass +class GetCredentialsForTraceDataDownloadResponse: + credential_info: Optional[ArtifactCredentialInfo] = None + """The artifact download credentials for the specified trace data.""" + + def as_dict(self) -> dict: + """Serializes the GetCredentialsForTraceDataDownloadResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.credential_info: + body["credential_info"] = self.credential_info.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GetCredentialsForTraceDataDownloadResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.credential_info: + body["credential_info"] = self.credential_info + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GetCredentialsForTraceDataDownloadResponse: + """Deserializes the GetCredentialsForTraceDataDownloadResponse from a dictionary.""" + return cls(credential_info=_from_dict(d, "credential_info", ArtifactCredentialInfo)) + + +@dataclass +class GetCredentialsForTraceDataUploadResponse: + credential_info: Optional[ArtifactCredentialInfo] = None + """The artifact upload credentials for the specified trace data.""" + + def as_dict(self) -> dict: + """Serializes the GetCredentialsForTraceDataUploadResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.credential_info: + body["credential_info"] = self.credential_info.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GetCredentialsForTraceDataUploadResponse into a shallow dictionary of its immediate attributes.""" + body = {} + if self.credential_info: + body["credential_info"] = self.credential_info + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GetCredentialsForTraceDataUploadResponse: + """Deserializes the GetCredentialsForTraceDataUploadResponse from a dictionary.""" + return cls(credential_info=_from_dict(d, "credential_info", ArtifactCredentialInfo)) + + @dataclass class GetExperimentByNameResponse: experiment: Optional[Experiment] = None @@ -2536,13 +2694,13 @@ def from_dict(cls, d: Dict[str, Any]) -> JobSpecWithoutSecret: @dataclass class ListArtifactsResponse: files: Optional[List[FileInfo]] = None - """File location and metadata for artifacts.""" + """The file location and metadata for artifacts.""" next_page_token: Optional[str] = None - """Token that can be used to retrieve the next page of artifact results""" + """The token that can be used to retrieve the next page of artifact results.""" root_uri: Optional[str] = None - """Root artifact directory for the run.""" + """The root artifact directory for the run.""" def as_dict(self) -> dict: """Serializes the ListArtifactsResponse into a dictionary suitable for use as a JSON request body.""" @@ -2784,11 +2942,16 @@ class LogInputs: datasets: Optional[List[DatasetInput]] = None """Dataset inputs""" + models: Optional[List[ModelInput]] = None + """Model inputs""" + def as_dict(self) -> dict: """Serializes the LogInputs into a dictionary suitable for use as a JSON request body.""" body = {} if self.datasets: body["datasets"] = [v.as_dict() for v in self.datasets] + if self.models: + body["models"] = [v.as_dict() for v in self.models] if self.run_id is not None: body["run_id"] = self.run_id return body @@ -2798,6 +2961,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.datasets: body["datasets"] = self.datasets + if self.models: + body["models"] = self.models if self.run_id is not None: body["run_id"] = self.run_id return body @@ -2805,7 +2970,11 @@ def as_shallow_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, Any]) -> LogInputs: """Deserializes the LogInputs from a dictionary.""" - return cls(datasets=_repeated_dict(d, "datasets", DatasetInput), run_id=d.get("run_id", None)) + return cls( + datasets=_repeated_dict(d, "datasets", DatasetInput), + models=_repeated_dict(d, "models", ModelInput), + run_id=d.get("run_id", None), + ) @dataclass @@ -2837,6 +3006,17 @@ class LogMetric: timestamp: int """Unix timestamp in milliseconds at the time metric was logged.""" + dataset_digest: Optional[str] = None + """Dataset digest of the dataset associated with the metric, e.g. an md5 hash of the dataset that + uniquely identifies it within datasets of the same name.""" + + dataset_name: Optional[str] = None + """The name of the dataset associated with the metric. E.g. “my.uc.table@2” + “nyc-taxi-dataset”, “fantastic-elk-3”""" + + model_id: Optional[str] = None + """ID of the logged model associated with the metric, if applicable""" + run_id: Optional[str] = None """ID of the run under which to log the metric. Must be provided.""" @@ -2850,8 +3030,14 @@ class LogMetric: def as_dict(self) -> dict: """Serializes the LogMetric into a dictionary suitable for use as a JSON request body.""" body = {} + if self.dataset_digest is not None: + body["dataset_digest"] = self.dataset_digest + if self.dataset_name is not None: + body["dataset_name"] = self.dataset_name if self.key is not None: body["key"] = self.key + if self.model_id is not None: + body["model_id"] = self.model_id if self.run_id is not None: body["run_id"] = self.run_id if self.run_uuid is not None: @@ -2867,8 +3053,14 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the LogMetric into a shallow dictionary of its immediate attributes.""" body = {} + if self.dataset_digest is not None: + body["dataset_digest"] = self.dataset_digest + if self.dataset_name is not None: + body["dataset_name"] = self.dataset_name if self.key is not None: body["key"] = self.key + if self.model_id is not None: + body["model_id"] = self.model_id if self.run_id is not None: body["run_id"] = self.run_id if self.run_uuid is not None: @@ -2885,7 +3077,10 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> LogMetric: """Deserializes the LogMetric from a dictionary.""" return cls( + dataset_digest=d.get("dataset_digest", None), + dataset_name=d.get("dataset_name", None), key=d.get("key", None), + model_id=d.get("model_id", None), run_id=d.get("run_id", None), run_uuid=d.get("run_uuid", None), step=d.get("step", None), @@ -3036,23 +3231,46 @@ def from_dict(cls, d: Dict[str, Any]) -> LogParamResponse: class Metric: """Metric associated with a run, represented as a key-value pair.""" + dataset_digest: Optional[str] = None + """The dataset digest of the dataset associated with the metric, e.g. an md5 hash of the dataset + that uniquely identifies it within datasets of the same name.""" + + dataset_name: Optional[str] = None + """The name of the dataset associated with the metric. E.g. “my.uc.table@2” + “nyc-taxi-dataset”, “fantastic-elk-3”""" + key: Optional[str] = None - """Key identifying this metric.""" + """The key identifying the metric.""" + + model_id: Optional[str] = None + """The ID of the logged model or registered model version associated with the metric, if + applicable.""" + + run_id: Optional[str] = None + """The ID of the run containing the metric.""" step: Optional[int] = None - """Step at which to log the metric.""" + """The step at which the metric was logged.""" timestamp: Optional[int] = None - """The timestamp at which this metric was recorded.""" + """The timestamp at which the metric was recorded.""" value: Optional[float] = None - """Value associated with this metric.""" + """The value of the metric.""" def as_dict(self) -> dict: """Serializes the Metric into a dictionary suitable for use as a JSON request body.""" body = {} + if self.dataset_digest is not None: + body["dataset_digest"] = self.dataset_digest + if self.dataset_name is not None: + body["dataset_name"] = self.dataset_name if self.key is not None: body["key"] = self.key + if self.model_id is not None: + body["model_id"] = self.model_id + if self.run_id is not None: + body["run_id"] = self.run_id if self.step is not None: body["step"] = self.step if self.timestamp is not None: @@ -3064,8 +3282,16 @@ def as_dict(self) -> dict: def as_shallow_dict(self) -> dict: """Serializes the Metric into a shallow dictionary of its immediate attributes.""" body = {} + if self.dataset_digest is not None: + body["dataset_digest"] = self.dataset_digest + if self.dataset_name is not None: + body["dataset_name"] = self.dataset_name if self.key is not None: body["key"] = self.key + if self.model_id is not None: + body["model_id"] = self.model_id + if self.run_id is not None: + body["run_id"] = self.run_id if self.step is not None: body["step"] = self.step if self.timestamp is not None: @@ -3078,7 +3304,11 @@ def as_shallow_dict(self) -> dict: def from_dict(cls, d: Dict[str, Any]) -> Metric: """Deserializes the Metric from a dictionary.""" return cls( + dataset_digest=d.get("dataset_digest", None), + dataset_name=d.get("dataset_name", None), key=d.get("key", None), + model_id=d.get("model_id", None), + run_id=d.get("run_id", None), step=d.get("step", None), timestamp=d.get("timestamp", None), value=d.get("value", None), @@ -3253,6 +3483,33 @@ def from_dict(cls, d: Dict[str, Any]) -> ModelDatabricks: ) +@dataclass +class ModelInput: + """Represents a LoggedModel or Registered Model Version input to a Run.""" + + model_id: str + """The unique identifier of the model.""" + + def as_dict(self) -> dict: + """Serializes the ModelInput into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.model_id is not None: + body["model_id"] = self.model_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ModelInput into a shallow dictionary of its immediate attributes.""" + body = {} + if self.model_id is not None: + body["model_id"] = self.model_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ModelInput: + """Deserializes the ModelInput from a dictionary.""" + return cls(model_id=d.get("model_id", None)) + + @dataclass class ModelTag: key: Optional[str] = None @@ -4571,11 +4828,19 @@ class RunInputs: dataset_inputs: Optional[List[DatasetInput]] = None """Run metrics.""" + model_inputs: Optional[List[ModelInput]] = None + """**NOTE**: Experimental: This API field may change or be removed in a future release without + warning. + + Model inputs to the Run.""" + def as_dict(self) -> dict: """Serializes the RunInputs into a dictionary suitable for use as a JSON request body.""" body = {} if self.dataset_inputs: body["dataset_inputs"] = [v.as_dict() for v in self.dataset_inputs] + if self.model_inputs: + body["model_inputs"] = [v.as_dict() for v in self.model_inputs] return body def as_shallow_dict(self) -> dict: @@ -4583,12 +4848,17 @@ def as_shallow_dict(self) -> dict: body = {} if self.dataset_inputs: body["dataset_inputs"] = self.dataset_inputs + if self.model_inputs: + body["model_inputs"] = self.model_inputs return body @classmethod def from_dict(cls, d: Dict[str, Any]) -> RunInputs: """Deserializes the RunInputs from a dictionary.""" - return cls(dataset_inputs=_repeated_dict(d, "dataset_inputs", DatasetInput)) + return cls( + dataset_inputs=_repeated_dict(d, "dataset_inputs", DatasetInput), + model_inputs=_repeated_dict(d, "model_inputs", ModelInput), + ) @dataclass @@ -6101,6 +6371,38 @@ def get_by_name(self, experiment_name: str) -> GetExperimentByNameResponse: res = self._api.do("GET", "/api/2.0/mlflow/experiments/get-by-name", query=query, headers=headers) return GetExperimentByNameResponse.from_dict(res) + def get_credentials_for_trace_data_download(self, request_id: str) -> GetCredentialsForTraceDataDownloadResponse: + """Get credentials to download trace data. + + :param request_id: str + The ID of the trace to fetch artifact download credentials for. + + :returns: :class:`GetCredentialsForTraceDataDownloadResponse` + """ + + headers = { + "Accept": "application/json", + } + + res = self._api.do("GET", f"/api/2.0/mlflow/traces/{request_id}/credentials-for-data-download", headers=headers) + return GetCredentialsForTraceDataDownloadResponse.from_dict(res) + + def get_credentials_for_trace_data_upload(self, request_id: str) -> GetCredentialsForTraceDataUploadResponse: + """Get credentials to upload trace data. + + :param request_id: str + The ID of the trace to fetch artifact upload credentials for. + + :returns: :class:`GetCredentialsForTraceDataUploadResponse` + """ + + headers = { + "Accept": "application/json", + } + + res = self._api.do("GET", f"/api/2.0/mlflow/traces/{request_id}/credentials-for-data-upload", headers=headers) + return GetCredentialsForTraceDataUploadResponse.from_dict(res) + def get_experiment(self, experiment_id: str) -> GetExperimentResponse: """Get an experiment. @@ -6257,10 +6559,11 @@ def list_artifacts( API](/api/workspace/files/listdirectorycontents). :param page_token: str (optional) - Token indicating the page of artifact results to fetch. `page_token` is not supported when listing - artifacts in UC Volumes. A maximum of 1000 artifacts will be retrieved for UC Volumes. Please call - `/api/2.0/fs/directories{directory_path}` for listing artifacts in UC Volumes, which supports - pagination. See [List directory contents | Files API](/api/workspace/files/listdirectorycontents). + The token indicating the page of artifact results to fetch. `page_token` is not supported when + listing artifacts in UC Volumes. A maximum of 1000 artifacts will be retrieved for UC Volumes. + Please call `/api/2.0/fs/directories{directory_path}` for listing artifacts in UC Volumes, which + supports pagination. See [List directory contents | Files + API](/api/workspace/files/listdirectorycontents). :param path: str (optional) Filter artifacts matching this path (a relative path from the root artifact directory). :param run_id: str (optional) @@ -6418,7 +6721,9 @@ def log_batch( self._api.do("POST", "/api/2.0/mlflow/runs/log-batch", body=body, headers=headers) - def log_inputs(self, run_id: str, *, datasets: Optional[List[DatasetInput]] = None): + def log_inputs( + self, run_id: str, *, datasets: Optional[List[DatasetInput]] = None, models: Optional[List[ModelInput]] = None + ): """Log inputs to a run. **NOTE:** Experimental: This API may change or be removed in a future release without warning. @@ -6429,12 +6734,16 @@ def log_inputs(self, run_id: str, *, datasets: Optional[List[DatasetInput]] = No ID of the run to log under :param datasets: List[:class:`DatasetInput`] (optional) Dataset inputs + :param models: List[:class:`ModelInput`] (optional) + Model inputs """ body = {} if datasets is not None: body["datasets"] = [v.as_dict() for v in datasets] + if models is not None: + body["models"] = [v.as_dict() for v in models] if run_id is not None: body["run_id"] = run_id headers = { @@ -6450,6 +6759,9 @@ def log_metric( value: float, timestamp: int, *, + dataset_digest: Optional[str] = None, + dataset_name: Optional[str] = None, + model_id: Optional[str] = None, run_id: Optional[str] = None, run_uuid: Optional[str] = None, step: Optional[int] = None, @@ -6466,6 +6778,14 @@ def log_metric( Double value of the metric being logged. :param timestamp: int Unix timestamp in milliseconds at the time metric was logged. + :param dataset_digest: str (optional) + Dataset digest of the dataset associated with the metric, e.g. an md5 hash of the dataset that + uniquely identifies it within datasets of the same name. + :param dataset_name: str (optional) + The name of the dataset associated with the metric. E.g. “my.uc.table@2” “nyc-taxi-dataset”, + “fantastic-elk-3” + :param model_id: str (optional) + ID of the logged model associated with the metric, if applicable :param run_id: str (optional) ID of the run under which to log the metric. Must be provided. :param run_uuid: str (optional) @@ -6477,8 +6797,14 @@ def log_metric( """ body = {} + if dataset_digest is not None: + body["dataset_digest"] = dataset_digest + if dataset_name is not None: + body["dataset_name"] = dataset_name if key is not None: body["key"] = key + if model_id is not None: + body["model_id"] = model_id if run_id is not None: body["run_id"] = run_id if run_uuid is not None: @@ -6947,41 +7273,6 @@ class ForecastingAPI: def __init__(self, api_client): self._api = api_client - def wait_get_experiment_forecasting_succeeded( - self, - experiment_id: str, - timeout=timedelta(minutes=120), - callback: Optional[Callable[[ForecastingExperiment], None]] = None, - ) -> ForecastingExperiment: - deadline = time.time() + timeout.total_seconds() - target_states = (ForecastingExperimentState.SUCCEEDED,) - failure_states = ( - ForecastingExperimentState.FAILED, - ForecastingExperimentState.CANCELLED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get_experiment(experiment_id=experiment_id) - status = poll.state - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach SUCCEEDED, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"experiment_id={experiment_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def create_experiment( self, train_data_path: str, @@ -6993,6 +7284,7 @@ def create_experiment( custom_weights_column: Optional[str] = None, experiment_path: Optional[str] = None, holiday_regions: Optional[List[str]] = None, + include_features: Optional[List[str]] = None, max_runtime: Optional[int] = None, prediction_data_path: Optional[str] = None, primary_metric: Optional[str] = None, @@ -7000,59 +7292,60 @@ def create_experiment( split_column: Optional[str] = None, timeseries_identifier_columns: Optional[List[str]] = None, training_frameworks: Optional[List[str]] = None, - ) -> Wait[ForecastingExperiment]: + ) -> CreateForecastingExperimentResponse: """Create a forecasting experiment. Creates a serverless forecasting experiment. Returns the experiment ID. :param train_data_path: str - The three-level (fully qualified) name of a unity catalog table. This table serves as the training - data for the forecasting model. + The fully qualified name of a Unity Catalog table, formatted as catalog_name.schema_name.table_name, + used as training data for the forecasting model. :param target_column: str - Name of the column in the input training table that serves as the prediction target. The values in - this column will be used as the ground truth for model training. + The column in the input training table used as the prediction target for model training. The values + in this column are used as the ground truth for model training. :param time_column: str - Name of the column in the input training table that represents the timestamp of each row. + The column in the input training table that represents each row's timestamp. :param forecast_granularity: str - The granularity of the forecast. This defines the time interval between consecutive rows in the time - series data. Possible values: '1 second', '1 minute', '5 minutes', '10 minutes', '15 minutes', '30 - minutes', 'Hourly', 'Daily', 'Weekly', 'Monthly', 'Quarterly', 'Yearly'. + The time interval between consecutive rows in the time series data. Possible values include: '1 + second', '1 minute', '5 minutes', '10 minutes', '15 minutes', '30 minutes', 'Hourly', 'Daily', + 'Weekly', 'Monthly', 'Quarterly', 'Yearly'. :param forecast_horizon: int - The number of time steps into the future for which predictions should be made. This value represents - a multiple of forecast_granularity determining how far ahead the model will forecast. + The number of time steps into the future to make predictions, calculated as a multiple of + forecast_granularity. This value represents how far ahead the model should forecast. :param custom_weights_column: str (optional) - Name of the column in the input training table used to customize the weight for each time series to - calculate weighted metrics. + The column in the training table used to customize weights for each time series. :param experiment_path: str (optional) - The path to the created experiment. This is the path where the experiment will be stored in the - workspace. + The path in the workspace to store the created experiment. :param holiday_regions: List[str] (optional) - Region code(s) to consider when automatically adding holiday features. When empty, no holiday - features are added. Only supports 1 holiday region for now. + The region code(s) to automatically add holiday features. Currently supports only one region. + :param include_features: List[str] (optional) + Specifies the list of feature columns to include in model training. These columns must exist in the + training data and be of type string, numerical, or boolean. If not specified, no additional features + will be included. Note: Certain columns are automatically handled: - Automatically excluded: + split_column, target_column, custom_weights_column. - Automatically included: time_column. :param max_runtime: int (optional) - The maximum duration in minutes for which the experiment is allowed to run. If the experiment - exceeds this time limit it will be stopped automatically. + The maximum duration for the experiment in minutes. The experiment stops automatically if it exceeds + this limit. :param prediction_data_path: str (optional) - The three-level (fully qualified) path to a unity catalog table. This table path serves to store the - predictions. + The fully qualified path of a Unity Catalog table, formatted as catalog_name.schema_name.table_name, + used to store predictions. :param primary_metric: str (optional) The evaluation metric used to optimize the forecasting model. :param register_to: str (optional) - The three-level (fully qualified) path to a unity catalog model. This model path serves to store the - best model. + The fully qualified path of a Unity Catalog model, formatted as catalog_name.schema_name.model_name, + used to store the best model. :param split_column: str (optional) - Name of the column in the input training table used for custom data splits. The values in this - column must be "train", "validate", or "test" to indicate which split each row belongs to. + // The column in the training table used for custom data splits. Values must be 'train', 'validate', + or 'test'. :param timeseries_identifier_columns: List[str] (optional) - Name of the column in the input training table used to group the dataset to predict individual time - series + The column in the training table used to group the dataset for predicting individual time series. :param training_frameworks: List[str] (optional) - The list of frameworks to include for model tuning. Possible values: 'Prophet', 'ARIMA', 'DeepAR'. - An empty list will include all supported frameworks. + List of frameworks to include for model tuning. Possible values are 'Prophet', 'ARIMA', 'DeepAR'. An + empty list includes all supported frameworks. :returns: Long-running operation waiter for :class:`ForecastingExperiment`. - See :method:wait_get_experiment_forecasting_succeeded for more details. + See :method:WaitGetExperimentForecastingSucceeded for more details. """ body = {} if custom_weights_column is not None: @@ -7065,6 +7358,8 @@ def create_experiment( body["forecast_horizon"] = forecast_horizon if holiday_regions is not None: body["holiday_regions"] = [v for v in holiday_regions] + if include_features is not None: + body["include_features"] = [v for v in include_features] if max_runtime is not None: body["max_runtime"] = max_runtime if prediction_data_path is not None: @@ -7090,50 +7385,8 @@ def create_experiment( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/automl/create-forecasting-experiment", body=body, headers=headers) - return Wait( - self.wait_get_experiment_forecasting_succeeded, - response=CreateForecastingExperimentResponse.from_dict(op_response), - experiment_id=op_response["experiment_id"], - ) - - def create_experiment_and_wait( - self, - train_data_path: str, - target_column: str, - time_column: str, - forecast_granularity: str, - forecast_horizon: int, - *, - custom_weights_column: Optional[str] = None, - experiment_path: Optional[str] = None, - holiday_regions: Optional[List[str]] = None, - max_runtime: Optional[int] = None, - prediction_data_path: Optional[str] = None, - primary_metric: Optional[str] = None, - register_to: Optional[str] = None, - split_column: Optional[str] = None, - timeseries_identifier_columns: Optional[List[str]] = None, - training_frameworks: Optional[List[str]] = None, - timeout=timedelta(minutes=120), - ) -> ForecastingExperiment: - return self.create_experiment( - custom_weights_column=custom_weights_column, - experiment_path=experiment_path, - forecast_granularity=forecast_granularity, - forecast_horizon=forecast_horizon, - holiday_regions=holiday_regions, - max_runtime=max_runtime, - prediction_data_path=prediction_data_path, - primary_metric=primary_metric, - register_to=register_to, - split_column=split_column, - target_column=target_column, - time_column=time_column, - timeseries_identifier_columns=timeseries_identifier_columns, - train_data_path=train_data_path, - training_frameworks=training_frameworks, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/automl/create-forecasting-experiment", body=body, headers=headers) + return CreateForecastingExperimentResponse.from_dict(res) def get_experiment(self, experiment_id: str) -> ForecastingExperiment: """Get a forecasting experiment. diff --git a/databricks/sdk/oauth2/__init__.py b/databricks/sdk/oauth2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/oauth2/v2/__init__.py b/databricks/sdk/oauth2/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/oauth2/v2/client.py b/databricks/sdk/oauth2/v2/client.py new file mode 100755 index 000000000..94a48913d --- /dev/null +++ b/databricks/sdk/oauth2/v2/client.py @@ -0,0 +1,508 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .oauth2 import (AccountFederationPolicyAPI, CustomAppIntegrationAPI, + OAuthPublishedAppsAPI, PublishedAppIntegrationAPI, + ServicePrincipalFederationPolicyAPI, + ServicePrincipalSecretsAPI) + +_LOG = logging.getLogger(__name__) + + +class AccountFederationPolicyClient(AccountFederationPolicyAPI): + """ + These APIs manage account federation policies. + + Account federation policies allow users and service principals in your Databricks account to + securely access Databricks APIs using tokens from your trusted identity providers (IdPs). + + With token federation, your users and service principals can exchange tokens from your IdP for + Databricks OAuth tokens, which can be used to access Databricks APIs. Token federation + eliminates the need to manage Databricks secrets, and allows you to centralize management of + token issuance policies in your IdP. Databricks token federation is typically used in + combination with [SCIM], so users in your IdP are synchronized into your Databricks account. + + Token federation is configured in your Databricks account using an account federation policy. An + account federation policy specifies: * which IdP, or issuer, your Databricks account should + accept tokens from * how to determine which Databricks user, or subject, a token is issued for + + To configure a federation policy, you provide the following: * The required token __issuer__, as + specified in the “iss” claim of your tokens. The issuer is an https URL that identifies your + IdP. * The allowed token __audiences__, as specified in the “aud” claim of your tokens. This + identifier is intended to represent the recipient of the token. As long as the audience in the + token matches at least one audience in the policy, the token is considered a match. If + unspecified, the default value is your Databricks account id. * The __subject claim__, which + indicates which token claim contains the Databricks username of the user the token was issued + for. If unspecified, the default value is “sub”. * Optionally, the public keys used to + validate the signature of your tokens, in JWKS format. If unspecified (recommended), Databricks + automatically fetches the public keys from your issuer’s well known endpoint. Databricks + strongly recommends relying on your issuer’s well known endpoint for discovering public keys. + + An example federation policy is: ``` issuer: "https://idp.mycompany.com/oidc" audiences: + ["databricks"] subject_claim: "sub" ``` + + An example JWT token body that matches this policy and could be used to authenticate to + Databricks as user `username@mycompany.com` is: ``` { "iss": "https://idp.mycompany.com/oidc", + "aud": "databricks", "sub": "username@mycompany.com" } ``` + + You may also need to configure your IdP to generate tokens for your users to exchange with + Databricks, if your users do not already have the ability to generate tokens that are compatible + with your federation policy. + + You do not need to configure an OAuth application in Databricks to use token federation. + + [SCIM]: https://docs.databricks.com/admin/users-groups/scim/index.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CustomAppIntegrationClient(CustomAppIntegrationAPI): + """ + These APIs enable administrators to manage custom OAuth app integrations, which is required for + adding/using Custom OAuth App Integration like Tableau Cloud for Databricks in AWS cloud. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class OAuthPublishedAppsClient(OAuthPublishedAppsAPI): + """ + These APIs enable administrators to view all the available published OAuth applications in + Databricks. Administrators can add the published OAuth applications to their account through the + OAuth Published App Integration APIs. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PublishedAppIntegrationClient(PublishedAppIntegrationAPI): + """ + These APIs enable administrators to manage published OAuth app integrations, which is required + for adding/using Published OAuth App Integration like Tableau Desktop for Databricks in AWS + cloud. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ServicePrincipalFederationPolicyClient(ServicePrincipalFederationPolicyAPI): + """ + These APIs manage service principal federation policies. + + Service principal federation, also known as Workload Identity Federation, allows your automated + workloads running outside of Databricks to securely access Databricks APIs without the need for + Databricks secrets. With Workload Identity Federation, your application (or workload) + authenticates to Databricks as a Databricks service principal, using tokens provided by the + workload runtime. + + Databricks strongly recommends using Workload Identity Federation to authenticate to Databricks + from automated workloads, over alternatives such as OAuth client secrets or Personal Access + Tokens, whenever possible. Workload Identity Federation is supported by many popular services, + including Github Actions, Azure DevOps, GitLab, Terraform Cloud, and Kubernetes clusters, among + others. + + Workload identity federation is configured in your Databricks account using a service principal + federation policy. A service principal federation policy specifies: * which IdP, or issuer, the + service principal is allowed to authenticate from * which workload identity, or subject, is + allowed to authenticate as the Databricks service principal + + To configure a federation policy, you provide the following: * The required token __issuer__, as + specified in the “iss” claim of workload identity tokens. The issuer is an https URL that + identifies the workload identity provider. * The required token __subject__, as specified in the + “sub” claim of workload identity tokens. The subject uniquely identifies the workload in the + workload runtime environment. * The allowed token __audiences__, as specified in the “aud” + claim of workload identity tokens. The audience is intended to represent the recipient of the + token. As long as the audience in the token matches at least one audience in the policy, the + token is considered a match. If unspecified, the default value is your Databricks account id. * + Optionally, the public keys used to validate the signature of the workload identity tokens, in + JWKS format. If unspecified (recommended), Databricks automatically fetches the public keys from + the issuer’s well known endpoint. Databricks strongly recommends relying on the issuer’s + well known endpoint for discovering public keys. + + An example service principal federation policy, for a Github Actions workload, is: ``` issuer: + "https://token.actions.githubusercontent.com" audiences: ["https://github.com/my-github-org"] + subject: "repo:my-github-org/my-repo:environment:prod" ``` + + An example JWT token body that matches this policy and could be used to authenticate to + Databricks is: ``` { "iss": "https://token.actions.githubusercontent.com", "aud": + "https://github.com/my-github-org", "sub": "repo:my-github-org/my-repo:environment:prod" } ``` + + You may also need to configure the workload runtime to generate tokens for your workloads. + + You do not need to configure an OAuth application in Databricks to use token federation. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ServicePrincipalSecretsClient(ServicePrincipalSecretsAPI): + """ + These APIs enable administrators to manage service principal secrets. + + You can use the generated secrets to obtain OAuth access tokens for a service principal, which + can then be used to access Databricks Accounts and Workspace APIs. For more information, see + [Authentication using OAuth tokens for service principals], + + In addition, the generated secrets can be used to configure the Databricks Terraform Provider to + authenticate with the service principal. For more information, see [Databricks Terraform + Provider]. + + [Authentication using OAuth tokens for service principals]: https://docs.databricks.com/dev-tools/authentication-oauth.html + [Databricks Terraform Provider]: https://github.com/databricks/terraform-provider-databricks/blob/master/docs/index.md#authenticating-with-service-principal + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/oauth2.py b/databricks/sdk/oauth2/v2/oauth2.py similarity index 99% rename from databricks/sdk/service/oauth2.py rename to databricks/sdk/oauth2/v2/oauth2.py index 366f282f4..3888b3176 100755 --- a/databricks/sdk/service/oauth2.py +++ b/databricks/sdk/oauth2/v2/oauth2.py @@ -6,11 +6,10 @@ from dataclasses import dataclass from typing import Any, Dict, Iterator, List, Optional -from ._internal import _from_dict, _repeated_dict +from ...service._internal import _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order diff --git a/databricks/sdk/pipelines/__init__.py b/databricks/sdk/pipelines/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/pipelines/v2/__init__.py b/databricks/sdk/pipelines/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/pipelines/v2/client.py b/databricks/sdk/pipelines/v2/client.py new file mode 100755 index 000000000..6403ca237 --- /dev/null +++ b/databricks/sdk/pipelines/v2/client.py @@ -0,0 +1,89 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .pipelines import PipelinesAPI + +_LOG = logging.getLogger(__name__) + + +class PipelinesClient(PipelinesAPI): + """ + The Delta Live Tables API allows you to create, edit, delete, start, and view details about + pipelines. + + Delta Live Tables is a framework for building reliable, maintainable, and testable data + processing pipelines. You define the transformations to perform on your data, and Delta Live + Tables manages task orchestration, cluster management, monitoring, data quality, and error + handling. + + Instead of defining your data pipelines using a series of separate Apache Spark tasks, Delta + Live Tables manages how your data is transformed based on a target schema you define for each + processing step. You can also enforce data quality with Delta Live Tables expectations. + Expectations allow you to define expected data quality and specify how to handle records that + fail those expectations. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/pipelines/v2/pipelines.py similarity index 80% rename from databricks/sdk/service/pipelines.py rename to databricks/sdk/pipelines/v2/pipelines.py index 3081ef8ee..d6e732f3a 100755 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/pipelines/v2/pipelines.py @@ -3,22 +3,323 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") +# all definitions in this file are in alphabetical order -from databricks.sdk.service import compute -# all definitions in this file are in alphabetical order +@dataclass +class Adlsgen2Info: + """A storage location in Adls Gen2""" + + destination: str + """abfss destination, e.g. + `abfss://@.dfs.core.windows.net/`.""" + + def as_dict(self) -> dict: + """Serializes the Adlsgen2Info into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the Adlsgen2Info into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> Adlsgen2Info: + """Deserializes the Adlsgen2Info from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +@dataclass +class AwsAttributes: + """Attributes set during cluster creation which are related to Amazon Web Services.""" + + availability: Optional[AwsAvailability] = None + """Availability type used for all subsequent nodes past the `first_on_demand` ones. + + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ebs_volume_count: Optional[int] = None + """The number of volumes launched for each instance. Users can choose up to 10 volumes. This + feature is only enabled for supported node types. Legacy node types cannot specify custom EBS + volumes. For node types with no instance store, at least one EBS volume needs to be specified; + otherwise, cluster creation will fail. + + These EBS volumes will be mounted at `/ebs0`, `/ebs1`, and etc. Instance store volumes will be + mounted at `/local_disk0`, `/local_disk1`, and etc. + + If EBS volumes are attached, Databricks will configure Spark to use only the EBS volumes for + scratch storage because heterogenously sized scratch devices can lead to inefficient disk + utilization. If no EBS volumes are attached, Databricks will configure Spark to use instance + store volumes. + + Please note that if EBS volumes are specified, then the Spark configuration `spark.local.dir` + will be overridden.""" + + ebs_volume_iops: Optional[int] = None + """If using gp3 volumes, what IOPS to use for the disk. If this is not set, the maximum performance + of a gp2 volume with the same volume size will be used.""" + + ebs_volume_size: Optional[int] = None + """The size of each EBS volume (in GiB) launched for each instance. For general purpose SSD, this + value must be within the range 100 - 4096. For throughput optimized HDD, this value must be + within the range 500 - 4096.""" + + ebs_volume_throughput: Optional[int] = None + """If using gp3 volumes, what throughput to use for the disk. If this is not set, the maximum + performance of a gp2 volume with the same volume size will be used.""" + + ebs_volume_type: Optional[EbsVolumeType] = None + """The type of EBS volumes that will be launched with this cluster.""" + + first_on_demand: Optional[int] = None + """The first `first_on_demand` nodes of the cluster will be placed on on-demand instances. If this + value is greater than 0, the cluster driver node in particular will be placed on an on-demand + instance. If this value is greater than or equal to the current cluster size, all nodes will be + placed on on-demand instances. If this value is less than the current cluster size, + `first_on_demand` nodes will be placed on on-demand instances and the remainder will be placed + on `availability` instances. Note that this value does not affect cluster size and cannot + currently be mutated over the lifetime of a cluster.""" + + instance_profile_arn: Optional[str] = None + """Nodes for this cluster will only be placed on AWS instances with this instance profile. If + ommitted, nodes will be placed on instances without an IAM instance profile. The instance + profile must have previously been added to the Databricks environment by an account + administrator. + + This feature may only be available to certain customer plans.""" + + spot_bid_price_percent: Optional[int] = None + """The bid price for AWS spot instances, as a percentage of the corresponding instance type's + on-demand price. For example, if this field is set to 50, and the cluster needs a new + `r3.xlarge` spot instance, then the bid price is half of the price of on-demand `r3.xlarge` + instances. Similarly, if this field is set to 200, the bid price is twice the price of on-demand + `r3.xlarge` instances. If not specified, the default value is 100. When spot instances are + requested for this cluster, only spot instances whose bid price percentage matches this field + will be considered. Note that, for safety, we enforce this field to be no more than 10000.""" + + zone_id: Optional[str] = None + """Identifier for the availability zone/datacenter in which the cluster resides. This string will + be of a form like "us-west-2a". The provided availability zone must be in the same region as the + Databricks deployment. For example, "us-west-2a" is not a valid zone id if the Databricks + deployment resides in the "us-east-1" region. This is an optional field at cluster creation, and + if not specified, a default zone will be used. If the zone specified is "auto", will try to + place cluster in a zone with high availability, and will retry placement in a different AZ if + there is not enough capacity. + + The list of available zones as well as the default value can be found by using the `List Zones` + method.""" + + def as_dict(self) -> dict: + """Serializes the AwsAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.ebs_volume_count is not None: + body["ebs_volume_count"] = self.ebs_volume_count + if self.ebs_volume_iops is not None: + body["ebs_volume_iops"] = self.ebs_volume_iops + if self.ebs_volume_size is not None: + body["ebs_volume_size"] = self.ebs_volume_size + if self.ebs_volume_throughput is not None: + body["ebs_volume_throughput"] = self.ebs_volume_throughput + if self.ebs_volume_type is not None: + body["ebs_volume_type"] = self.ebs_volume_type.value + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.instance_profile_arn is not None: + body["instance_profile_arn"] = self.instance_profile_arn + if self.spot_bid_price_percent is not None: + body["spot_bid_price_percent"] = self.spot_bid_price_percent + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AwsAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.ebs_volume_count is not None: + body["ebs_volume_count"] = self.ebs_volume_count + if self.ebs_volume_iops is not None: + body["ebs_volume_iops"] = self.ebs_volume_iops + if self.ebs_volume_size is not None: + body["ebs_volume_size"] = self.ebs_volume_size + if self.ebs_volume_throughput is not None: + body["ebs_volume_throughput"] = self.ebs_volume_throughput + if self.ebs_volume_type is not None: + body["ebs_volume_type"] = self.ebs_volume_type + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.instance_profile_arn is not None: + body["instance_profile_arn"] = self.instance_profile_arn + if self.spot_bid_price_percent is not None: + body["spot_bid_price_percent"] = self.spot_bid_price_percent + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AwsAttributes: + """Deserializes the AwsAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", AwsAvailability), + ebs_volume_count=d.get("ebs_volume_count", None), + ebs_volume_iops=d.get("ebs_volume_iops", None), + ebs_volume_size=d.get("ebs_volume_size", None), + ebs_volume_throughput=d.get("ebs_volume_throughput", None), + ebs_volume_type=_enum(d, "ebs_volume_type", EbsVolumeType), + first_on_demand=d.get("first_on_demand", None), + instance_profile_arn=d.get("instance_profile_arn", None), + spot_bid_price_percent=d.get("spot_bid_price_percent", None), + zone_id=d.get("zone_id", None), + ) + + +class AwsAvailability(Enum): + """Availability type used for all subsequent nodes past the `first_on_demand` ones. + + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ON_DEMAND = "ON_DEMAND" + SPOT = "SPOT" + SPOT_WITH_FALLBACK = "SPOT_WITH_FALLBACK" + + +@dataclass +class AzureAttributes: + """Attributes set during cluster creation which are related to Microsoft Azure.""" + + availability: Optional[AzureAvailability] = None + """Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If + `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + first_on_demand: Optional[int] = None + """The first `first_on_demand` nodes of the cluster will be placed on on-demand instances. This + value should be greater than 0, to make sure the cluster driver node is placed on an on-demand + instance. If this value is greater than or equal to the current cluster size, all nodes will be + placed on on-demand instances. If this value is less than the current cluster size, + `first_on_demand` nodes will be placed on on-demand instances and the remainder will be placed + on `availability` instances. Note that this value does not affect cluster size and cannot + currently be mutated over the lifetime of a cluster.""" + + log_analytics_info: Optional[LogAnalyticsInfo] = None + """Defines values necessary to configure and run Azure Log Analytics agent""" + + spot_bid_max_price: Optional[float] = None + """The max bid price to be used for Azure spot instances. The Max price for the bid cannot be + higher than the on-demand price of the instance. If not specified, the default value is -1, + which specifies that the instance cannot be evicted on the basis of price, and only on the basis + of availability. Further, the value should > 0 or -1.""" + + def as_dict(self) -> dict: + """Serializes the AzureAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.log_analytics_info: + body["log_analytics_info"] = self.log_analytics_info.as_dict() + if self.spot_bid_max_price is not None: + body["spot_bid_max_price"] = self.spot_bid_max_price + return body + + def as_shallow_dict(self) -> dict: + """Serializes the AzureAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.first_on_demand is not None: + body["first_on_demand"] = self.first_on_demand + if self.log_analytics_info: + body["log_analytics_info"] = self.log_analytics_info + if self.spot_bid_max_price is not None: + body["spot_bid_max_price"] = self.spot_bid_max_price + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> AzureAttributes: + """Deserializes the AzureAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", AzureAvailability), + first_on_demand=d.get("first_on_demand", None), + log_analytics_info=_from_dict(d, "log_analytics_info", LogAnalyticsInfo), + spot_bid_max_price=d.get("spot_bid_max_price", None), + ) + + +class AzureAvailability(Enum): + """Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If + `first_on_demand` is zero, this availability type will be used for the entire cluster.""" + + ON_DEMAND_AZURE = "ON_DEMAND_AZURE" + SPOT_AZURE = "SPOT_AZURE" + SPOT_WITH_FALLBACK_AZURE = "SPOT_WITH_FALLBACK_AZURE" + + +@dataclass +class ClusterLogConf: + """Cluster log delivery config""" + + dbfs: Optional[DbfsStorageInfo] = None + """destination needs to be provided. e.g. `{ "dbfs" : { "destination" : "dbfs:/home/cluster_log" } + }`""" + + s3: Optional[S3StorageInfo] = None + """destination and either the region or endpoint need to be provided. e.g. `{ "s3": { "destination" + : "s3://cluster_log_bucket/prefix", "region" : "us-west-2" } }` Cluster iam role is used to + access s3, please make sure the cluster iam role in `instance_profile_arn` has permission to + write data to the s3 destination.""" + + volumes: Optional[VolumesStorageInfo] = None + """destination needs to be provided, e.g. `{ "volumes": { "destination": + "/Volumes/catalog/schema/volume/cluster_log" } }`""" + + def as_dict(self) -> dict: + """Serializes the ClusterLogConf into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dbfs: + body["dbfs"] = self.dbfs.as_dict() + if self.s3: + body["s3"] = self.s3.as_dict() + if self.volumes: + body["volumes"] = self.volumes.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ClusterLogConf into a shallow dictionary of its immediate attributes.""" + body = {} + if self.dbfs: + body["dbfs"] = self.dbfs + if self.s3: + body["s3"] = self.s3 + if self.volumes: + body["volumes"] = self.volumes + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ClusterLogConf: + """Deserializes the ClusterLogConf from a dictionary.""" + return cls( + dbfs=_from_dict(d, "dbfs", DbfsStorageInfo), + s3=_from_dict(d, "s3", S3StorageInfo), + volumes=_from_dict(d, "volumes", VolumesStorageInfo), + ) @dataclass @@ -58,6 +359,9 @@ class CreatePipeline: edition: Optional[str] = None """Pipeline product edition.""" + event_log: Optional[EventLogSpec] = None + """Event log configuration for this pipeline""" + filters: Optional[Filters] = None """Filters on which Pipeline packages to include in the deployed graph.""" @@ -136,6 +440,8 @@ def as_dict(self) -> dict: body["dry_run"] = self.dry_run if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log.as_dict() if self.filters: body["filters"] = self.filters.as_dict() if self.gateway_definition: @@ -193,6 +499,8 @@ def as_shallow_dict(self) -> dict: body["dry_run"] = self.dry_run if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log if self.filters: body["filters"] = self.filters if self.gateway_definition: @@ -240,6 +548,7 @@ def from_dict(cls, d: Dict[str, Any]) -> CreatePipeline: development=d.get("development", None), dry_run=d.get("dry_run", None), edition=d.get("edition", None), + event_log=_from_dict(d, "event_log", EventLogSpec), filters=_from_dict(d, "filters", Filters), gateway_definition=_from_dict(d, "gateway_definition", IngestionGatewayPipelineDefinition), id=d.get("id", None), @@ -367,6 +676,33 @@ class DayOfWeek(Enum): WEDNESDAY = "WEDNESDAY" +@dataclass +class DbfsStorageInfo: + """A storage location in DBFS""" + + destination: str + """dbfs destination, e.g. `dbfs:/my/path`""" + + def as_dict(self) -> dict: + """Serializes the DbfsStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the DbfsStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> DbfsStorageInfo: + """Deserializes the DbfsStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + @dataclass class DeletePipelineResponse: def as_dict(self) -> dict: @@ -392,6 +728,14 @@ class DeploymentKind(Enum): BUNDLE = "BUNDLE" +class EbsVolumeType(Enum): + """All EBS volume types that Databricks supports. See https://aws.amazon.com/ebs/details/ for + details.""" + + GENERAL_PURPOSE_SSD = "GENERAL_PURPOSE_SSD" + THROUGHPUT_OPTIMIZED_HDD = "THROUGHPUT_OPTIMIZED_HDD" + + @dataclass class EditPipeline: allow_duplicate_names: Optional[bool] = None @@ -427,6 +771,9 @@ class EditPipeline: edition: Optional[str] = None """Pipeline product edition.""" + event_log: Optional[EventLogSpec] = None + """Event log configuration for this pipeline""" + expected_last_modified: Optional[int] = None """If present, the last-modified time of the pipeline settings before the edit. If the settings were modified after that time, then the request will fail with a conflict.""" @@ -510,6 +857,8 @@ def as_dict(self) -> dict: body["development"] = self.development if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log.as_dict() if self.expected_last_modified is not None: body["expected_last_modified"] = self.expected_last_modified if self.filters: @@ -569,6 +918,8 @@ def as_shallow_dict(self) -> dict: body["development"] = self.development if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log if self.expected_last_modified is not None: body["expected_last_modified"] = self.expected_last_modified if self.filters: @@ -619,6 +970,7 @@ def from_dict(cls, d: Dict[str, Any]) -> EditPipeline: deployment=_from_dict(d, "deployment", PipelineDeployment), development=d.get("development", None), edition=d.get("edition", None), + event_log=_from_dict(d, "event_log", EventLogSpec), expected_last_modified=d.get("expected_last_modified", None), filters=_from_dict(d, "filters", Filters), gateway_definition=_from_dict(d, "gateway_definition", IngestionGatewayPipelineDefinition), @@ -698,6 +1050,47 @@ class EventLevel(Enum): WARN = "WARN" +@dataclass +class EventLogSpec: + """Configurable event log parameters.""" + + catalog: Optional[str] = None + """The UC catalog the event log is published under.""" + + name: Optional[str] = None + """The name the event log is published to in UC.""" + + schema: Optional[str] = None + """The UC schema the event log is published under.""" + + def as_dict(self) -> dict: + """Serializes the EventLogSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EventLogSpec into a shallow dictionary of its immediate attributes.""" + body = {} + if self.catalog is not None: + body["catalog"] = self.catalog + if self.name is not None: + body["name"] = self.name + if self.schema is not None: + body["schema"] = self.schema + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EventLogSpec: + """Deserializes the EventLogSpec from a dictionary.""" + return cls(catalog=d.get("catalog", None), name=d.get("name", None), schema=d.get("schema", None)) + + @dataclass class FileLibrary: path: Optional[str] = None @@ -755,6 +1148,124 @@ def from_dict(cls, d: Dict[str, Any]) -> Filters: return cls(exclude=d.get("exclude", None), include=d.get("include", None)) +@dataclass +class GcpAttributes: + """Attributes set during cluster creation which are related to GCP.""" + + availability: Optional[GcpAvailability] = None + """This field determines whether the spark executors will be scheduled to run on preemptible VMs, + on-demand VMs, or preemptible VMs with a fallback to on-demand VMs if the former is unavailable.""" + + boot_disk_size: Optional[int] = None + """Boot disk size in GB""" + + google_service_account: Optional[str] = None + """If provided, the cluster will impersonate the google service account when accessing gcloud + services (like GCS). The google service account must have previously been added to the + Databricks environment by an account administrator.""" + + local_ssd_count: Optional[int] = None + """If provided, each node (workers and driver) in the cluster will have this number of local SSDs + attached. Each local SSD is 375GB in size. Refer to [GCP documentation] for the supported number + of local SSDs for each instance type. + + [GCP documentation]: https://cloud.google.com/compute/docs/disks/local-ssd#choose_number_local_ssds""" + + use_preemptible_executors: Optional[bool] = None + """This field determines whether the spark executors will be scheduled to run on preemptible VMs + (when set to true) versus standard compute engine VMs (when set to false; default). Note: Soon + to be deprecated, use the 'availability' field instead.""" + + zone_id: Optional[str] = None + """Identifier for the availability zone in which the cluster resides. This can be one of the + following: - "HA" => High availability, spread nodes across availability zones for a Databricks + deployment region [default]. - "AUTO" => Databricks picks an availability zone to schedule the + cluster on. - A GCP availability zone => Pick One of the available zones for (machine type + + region) from https://cloud.google.com/compute/docs/regions-zones.""" + + def as_dict(self) -> dict: + """Serializes the GcpAttributes into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability.value + if self.boot_disk_size is not None: + body["boot_disk_size"] = self.boot_disk_size + if self.google_service_account is not None: + body["google_service_account"] = self.google_service_account + if self.local_ssd_count is not None: + body["local_ssd_count"] = self.local_ssd_count + if self.use_preemptible_executors is not None: + body["use_preemptible_executors"] = self.use_preemptible_executors + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GcpAttributes into a shallow dictionary of its immediate attributes.""" + body = {} + if self.availability is not None: + body["availability"] = self.availability + if self.boot_disk_size is not None: + body["boot_disk_size"] = self.boot_disk_size + if self.google_service_account is not None: + body["google_service_account"] = self.google_service_account + if self.local_ssd_count is not None: + body["local_ssd_count"] = self.local_ssd_count + if self.use_preemptible_executors is not None: + body["use_preemptible_executors"] = self.use_preemptible_executors + if self.zone_id is not None: + body["zone_id"] = self.zone_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GcpAttributes: + """Deserializes the GcpAttributes from a dictionary.""" + return cls( + availability=_enum(d, "availability", GcpAvailability), + boot_disk_size=d.get("boot_disk_size", None), + google_service_account=d.get("google_service_account", None), + local_ssd_count=d.get("local_ssd_count", None), + use_preemptible_executors=d.get("use_preemptible_executors", None), + zone_id=d.get("zone_id", None), + ) + + +class GcpAvailability(Enum): + """This field determines whether the instance pool will contain preemptible VMs, on-demand VMs, or + preemptible VMs with a fallback to on-demand VMs if the former is unavailable.""" + + ON_DEMAND_GCP = "ON_DEMAND_GCP" + PREEMPTIBLE_GCP = "PREEMPTIBLE_GCP" + PREEMPTIBLE_WITH_FALLBACK_GCP = "PREEMPTIBLE_WITH_FALLBACK_GCP" + + +@dataclass +class GcsStorageInfo: + """A storage location in Google Cloud Platform's GCS""" + + destination: str + """GCS destination/URI, e.g. `gs://my-bucket/some-prefix`""" + + def as_dict(self) -> dict: + """Serializes the GcsStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the GcsStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> GcsStorageInfo: + """Deserializes the GcsStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + @dataclass class GetPipelinePermissionLevelsResponse: permission_levels: Optional[List[PipelinePermissionsDescription]] = None @@ -1087,6 +1598,90 @@ def from_dict(cls, d: Dict[str, Any]) -> IngestionPipelineDefinition: ) +@dataclass +class InitScriptInfo: + """Config for an individual init script Next ID: 11""" + + abfss: Optional[Adlsgen2Info] = None + """destination needs to be provided, e.g. + `abfss://@.dfs.core.windows.net/`""" + + dbfs: Optional[DbfsStorageInfo] = None + """destination needs to be provided. e.g. `{ "dbfs": { "destination" : "dbfs:/home/cluster_log" } + }`""" + + file: Optional[LocalFileInfo] = None + """destination needs to be provided, e.g. `{ "file": { "destination": "file:/my/local/file.sh" } }`""" + + gcs: Optional[GcsStorageInfo] = None + """destination needs to be provided, e.g. `{ "gcs": { "destination": "gs://my-bucket/file.sh" } }`""" + + s3: Optional[S3StorageInfo] = None + """destination and either the region or endpoint need to be provided. e.g. `{ \"s3\": { + \"destination\": \"s3://cluster_log_bucket/prefix\", \"region\": \"us-west-2\" } }` Cluster iam + role is used to access s3, please make sure the cluster iam role in `instance_profile_arn` has + permission to write data to the s3 destination.""" + + volumes: Optional[VolumesStorageInfo] = None + """destination needs to be provided. e.g. `{ \"volumes\" : { \"destination\" : + \"/Volumes/my-init.sh\" } }`""" + + workspace: Optional[WorkspaceStorageInfo] = None + """destination needs to be provided, e.g. `{ "workspace": { "destination": + "/cluster-init-scripts/setup-datadog.sh" } }`""" + + def as_dict(self) -> dict: + """Serializes the InitScriptInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.abfss: + body["abfss"] = self.abfss.as_dict() + if self.dbfs: + body["dbfs"] = self.dbfs.as_dict() + if self.file: + body["file"] = self.file.as_dict() + if self.gcs: + body["gcs"] = self.gcs.as_dict() + if self.s3: + body["s3"] = self.s3.as_dict() + if self.volumes: + body["volumes"] = self.volumes.as_dict() + if self.workspace: + body["workspace"] = self.workspace.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the InitScriptInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.abfss: + body["abfss"] = self.abfss + if self.dbfs: + body["dbfs"] = self.dbfs + if self.file: + body["file"] = self.file + if self.gcs: + body["gcs"] = self.gcs + if self.s3: + body["s3"] = self.s3 + if self.volumes: + body["volumes"] = self.volumes + if self.workspace: + body["workspace"] = self.workspace + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> InitScriptInfo: + """Deserializes the InitScriptInfo from a dictionary.""" + return cls( + abfss=_from_dict(d, "abfss", Adlsgen2Info), + dbfs=_from_dict(d, "dbfs", DbfsStorageInfo), + file=_from_dict(d, "file", LocalFileInfo), + gcs=_from_dict(d, "gcs", GcsStorageInfo), + s3=_from_dict(d, "s3", S3StorageInfo), + volumes=_from_dict(d, "volumes", VolumesStorageInfo), + workspace=_from_dict(d, "workspace", WorkspaceStorageInfo), + ) + + @dataclass class ListPipelineEventsResponse: events: Optional[List[PipelineEvent]] = None @@ -1207,6 +1802,64 @@ def from_dict(cls, d: Dict[str, Any]) -> ListUpdatesResponse: ) +@dataclass +class LocalFileInfo: + destination: str + """local file destination, e.g. `file:/my/local/file.sh`""" + + def as_dict(self) -> dict: + """Serializes the LocalFileInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the LocalFileInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> LocalFileInfo: + """Deserializes the LocalFileInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +@dataclass +class LogAnalyticsInfo: + log_analytics_primary_key: Optional[str] = None + + log_analytics_workspace_id: Optional[str] = None + + def as_dict(self) -> dict: + """Serializes the LogAnalyticsInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.log_analytics_primary_key is not None: + body["log_analytics_primary_key"] = self.log_analytics_primary_key + if self.log_analytics_workspace_id is not None: + body["log_analytics_workspace_id"] = self.log_analytics_workspace_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the LogAnalyticsInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.log_analytics_primary_key is not None: + body["log_analytics_primary_key"] = self.log_analytics_primary_key + if self.log_analytics_workspace_id is not None: + body["log_analytics_workspace_id"] = self.log_analytics_workspace_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> LogAnalyticsInfo: + """Deserializes the LogAnalyticsInfo from a dictionary.""" + return cls( + log_analytics_primary_key=d.get("log_analytics_primary_key", None), + log_analytics_workspace_id=d.get("log_analytics_workspace_id", None), + ) + + @dataclass class ManualTrigger: def as_dict(self) -> dict: @@ -1233,6 +1886,51 @@ class MaturityLevel(Enum): STABLE = "STABLE" +@dataclass +class MavenLibrary: + coordinates: str + """Gradle-style maven coordinates. For example: "org.jsoup:jsoup:1.7.2".""" + + exclusions: Optional[List[str]] = None + """List of dependences to exclude. For example: `["slf4j:slf4j", "*:hadoop-client"]`. + + Maven dependency exclusions: + https://maven.apache.org/guides/introduction/introduction-to-optional-and-excludes-dependencies.html.""" + + repo: Optional[str] = None + """Maven repo to install the Maven package from. If omitted, both Maven Central Repository and + Spark Packages are searched.""" + + def as_dict(self) -> dict: + """Serializes the MavenLibrary into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.coordinates is not None: + body["coordinates"] = self.coordinates + if self.exclusions: + body["exclusions"] = [v for v in self.exclusions] + if self.repo is not None: + body["repo"] = self.repo + return body + + def as_shallow_dict(self) -> dict: + """Serializes the MavenLibrary into a shallow dictionary of its immediate attributes.""" + body = {} + if self.coordinates is not None: + body["coordinates"] = self.coordinates + if self.exclusions: + body["exclusions"] = self.exclusions + if self.repo is not None: + body["repo"] = self.repo + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> MavenLibrary: + """Deserializes the MavenLibrary from a dictionary.""" + return cls( + coordinates=d.get("coordinates", None), exclusions=d.get("exclusions", None), repo=d.get("repo", None) + ) + + @dataclass class NotebookLibrary: path: Optional[str] = None @@ -1570,15 +2268,15 @@ class PipelineCluster: """Parameters needed in order to automatically scale clusters up and down based on load. Note: autoscaling works best with DB runtime versions 3.0 or later.""" - aws_attributes: Optional[compute.AwsAttributes] = None + aws_attributes: Optional[AwsAttributes] = None """Attributes related to clusters running on Amazon Web Services. If not specified at cluster creation, a set of default values will be used.""" - azure_attributes: Optional[compute.AzureAttributes] = None + azure_attributes: Optional[AzureAttributes] = None """Attributes related to clusters running on Microsoft Azure. If not specified at cluster creation, a set of default values will be used.""" - cluster_log_conf: Optional[compute.ClusterLogConf] = None + cluster_log_conf: Optional[ClusterLogConf] = None """The configuration for delivering spark logs to a long-term storage destination. Only dbfs destinations are supported. Only one destination can be specified for one cluster. If the conf is given, the logs will be delivered to the destination every `5 mins`. The destination of @@ -1605,11 +2303,11 @@ class PipelineCluster: enable_local_disk_encryption: Optional[bool] = None """Whether to enable local disk encryption for the cluster.""" - gcp_attributes: Optional[compute.GcpAttributes] = None + gcp_attributes: Optional[GcpAttributes] = None """Attributes related to clusters running on Google Cloud Platform. If not specified at cluster creation, a set of default values will be used.""" - init_scripts: Optional[List[compute.InitScriptInfo]] = None + init_scripts: Optional[List[InitScriptInfo]] = None """The configuration for storing init scripts. Any number of destinations can be specified. The scripts are executed sequentially in the order provided. If `cluster_log_conf` is specified, init script logs are sent to `//init_scripts`.""" @@ -1755,15 +2453,15 @@ def from_dict(cls, d: Dict[str, Any]) -> PipelineCluster: return cls( apply_policy_default_values=d.get("apply_policy_default_values", None), autoscale=_from_dict(d, "autoscale", PipelineClusterAutoscale), - aws_attributes=_from_dict(d, "aws_attributes", compute.AwsAttributes), - azure_attributes=_from_dict(d, "azure_attributes", compute.AzureAttributes), - cluster_log_conf=_from_dict(d, "cluster_log_conf", compute.ClusterLogConf), + aws_attributes=_from_dict(d, "aws_attributes", AwsAttributes), + azure_attributes=_from_dict(d, "azure_attributes", AzureAttributes), + cluster_log_conf=_from_dict(d, "cluster_log_conf", ClusterLogConf), custom_tags=d.get("custom_tags", None), driver_instance_pool_id=d.get("driver_instance_pool_id", None), driver_node_type_id=d.get("driver_node_type_id", None), enable_local_disk_encryption=d.get("enable_local_disk_encryption", None), - gcp_attributes=_from_dict(d, "gcp_attributes", compute.GcpAttributes), - init_scripts=_repeated_dict(d, "init_scripts", compute.InitScriptInfo), + gcp_attributes=_from_dict(d, "gcp_attributes", GcpAttributes), + init_scripts=_repeated_dict(d, "init_scripts", InitScriptInfo), instance_pool_id=d.get("instance_pool_id", None), label=d.get("label", None), node_type_id=d.get("node_type_id", None), @@ -1964,7 +2662,7 @@ class PipelineLibrary: jar: Optional[str] = None """URI of the jar to be installed. Currently only DBFS is supported.""" - maven: Optional[compute.MavenLibrary] = None + maven: Optional[MavenLibrary] = None """Specification of a maven library to be installed.""" notebook: Optional[NotebookLibrary] = None @@ -2009,7 +2707,7 @@ def from_dict(cls, d: Dict[str, Any]) -> PipelineLibrary: return cls( file=_from_dict(d, "file", FileLibrary), jar=d.get("jar", None), - maven=_from_dict(d, "maven", compute.MavenLibrary), + maven=_from_dict(d, "maven", MavenLibrary), notebook=_from_dict(d, "notebook", NotebookLibrary), whl=d.get("whl", None), ) @@ -2205,6 +2903,9 @@ class PipelineSpec: edition: Optional[str] = None """Pipeline product edition.""" + event_log: Optional[EventLogSpec] = None + """Event log configuration for this pipeline""" + filters: Optional[Filters] = None """Filters on which Pipeline packages to include in the deployed graph.""" @@ -2271,6 +2972,8 @@ def as_dict(self) -> dict: body["development"] = self.development if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log.as_dict() if self.filters: body["filters"] = self.filters.as_dict() if self.gateway_definition: @@ -2322,6 +3025,8 @@ def as_shallow_dict(self) -> dict: body["development"] = self.development if self.edition is not None: body["edition"] = self.edition + if self.event_log: + body["event_log"] = self.event_log if self.filters: body["filters"] = self.filters if self.gateway_definition: @@ -2365,6 +3070,7 @@ def from_dict(cls, d: Dict[str, Any]) -> PipelineSpec: deployment=_from_dict(d, "deployment", PipelineDeployment), development=d.get("development", None), edition=d.get("edition", None), + event_log=_from_dict(d, "event_log", EventLogSpec), filters=_from_dict(d, "filters", Filters), gateway_definition=_from_dict(d, "gateway_definition", IngestionGatewayPipelineDefinition), id=d.get("id", None), @@ -2664,6 +3370,95 @@ def from_dict(cls, d: Dict[str, Any]) -> RunAs: return cls(service_principal_name=d.get("service_principal_name", None), user_name=d.get("user_name", None)) +@dataclass +class S3StorageInfo: + """A storage location in Amazon S3""" + + destination: str + """S3 destination, e.g. `s3://my-bucket/some-prefix` Note that logs will be delivered using cluster + iam role, please make sure you set cluster iam role and the role has write access to the + destination. Please also note that you cannot use AWS keys to deliver logs.""" + + canned_acl: Optional[str] = None + """(Optional) Set canned access control list for the logs, e.g. `bucket-owner-full-control`. If + `canned_cal` is set, please make sure the cluster iam role has `s3:PutObjectAcl` permission on + the destination bucket and prefix. The full list of possible canned acl can be found at + http://docs.aws.amazon.com/AmazonS3/latest/dev/acl-overview.html#canned-acl. Please also note + that by default only the object owner gets full controls. If you are using cross account role + for writing data, you may want to set `bucket-owner-full-control` to make bucket owner able to + read the logs.""" + + enable_encryption: Optional[bool] = None + """(Optional) Flag to enable server side encryption, `false` by default.""" + + encryption_type: Optional[str] = None + """(Optional) The encryption type, it could be `sse-s3` or `sse-kms`. It will be used only when + encryption is enabled and the default type is `sse-s3`.""" + + endpoint: Optional[str] = None + """S3 endpoint, e.g. `https://s3-us-west-2.amazonaws.com`. Either region or endpoint needs to be + set. If both are set, endpoint will be used.""" + + kms_key: Optional[str] = None + """(Optional) Kms key which will be used if encryption is enabled and encryption type is set to + `sse-kms`.""" + + region: Optional[str] = None + """S3 region, e.g. `us-west-2`. Either region or endpoint needs to be set. If both are set, + endpoint will be used.""" + + def as_dict(self) -> dict: + """Serializes the S3StorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.canned_acl is not None: + body["canned_acl"] = self.canned_acl + if self.destination is not None: + body["destination"] = self.destination + if self.enable_encryption is not None: + body["enable_encryption"] = self.enable_encryption + if self.encryption_type is not None: + body["encryption_type"] = self.encryption_type + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.kms_key is not None: + body["kms_key"] = self.kms_key + if self.region is not None: + body["region"] = self.region + return body + + def as_shallow_dict(self) -> dict: + """Serializes the S3StorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.canned_acl is not None: + body["canned_acl"] = self.canned_acl + if self.destination is not None: + body["destination"] = self.destination + if self.enable_encryption is not None: + body["enable_encryption"] = self.enable_encryption + if self.encryption_type is not None: + body["encryption_type"] = self.encryption_type + if self.endpoint is not None: + body["endpoint"] = self.endpoint + if self.kms_key is not None: + body["kms_key"] = self.kms_key + if self.region is not None: + body["region"] = self.region + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> S3StorageInfo: + """Deserializes the S3StorageInfo from a dictionary.""" + return cls( + canned_acl=d.get("canned_acl", None), + destination=d.get("destination", None), + enable_encryption=d.get("enable_encryption", None), + encryption_type=d.get("encryption_type", None), + endpoint=d.get("endpoint", None), + kms_key=d.get("kms_key", None), + region=d.get("region", None), + ) + + @dataclass class SchemaSpec: destination_catalog: Optional[str] = None @@ -3310,6 +4105,61 @@ class UpdateStateInfoState(Enum): WAITING_FOR_RESOURCES = "WAITING_FOR_RESOURCES" +@dataclass +class VolumesStorageInfo: + """A storage location back by UC Volumes.""" + + destination: str + """UC Volumes destination, e.g. `/Volumes/catalog/schema/vol1/init-scripts/setup-datadog.sh` or + `dbfs:/Volumes/catalog/schema/vol1/init-scripts/setup-datadog.sh`""" + + def as_dict(self) -> dict: + """Serializes the VolumesStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the VolumesStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> VolumesStorageInfo: + """Deserializes the VolumesStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + +@dataclass +class WorkspaceStorageInfo: + """A storage location in Workspace Filesystem (WSFS)""" + + destination: str + """wsfs destination, e.g. `workspace:/cluster-init-scripts/setup-datadog.sh`""" + + def as_dict(self) -> dict: + """Serializes the WorkspaceStorageInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + def as_shallow_dict(self) -> dict: + """Serializes the WorkspaceStorageInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.destination is not None: + body["destination"] = self.destination + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> WorkspaceStorageInfo: + """Deserializes the WorkspaceStorageInfo from a dictionary.""" + return cls(destination=d.get("destination", None)) + + class PipelinesAPI: """The Delta Live Tables API allows you to create, edit, delete, start, and view details about pipelines. @@ -3325,70 +4175,6 @@ class PipelinesAPI: def __init__(self, api_client): self._api = api_client - def wait_get_pipeline_running( - self, - pipeline_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[GetPipelineResponse], None]] = None, - ) -> GetPipelineResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (PipelineState.RUNNING,) - failure_states = (PipelineState.FAILED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(pipeline_id=pipeline_id) - status = poll.state - status_message = poll.cause - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach RUNNING, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"pipeline_id={pipeline_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_get_pipeline_idle( - self, - pipeline_id: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[GetPipelineResponse], None]] = None, - ) -> GetPipelineResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (PipelineState.IDLE,) - failure_states = (PipelineState.FAILED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(pipeline_id=pipeline_id) - status = poll.state - status_message = poll.cause - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach IDLE, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"pipeline_id={pipeline_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def create( self, *, @@ -3403,6 +4189,7 @@ def create( development: Optional[bool] = None, dry_run: Optional[bool] = None, edition: Optional[str] = None, + event_log: Optional[EventLogSpec] = None, filters: Optional[Filters] = None, gateway_definition: Optional[IngestionGatewayPipelineDefinition] = None, id: Optional[str] = None, @@ -3447,6 +4234,8 @@ def create( :param dry_run: bool (optional) :param edition: str (optional) Pipeline product edition. + :param event_log: :class:`EventLogSpec` (optional) + Event log configuration for this pipeline :param filters: :class:`Filters` (optional) Filters on which Pipeline packages to include in the deployed graph. :param gateway_definition: :class:`IngestionGatewayPipelineDefinition` (optional) @@ -3510,6 +4299,8 @@ def create( body["dry_run"] = dry_run if edition is not None: body["edition"] = edition + if event_log is not None: + body["event_log"] = event_log.as_dict() if filters is not None: body["filters"] = filters.as_dict() if gateway_definition is not None: @@ -3864,7 +4655,7 @@ def start_update( res = self._api.do("POST", f"/api/2.0/pipelines/{pipeline_id}/updates", body=body, headers=headers) return StartUpdateResponse.from_dict(res) - def stop(self, pipeline_id: str) -> Wait[GetPipelineResponse]: + def stop(self, pipeline_id: str): """Stop a pipeline. Stops the pipeline by canceling the active update. If there is no active update for the pipeline, this @@ -3874,20 +4665,14 @@ def stop(self, pipeline_id: str) -> Wait[GetPipelineResponse]: :returns: Long-running operation waiter for :class:`GetPipelineResponse`. - See :method:wait_get_pipeline_idle for more details. + See :method:WaitGetPipelineIdle for more details. """ headers = { "Accept": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/pipelines/{pipeline_id}/stop", headers=headers) - return Wait( - self.wait_get_pipeline_idle, response=StopPipelineResponse.from_dict(op_response), pipeline_id=pipeline_id - ) - - def stop_and_wait(self, pipeline_id: str, timeout=timedelta(minutes=20)) -> GetPipelineResponse: - return self.stop(pipeline_id=pipeline_id).result(timeout=timeout) + self._api.do("POST", f"/api/2.0/pipelines/{pipeline_id}/stop", headers=headers) def update( self, @@ -3903,6 +4688,7 @@ def update( deployment: Optional[PipelineDeployment] = None, development: Optional[bool] = None, edition: Optional[str] = None, + event_log: Optional[EventLogSpec] = None, expected_last_modified: Optional[int] = None, filters: Optional[Filters] = None, gateway_definition: Optional[IngestionGatewayPipelineDefinition] = None, @@ -3948,6 +4734,8 @@ def update( Whether the pipeline is in Development mode. Defaults to false. :param edition: str (optional) Pipeline product edition. + :param event_log: :class:`EventLogSpec` (optional) + Event log configuration for this pipeline :param expected_last_modified: int (optional) If present, the last-modified time of the pipeline settings before the edit. If the settings were modified after that time, then the request will fail with a conflict. @@ -4012,6 +4800,8 @@ def update( body["development"] = development if edition is not None: body["edition"] = edition + if event_log is not None: + body["event_log"] = event_log.as_dict() if expected_last_modified is not None: body["expected_last_modified"] = expected_last_modified if filters is not None: diff --git a/databricks/sdk/provisioning/__init__.py b/databricks/sdk/provisioning/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/provisioning/v2/__init__.py b/databricks/sdk/provisioning/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/provisioning/v2/client.py b/databricks/sdk/provisioning/v2/client.py new file mode 100755 index 000000000..dd0136540 --- /dev/null +++ b/databricks/sdk/provisioning/v2/client.py @@ -0,0 +1,501 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .provisioning import (CredentialsAPI, EncryptionKeysAPI, NetworksAPI, + PrivateAccessAPI, StorageAPI, VpcEndpointsAPI, + WorkspacesAPI) + +_LOG = logging.getLogger(__name__) + + +class CredentialsClient(CredentialsAPI): + """ + These APIs manage credential configurations for this workspace. Databricks needs access to a + cross-account service IAM role in your AWS account so that Databricks can deploy clusters in the + appropriate VPC for the new workspace. A credential configuration encapsulates this role + information, and its ID is used when creating a new workspace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class EncryptionKeysClient(EncryptionKeysAPI): + """ + These APIs manage encryption key configurations for this workspace (optional). A key + configuration encapsulates the AWS KMS key information and some information about how the key + configuration can be used. There are two possible uses for key configurations: + + * Managed services: A key configuration can be used to encrypt a workspace's notebook and secret + data in the control plane, as well as Databricks SQL queries and query history. * Storage: A key + configuration can be used to encrypt a workspace's DBFS and EBS data in the data plane. + + In both of these cases, the key configuration's ID is used when creating a new workspace. This + Preview feature is available if your account is on the E2 version of the platform. Updating a + running workspace with workspace storage encryption requires that the workspace is on the E2 + version of the platform. If you have an older workspace, it might not be on the E2 version of + the platform. If you are not sure, contact your Databricks representative. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class NetworksClient(NetworksAPI): + """ + These APIs manage network configurations for customer-managed VPCs (optional). Its ID is used + when creating a new workspace if you use customer-managed VPCs. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PrivateAccessClient(PrivateAccessAPI): + """ + These APIs manage private access settings for this account. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class StorageClient(StorageAPI): + """ + These APIs manage storage configurations for this workspace. A root storage S3 bucket in your + account is required to store objects like cluster logs, notebook revisions, and job results. You + can also use the root storage S3 bucket for storage of non-production DBFS data. A storage + configuration encapsulates this bucket information, and its ID is used when creating a new + workspace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class VpcEndpointsClient(VpcEndpointsAPI): + """ + These APIs manage VPC endpoint configurations for this account. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WorkspacesClient(WorkspacesAPI): + """ + These APIs manage workspaces for this account. A Databricks workspace is an environment for + accessing all of your Databricks assets. The workspace organizes objects (notebooks, libraries, + and experiments) into folders, and provides access to data and computational resources such as + clusters and jobs. + + These endpoints are available if your account is on the E2 version of the platform or on a + select custom plan that allows multiple workspaces per account. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/provisioning.py b/databricks/sdk/provisioning/v2/provisioning.py similarity index 96% rename from databricks/sdk/service/provisioning.py rename to databricks/sdk/provisioning/v2/provisioning.py index d98c44b17..263203849 100755 --- a/databricks/sdk/service/provisioning.py +++ b/databricks/sdk/provisioning/v2/provisioning.py @@ -3,19 +3,15 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -3166,38 +3162,6 @@ class WorkspacesAPI: def __init__(self, api_client): self._api = api_client - def wait_get_workspace_running( - self, workspace_id: int, timeout=timedelta(minutes=20), callback: Optional[Callable[[Workspace], None]] = None - ) -> Workspace: - deadline = time.time() + timeout.total_seconds() - target_states = (WorkspaceStatus.RUNNING,) - failure_states = ( - WorkspaceStatus.BANNED, - WorkspaceStatus.FAILED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(workspace_id=workspace_id) - status = poll.workspace_status - status_message = poll.workspace_status_message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach RUNNING, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"workspace_id={workspace_id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def create( self, workspace_name: str, @@ -3218,7 +3182,7 @@ def create( private_access_settings_id: Optional[str] = None, storage_configuration_id: Optional[str] = None, storage_customer_managed_key_id: Optional[str] = None, - ) -> Wait[Workspace]: + ) -> Workspace: """Create a new workspace. Creates a new workspace. @@ -3321,7 +3285,7 @@ def create( :returns: Long-running operation waiter for :class:`Workspace`. - See :method:wait_get_workspace_running for more details. + See :method:WaitGetWorkspaceRunning for more details. """ body = {} if aws_region is not None: @@ -3363,56 +3327,8 @@ def create( "Content-Type": "application/json", } - op_response = self._api.do( - "POST", f"/api/2.0/accounts/{self._api.account_id}/workspaces", body=body, headers=headers - ) - return Wait( - self.wait_get_workspace_running, - response=Workspace.from_dict(op_response), - workspace_id=op_response["workspace_id"], - ) - - def create_and_wait( - self, - workspace_name: str, - *, - aws_region: Optional[str] = None, - cloud: Optional[str] = None, - cloud_resource_container: Optional[CloudResourceContainer] = None, - credentials_id: Optional[str] = None, - custom_tags: Optional[Dict[str, str]] = None, - deployment_name: Optional[str] = None, - gcp_managed_network_config: Optional[GcpManagedNetworkConfig] = None, - gke_config: Optional[GkeConfig] = None, - is_no_public_ip_enabled: Optional[bool] = None, - location: Optional[str] = None, - managed_services_customer_managed_key_id: Optional[str] = None, - network_id: Optional[str] = None, - pricing_tier: Optional[PricingTier] = None, - private_access_settings_id: Optional[str] = None, - storage_configuration_id: Optional[str] = None, - storage_customer_managed_key_id: Optional[str] = None, - timeout=timedelta(minutes=20), - ) -> Workspace: - return self.create( - aws_region=aws_region, - cloud=cloud, - cloud_resource_container=cloud_resource_container, - credentials_id=credentials_id, - custom_tags=custom_tags, - deployment_name=deployment_name, - gcp_managed_network_config=gcp_managed_network_config, - gke_config=gke_config, - is_no_public_ip_enabled=is_no_public_ip_enabled, - location=location, - managed_services_customer_managed_key_id=managed_services_customer_managed_key_id, - network_id=network_id, - pricing_tier=pricing_tier, - private_access_settings_id=private_access_settings_id, - storage_configuration_id=storage_configuration_id, - storage_customer_managed_key_id=storage_customer_managed_key_id, - workspace_name=workspace_name, - ).result(timeout=timeout) + res = self._api.do("POST", f"/api/2.0/accounts/{self._api.account_id}/workspaces", body=body, headers=headers) + return Workspace.from_dict(res) def delete(self, workspace_id: int): """Delete a workspace. @@ -3498,7 +3414,7 @@ def update( private_access_settings_id: Optional[str] = None, storage_configuration_id: Optional[str] = None, storage_customer_managed_key_id: Optional[str] = None, - ) -> Wait[Workspace]: + ): """Update workspace configuration. Updates a workspace configuration for either a running workspace or a failed workspace. The elements @@ -3627,7 +3543,7 @@ def update( :returns: Long-running operation waiter for :class:`Workspace`. - See :method:wait_get_workspace_running for more details. + See :method:WaitGetWorkspaceRunning for more details. """ body = {} if aws_region is not None: @@ -3653,37 +3569,6 @@ def update( "Content-Type": "application/json", } - op_response = self._api.do( + self._api.do( "PATCH", f"/api/2.0/accounts/{self._api.account_id}/workspaces/{workspace_id}", body=body, headers=headers ) - return Wait( - self.wait_get_workspace_running, response=UpdateResponse.from_dict(op_response), workspace_id=workspace_id - ) - - def update_and_wait( - self, - workspace_id: int, - *, - aws_region: Optional[str] = None, - credentials_id: Optional[str] = None, - custom_tags: Optional[Dict[str, str]] = None, - managed_services_customer_managed_key_id: Optional[str] = None, - network_connectivity_config_id: Optional[str] = None, - network_id: Optional[str] = None, - private_access_settings_id: Optional[str] = None, - storage_configuration_id: Optional[str] = None, - storage_customer_managed_key_id: Optional[str] = None, - timeout=timedelta(minutes=20), - ) -> Workspace: - return self.update( - aws_region=aws_region, - credentials_id=credentials_id, - custom_tags=custom_tags, - managed_services_customer_managed_key_id=managed_services_customer_managed_key_id, - network_connectivity_config_id=network_connectivity_config_id, - network_id=network_id, - private_access_settings_id=private_access_settings_id, - storage_configuration_id=storage_configuration_id, - storage_customer_managed_key_id=storage_customer_managed_key_id, - workspace_id=workspace_id, - ).result(timeout=timeout) diff --git a/databricks/sdk/service/_internal.py b/databricks/sdk/service/_internal.py index 1e501e0e0..d1c124832 100644 --- a/databricks/sdk/service/_internal.py +++ b/databricks/sdk/service/_internal.py @@ -1,15 +1,15 @@ import datetime import urllib.parse -from typing import Callable, Dict, Generic, Optional, Type, TypeVar +from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar -def _from_dict(d: Dict[str, any], field: str, cls: Type) -> any: +def _from_dict(d: Dict[str, Any], field: str, cls: Type) -> Any: if field not in d or d[field] is None: return None return getattr(cls, "from_dict")(d[field]) -def _repeated_dict(d: Dict[str, any], field: str, cls: Type) -> any: +def _repeated_dict(d: Dict[str, Any], field: str, cls: Type) -> Any: if field not in d or not d[field]: return [] from_dict = getattr(cls, "from_dict") @@ -23,14 +23,14 @@ def _get_enum_value(cls: Type, value: str) -> Optional[Type]: ) -def _enum(d: Dict[str, any], field: str, cls: Type) -> any: +def _enum(d: Dict[str, Any], field: str, cls: Type) -> Any: """Unknown enum values are returned as None.""" if field not in d or not d[field]: return None return _get_enum_value(cls, d[field]) -def _repeated_enum(d: Dict[str, any], field: str, cls: Type) -> any: +def _repeated_enum(d: Dict[str, Any], field: str, cls: Type) -> Any: """For now, unknown enum values are not included in the response.""" if field not in d or not d[field]: return None @@ -51,13 +51,13 @@ def _escape_multi_segment_path_parameter(param: str) -> str: class Wait(Generic[ReturnType]): - def __init__(self, waiter: Callable, response: any = None, **kwargs) -> None: + def __init__(self, waiter: Callable, response: Any = None, **kwargs) -> None: self.response = response self._waiter = waiter self._bind = kwargs - def __getattr__(self, key) -> any: + def __getattr__(self, key) -> Any: return self._bind[key] def bind(self) -> dict: diff --git a/databricks/sdk/serving/__init__.py b/databricks/sdk/serving/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/serving/v2/__init__.py b/databricks/sdk/serving/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/serving/v2/client.py b/databricks/sdk/serving/v2/client.py new file mode 100755 index 000000000..e8cebcb09 --- /dev/null +++ b/databricks/sdk/serving/v2/client.py @@ -0,0 +1,154 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .mixin import ServingEndpointsExt +from .serving import ServingEndpointsDataPlaneAPI + +_LOG = logging.getLogger(__name__) + + +class ServingEndpointsClient(ServingEndpointsExt): + """ + The Serving Endpoints API allows you to create, update, and delete model serving endpoints. + + You can use a serving endpoint to serve models from the Databricks Model Registry or from Unity + Catalog. Endpoints expose the underlying models as scalable REST API endpoints using serverless + compute. This means the endpoints and associated compute resources are fully managed by + Databricks and will not appear in your cloud account. A serving endpoint can consist of one or + more MLflow models from the Databricks Model Registry, called served entities. A serving + endpoint can have at most ten served entities. You can configure traffic settings to define how + requests should be routed to your served entities behind an endpoint. Additionally, you can + configure the scale of resources that should be applied to each served entity. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ServingEndpointsDataPlaneClient(ServingEndpointsDataPlaneAPI): + """ + Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints + for Serving endpoints service. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/mixins/open_ai_client.py b/databricks/sdk/serving/v2/mixin.py similarity index 96% rename from databricks/sdk/mixins/open_ai_client.py rename to databricks/sdk/serving/v2/mixin.py index 62835dcec..0fb198565 100644 --- a/databricks/sdk/mixins/open_ai_client.py +++ b/databricks/sdk/serving/v2/mixin.py @@ -3,8 +3,7 @@ from requests import Response -from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod, - ServingEndpointsAPI) +from .serving import ExternalFunctionRequestHttpMethod, ServingEndpointsAPI class ServingEndpointsExt(ServingEndpointsAPI): diff --git a/databricks/sdk/service/serving.py b/databricks/sdk/serving/v2/serving.py similarity index 97% rename from databricks/sdk/service/serving.py rename to databricks/sdk/serving/v2/serving.py index 0a66ec777..4ca074a96 100755 --- a/databricks/sdk/service/serving.py +++ b/databricks/sdk/serving/v2/serving.py @@ -3,22 +3,14 @@ from __future__ import annotations import logging -import random -import threading -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, BinaryIO, Callable, Dict, Iterator, List, Optional +from typing import Any, BinaryIO, Dict, Iterator, List, Optional -import requests - -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -849,8 +841,9 @@ class CreateServingEndpoint: workspace. An endpoint name can consist of alphanumeric characters, dashes, and underscores.""" ai_gateway: Optional[AiGatewayConfig] = None - """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned - throughput endpoints are currently supported.""" + """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned + throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only + support inference tables.""" budget_policy_id: Optional[str] = None """The budget policy to be applied to the serving endpoint.""" @@ -3294,6 +3287,7 @@ class ServedModelInputWorkloadSize(Enum): class ServedModelInputWorkloadType(Enum): + """Please keep this in sync with with workload types in InferenceEndpointEntities.scala""" CPU = "CPU" GPU_LARGE = "GPU_LARGE" @@ -3530,8 +3524,9 @@ def from_dict(cls, d: Dict[str, Any]) -> ServerLogsResponse: @dataclass class ServingEndpoint: ai_gateway: Optional[AiGatewayConfig] = None - """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned - throughput endpoints are currently supported.""" + """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned + throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only + support inference tables.""" budget_policy_id: Optional[str] = None """The budget policy associated with the endpoint.""" @@ -3748,8 +3743,9 @@ def from_dict(cls, d: Dict[str, Any]) -> ServingEndpointAccessControlResponse: @dataclass class ServingEndpointDetailed: ai_gateway: Optional[AiGatewayConfig] = None - """The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned - throughput endpoints are currently supported.""" + """The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned + throughput, and pay-per-token endpoints are fully supported; agent endpoints currently only + support inference tables.""" budget_policy_id: Optional[str] = None """The budget policy associated with the endpoint.""" @@ -4059,6 +4055,7 @@ def from_dict(cls, d: Dict[str, Any]) -> ServingEndpointPermissionsRequest: class ServingModelWorkloadType(Enum): + """Please keep this in sync with with workload types in InferenceEndpointEntities.scala""" CPU = "CPU" GPU_LARGE = "GPU_LARGE" @@ -4166,41 +4163,6 @@ class ServingEndpointsAPI: def __init__(self, api_client): self._api = api_client - def wait_get_serving_endpoint_not_updating( - self, - name: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[ServingEndpointDetailed], None]] = None, - ) -> ServingEndpointDetailed: - deadline = time.time() + timeout.total_seconds() - target_states = (EndpointStateConfigUpdate.NOT_UPDATING,) - failure_states = ( - EndpointStateConfigUpdate.UPDATE_FAILED, - EndpointStateConfigUpdate.UPDATE_CANCELED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(name=name) - status = poll.state.config_update - status_message = f"current status: {status}" - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach NOT_UPDATING, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"name={name}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def build_logs(self, name: str, served_model_name: str) -> BuildLogsResponse: """Get build logs for a served model. @@ -4233,15 +4195,16 @@ def create( rate_limits: Optional[List[RateLimit]] = None, route_optimized: Optional[bool] = None, tags: Optional[List[EndpointTag]] = None, - ) -> Wait[ServingEndpointDetailed]: + ) -> ServingEndpointDetailed: """Create a new serving endpoint. :param name: str The name of the serving endpoint. This field is required and must be unique across a Databricks workspace. An endpoint name can consist of alphanumeric characters, dashes, and underscores. :param ai_gateway: :class:`AiGatewayConfig` (optional) - The AI Gateway configuration for the serving endpoint. NOTE: Only external model and provisioned - throughput endpoints are currently supported. + The AI Gateway configuration for the serving endpoint. NOTE: External model, provisioned throughput, + and pay-per-token endpoints are fully supported; agent endpoints currently only support inference + tables. :param budget_policy_id: str (optional) The budget policy to be applied to the serving endpoint. :param config: :class:`EndpointCoreConfigInput` (optional) @@ -4256,7 +4219,7 @@ def create( :returns: Long-running operation waiter for :class:`ServingEndpointDetailed`. - See :method:wait_get_serving_endpoint_not_updating for more details. + See :method:WaitGetServingEndpointNotUpdating for more details. """ body = {} if ai_gateway is not None: @@ -4278,34 +4241,8 @@ def create( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/serving-endpoints", body=body, headers=headers) - return Wait( - self.wait_get_serving_endpoint_not_updating, - response=ServingEndpointDetailed.from_dict(op_response), - name=op_response["name"], - ) - - def create_and_wait( - self, - name: str, - *, - ai_gateway: Optional[AiGatewayConfig] = None, - budget_policy_id: Optional[str] = None, - config: Optional[EndpointCoreConfigInput] = None, - rate_limits: Optional[List[RateLimit]] = None, - route_optimized: Optional[bool] = None, - tags: Optional[List[EndpointTag]] = None, - timeout=timedelta(minutes=20), - ) -> ServingEndpointDetailed: - return self.create( - ai_gateway=ai_gateway, - budget_policy_id=budget_policy_id, - config=config, - name=name, - rate_limits=rate_limits, - route_optimized=route_optimized, - tags=tags, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/serving-endpoints", body=body, headers=headers) + return ServingEndpointDetailed.from_dict(res) def delete(self, name: str): """Delete a serving endpoint. @@ -4531,8 +4468,7 @@ def patch( def put(self, name: str, *, rate_limits: Optional[List[RateLimit]] = None) -> PutResponse: """Update rate limits of a serving endpoint. - Used to update the rate limits of a serving endpoint. NOTE: Only foundation model endpoints are - currently supported. For external models, use AI Gateway to manage rate limits. + Deprecated: Please use AI Gateway to manage rate limits instead. :param name: str The name of the serving endpoint whose rate limits are being updated. This field is required. @@ -4564,8 +4500,8 @@ def put_ai_gateway( ) -> PutAiGatewayResponse: """Update AI Gateway of a serving endpoint. - Used to update the AI Gateway of a serving endpoint. NOTE: Only external model and provisioned - throughput endpoints are currently supported. + Used to update the AI Gateway of a serving endpoint. NOTE: External model, provisioned throughput, and + pay-per-token endpoints are fully supported; agent endpoints currently only support inference tables. :param name: str The name of the serving endpoint whose AI Gateway is being updated. This field is required. @@ -4751,7 +4687,7 @@ def update_config( served_entities: Optional[List[ServedEntityInput]] = None, served_models: Optional[List[ServedModelInput]] = None, traffic_config: Optional[TrafficConfig] = None, - ) -> Wait[ServingEndpointDetailed]: + ) -> ServingEndpointDetailed: """Update config of a serving endpoint. Updates any combination of the serving endpoint's served entities, the compute configuration of those @@ -4775,7 +4711,7 @@ def update_config( :returns: Long-running operation waiter for :class:`ServingEndpointDetailed`. - See :method:wait_get_serving_endpoint_not_updating for more details. + See :method:WaitGetServingEndpointNotUpdating for more details. """ body = {} if auto_capture_config is not None: @@ -4791,30 +4727,8 @@ def update_config( "Content-Type": "application/json", } - op_response = self._api.do("PUT", f"/api/2.0/serving-endpoints/{name}/config", body=body, headers=headers) - return Wait( - self.wait_get_serving_endpoint_not_updating, - response=ServingEndpointDetailed.from_dict(op_response), - name=op_response["name"], - ) - - def update_config_and_wait( - self, - name: str, - *, - auto_capture_config: Optional[AutoCaptureConfigInput] = None, - served_entities: Optional[List[ServedEntityInput]] = None, - served_models: Optional[List[ServedModelInput]] = None, - traffic_config: Optional[TrafficConfig] = None, - timeout=timedelta(minutes=20), - ) -> ServingEndpointDetailed: - return self.update_config( - auto_capture_config=auto_capture_config, - name=name, - served_entities=served_entities, - served_models=served_models, - traffic_config=traffic_config, - ).result(timeout=timeout) + res = self._api.do("PUT", f"/api/2.0/serving-endpoints/{name}/config", body=body, headers=headers) + return ServingEndpointDetailed.from_dict(res) def update_permissions( self, @@ -4851,31 +4765,8 @@ class ServingEndpointsDataPlaneAPI: """Serving endpoints DataPlane provides a set of operations to interact with data plane endpoints for Serving endpoints service.""" - def __init__(self, api_client, control_plane_service, dpts): + def __init__(self, api_client): self._api = api_client - self._lock = threading.Lock() - self._control_plane_service = control_plane_service - self._dpts = dpts - self._data_plane_details = {} - - def _data_plane_info_query(self, name: str) -> DataPlaneInfo: - key = "query" + "/".join( - [ - str(name), - ] - ) - with self._lock: - if key in self._data_plane_details: - return self._data_plane_details[key] - response = self._control_plane_service.get( - name=name, - ) - if response.data_plane_info is None: - raise Exception("Resource does not support direct Data Plane access") - result = response.data_plane_info.query_info - with self._lock: - self._data_plane_details[key] = result - return result def query( self, @@ -4970,16 +4861,6 @@ def query( body["stream"] = stream if temperature is not None: body["temperature"] = temperature - data_plane_info = self._data_plane_info_query( - name=name, - ) - token = self._dpts.token(data_plane_info.endpoint_url, data_plane_info.authorization_details) - - def auth(r: requests.PreparedRequest) -> requests.PreparedRequest: - authorization = f"{token.token_type} {token.access_token}" - r.headers["Authorization"] = authorization - return r - headers = { "Accept": "application/json", "Content-Type": "application/json", @@ -4989,10 +4870,9 @@ def auth(r: requests.PreparedRequest) -> requests.PreparedRequest: ] res = self._api.do( "POST", - url=data_plane_info.endpoint_url, + f"/serving-endpoints/{name}/invocations", body=body, headers=headers, response_headers=response_headers, - auth=auth, ) return QueryEndpointResponse.from_dict(res) diff --git a/databricks/sdk/settings/__init__.py b/databricks/sdk/settings/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/settings/v2/__init__.py b/databricks/sdk/settings/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/settings/v2/client.py b/databricks/sdk/settings/v2/client.py new file mode 100755 index 000000000..8c1642bb8 --- /dev/null +++ b/databricks/sdk/settings/v2/client.py @@ -0,0 +1,1775 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .settings import (AccountIpAccessListsAPI, AccountSettingsAPI, + AibiDashboardEmbeddingAccessPolicyAPI, + AibiDashboardEmbeddingApprovedDomainsAPI, + AutomaticClusterUpdateAPI, ComplianceSecurityProfileAPI, + CredentialsManagerAPI, CspEnablementAccountAPI, + DefaultNamespaceAPI, DisableLegacyAccessAPI, + DisableLegacyDbfsAPI, DisableLegacyFeaturesAPI, + EnableIpAccessListsAPI, EnableResultsDownloadingAPI, + EnhancedSecurityMonitoringAPI, EsmEnablementAccountAPI, + IpAccessListsAPI, NetworkConnectivityAPI, + NotificationDestinationsAPI, PersonalComputeAPI, + RestrictWorkspaceAdminsAPI, SettingsAPI, + TokenManagementAPI, TokensAPI, WorkspaceConfAPI) + +_LOG = logging.getLogger(__name__) + + +class AccountIpAccessListsClient(AccountIpAccessListsAPI): + """ + The Accounts IP Access List API enables account admins to configure IP access lists for access + to the account console. + + Account IP Access Lists affect web application access and REST API access to the account console + and account APIs. If the feature is disabled for the account, all access is allowed for this + account. There is support for allow lists (inclusion) and block lists (exclusion). + + When a connection is attempted: 1. **First, all block lists are checked.** If the connection IP + address matches any block list, the connection is rejected. 2. **If the connection was not + rejected by block lists**, the IP address is compared with the allow lists. + + If there is at least one allow list for the account, the connection is allowed only if the IP + address matches an allow list. If there are no allow lists for the account, all IP addresses are + allowed. + + For all allow lists and block lists combined, the account supports a maximum of 1000 IP/CIDR + values, where one CIDR counts as a single value. + + After changes to the account-level IP access lists, it can take a few minutes for changes to + take effect. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AccountSettingsClient(AccountSettingsAPI): + """ + Accounts Settings API allows users to manage settings at the account level. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AibiDashboardEmbeddingAccessPolicyClient(AibiDashboardEmbeddingAccessPolicyAPI): + """ + Controls whether AI/BI published dashboard embedding is enabled, conditionally enabled, or + disabled at the workspace level. By default, this setting is conditionally enabled + (ALLOW_APPROVED_DOMAINS). + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AibiDashboardEmbeddingApprovedDomainsClient(AibiDashboardEmbeddingApprovedDomainsAPI): + """ + Controls the list of domains approved to host the embedded AI/BI dashboards. The approved + domains list can't be mutated when the current access policy is not set to + ALLOW_APPROVED_DOMAINS. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AutomaticClusterUpdateClient(AutomaticClusterUpdateAPI): + """ + Controls whether automatic cluster update is enabled for the current workspace. By default, it + is turned off. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ComplianceSecurityProfileClient(ComplianceSecurityProfileAPI): + """ + Controls whether to enable the compliance security profile for the current workspace. Enabling + it on a workspace is permanent. By default, it is turned off. + + This settings can NOT be disabled once it is enabled. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CredentialsManagerClient(CredentialsManagerAPI): + """ + Credentials manager interacts with with Identity Providers to to perform token exchanges using + stored credentials and refresh tokens. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class CspEnablementAccountClient(CspEnablementAccountAPI): + """ + The compliance security profile settings at the account level control whether to enable it for + new workspaces. By default, this account-level setting is disabled for new workspaces. After + workspace creation, account admins can enable the compliance security profile individually for + each workspace. + + This settings can be disabled so that new workspaces do not have compliance security profile + enabled by default. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DefaultNamespaceClient(DefaultNamespaceAPI): + """ + The default namespace setting API allows users to configure the default namespace for a + Databricks workspace. + + Through this API, users can retrieve, set, or modify the default namespace used when queries do + not reference a fully qualified three-level name. For example, if you use the API to set + 'retail_prod' as the default catalog, then a query 'SELECT * FROM myTable' would reference the + object 'retail_prod.default.myTable' (the schema 'default' is always assumed). + + This setting requires a restart of clusters and SQL warehouses to take effect. Additionally, the + default namespace only applies when using Unity Catalog-enabled compute. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DisableLegacyAccessClient(DisableLegacyAccessAPI): + """ + 'Disabling legacy access' has the following impacts: + + 1. Disables direct access to Hive Metastores from the workspace. However, you can still access a + Hive Metastore through Hive Metastore federation. 2. Disables fallback mode on external location + access from the workspace. 3. Disables Databricks Runtime versions prior to 13.3LTS. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DisableLegacyDbfsClient(DisableLegacyDbfsAPI): + """ + When this setting is on, access to DBFS root and DBFS mounts is disallowed (as well as creation + of new mounts). When the setting is off, all DBFS functionality is enabled + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DisableLegacyFeaturesClient(DisableLegacyFeaturesAPI): + """ + Disable legacy features for new Databricks workspaces. + + For newly created workspaces: 1. Disables the use of DBFS root and mounts. 2. Hive Metastore + will not be provisioned. 3. Disables the use of ‘No-isolation clusters’. 4. Disables + Databricks Runtime versions prior to 13.3LTS. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class EnableIpAccessListsClient(EnableIpAccessListsAPI): + """ + Controls the enforcement of IP access lists for accessing the account console. Allowing you to + enable or disable restricted access based on IP addresses. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class EnableResultsDownloadingClient(EnableResultsDownloadingAPI): + """ + The Enable Results Downloading API controls the workspace level conf for the enablement of + downloading results. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class EnhancedSecurityMonitoringClient(EnhancedSecurityMonitoringAPI): + """ + Controls whether enhanced security monitoring is enabled for the current workspace. If the + compliance security profile is enabled, this is automatically enabled. By default, it is + disabled. However, if the compliance security profile is enabled, this is automatically enabled. + + If the compliance security profile is disabled, you can enable or disable this setting and it is + not permanent. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class EsmEnablementAccountClient(EsmEnablementAccountAPI): + """ + The enhanced security monitoring setting at the account level controls whether to enable the + feature on new workspaces. By default, this account-level setting is disabled for new + workspaces. After workspace creation, account admins can enable enhanced security monitoring + individually for each workspace. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class IpAccessListsClient(IpAccessListsAPI): + """ + IP Access List enables admins to configure IP access lists. + + IP access lists affect web application access and REST API access to this workspace only. If the + feature is disabled for a workspace, all access is allowed for this workspace. There is support + for allow lists (inclusion) and block lists (exclusion). + + When a connection is attempted: 1. **First, all block lists are checked.** If the connection IP + address matches any block list, the connection is rejected. 2. **If the connection was not + rejected by block lists**, the IP address is compared with the allow lists. + + If there is at least one allow list for the workspace, the connection is allowed only if the IP + address matches an allow list. If there are no allow lists for the workspace, all IP addresses + are allowed. + + For all allow lists and block lists combined, the workspace supports a maximum of 1000 IP/CIDR + values, where one CIDR counts as a single value. + + After changes to the IP access list feature, it can take a few minutes for changes to take + effect. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class NetworkConnectivityClient(NetworkConnectivityAPI): + """ + These APIs provide configurations for the network connectivity of your workspaces for serverless + compute resources. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class NotificationDestinationsClient(NotificationDestinationsAPI): + """ + The notification destinations API lets you programmatically manage a workspace's notification + destinations. Notification destinations are used to send notifications for query alerts and jobs + to destinations outside of Databricks. Only workspace admins can create, update, and delete + notification destinations. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class PersonalComputeClient(PersonalComputeAPI): + """ + The Personal Compute enablement setting lets you control which users can use the Personal + Compute default policy to create compute resources. By default all users in all workspaces have + access (ON), but you can change the setting to instead let individual workspaces configure + access control (DELEGATE). + + There is only one instance of this setting per account. Since this setting has a default value, + this setting is present on all accounts even though it's never set on a given account. Deletion + reverts the value of the setting back to the default value. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class RestrictWorkspaceAdminsClient(RestrictWorkspaceAdminsAPI): + """ + The Restrict Workspace Admins setting lets you control the capabilities of workspace admins. + With the setting status set to ALLOW_ALL, workspace admins can create service principal personal + access tokens on behalf of any service principal in their workspace. Workspace admins can also + change a job owner to any user in their workspace. And they can change the job run_as setting to + any user in their workspace or to a service principal on which they have the Service Principal + User role. With the setting status set to RESTRICT_TOKENS_AND_JOB_RUN_AS, workspace admins can + only create personal access tokens on behalf of service principals they have the Service + Principal User role on. They can also only change a job owner to themselves. And they can change + the job run_as setting to themselves or to a service principal on which they have the Service + Principal User role. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class SettingsClient(SettingsAPI): + """ + Workspace Settings API allows users to manage settings at the workspace level. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class TokenManagementClient(TokenManagementAPI): + """ + Enables administrators to get all tokens and delete tokens for other users. Admins can either + get every token, get a specific token by ID, or get all tokens for a particular user. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class TokensClient(TokensAPI): + """ + The Token API allows you to create, list, and revoke tokens that can be used to authenticate and + access Databricks REST APIs. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WorkspaceConfClient(WorkspaceConfAPI): + """ + This API allows updating known workspace settings for advanced users. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/settings.py b/databricks/sdk/settings/v2/settings.py similarity index 94% rename from databricks/sdk/service/settings.py rename to databricks/sdk/settings/v2/settings.py index 317219e7f..67ae16b9b 100755 --- a/databricks/sdk/service/settings.py +++ b/databricks/sdk/settings/v2/settings.py @@ -7,11 +7,11 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -1662,305 +1662,6 @@ def from_dict(cls, d: Dict[str, Any]) -> DisableLegacyFeatures: ) -@dataclass -class EgressNetworkPolicy: - """The network policies applying for egress traffic. This message is used by the UI/REST API. We - translate this message to the format expected by the dataplane in Lakehouse Network Manager (for - the format expected by the dataplane, see networkconfig.textproto).""" - - internet_access: Optional[EgressNetworkPolicyInternetAccessPolicy] = None - """The access policy enforced for egress traffic to the internet.""" - - def as_dict(self) -> dict: - """Serializes the EgressNetworkPolicy into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.internet_access: - body["internet_access"] = self.internet_access.as_dict() - return body - - def as_shallow_dict(self) -> dict: - """Serializes the EgressNetworkPolicy into a shallow dictionary of its immediate attributes.""" - body = {} - if self.internet_access: - body["internet_access"] = self.internet_access - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicy: - """Deserializes the EgressNetworkPolicy from a dictionary.""" - return cls(internet_access=_from_dict(d, "internet_access", EgressNetworkPolicyInternetAccessPolicy)) - - -@dataclass -class EgressNetworkPolicyInternetAccessPolicy: - allowed_internet_destinations: Optional[List[EgressNetworkPolicyInternetAccessPolicyInternetDestination]] = None - - allowed_storage_destinations: Optional[List[EgressNetworkPolicyInternetAccessPolicyStorageDestination]] = None - - log_only_mode: Optional[EgressNetworkPolicyInternetAccessPolicyLogOnlyMode] = None - """Optional. If not specified, assume the policy is enforced for all workloads.""" - - restriction_mode: Optional[EgressNetworkPolicyInternetAccessPolicyRestrictionMode] = None - """At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: - Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can - only access explicitly allowed internet and storage destinations, as well as UC connections and - external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via - private link.""" - - def as_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicy into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.allowed_internet_destinations: - body["allowed_internet_destinations"] = [v.as_dict() for v in self.allowed_internet_destinations] - if self.allowed_storage_destinations: - body["allowed_storage_destinations"] = [v.as_dict() for v in self.allowed_storage_destinations] - if self.log_only_mode: - body["log_only_mode"] = self.log_only_mode.as_dict() - if self.restriction_mode is not None: - body["restriction_mode"] = self.restriction_mode.value - return body - - def as_shallow_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicy into a shallow dictionary of its immediate attributes.""" - body = {} - if self.allowed_internet_destinations: - body["allowed_internet_destinations"] = self.allowed_internet_destinations - if self.allowed_storage_destinations: - body["allowed_storage_destinations"] = self.allowed_storage_destinations - if self.log_only_mode: - body["log_only_mode"] = self.log_only_mode - if self.restriction_mode is not None: - body["restriction_mode"] = self.restriction_mode - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicy: - """Deserializes the EgressNetworkPolicyInternetAccessPolicy from a dictionary.""" - return cls( - allowed_internet_destinations=_repeated_dict( - d, "allowed_internet_destinations", EgressNetworkPolicyInternetAccessPolicyInternetDestination - ), - allowed_storage_destinations=_repeated_dict( - d, "allowed_storage_destinations", EgressNetworkPolicyInternetAccessPolicyStorageDestination - ), - log_only_mode=_from_dict(d, "log_only_mode", EgressNetworkPolicyInternetAccessPolicyLogOnlyMode), - restriction_mode=_enum(d, "restriction_mode", EgressNetworkPolicyInternetAccessPolicyRestrictionMode), - ) - - -@dataclass -class EgressNetworkPolicyInternetAccessPolicyInternetDestination: - """Users can specify accessible internet destinations when outbound access is restricted. We only - support domain name (FQDN) destinations for the time being, though going forwards we want to - support host names and IP addresses.""" - - destination: Optional[str] = None - - protocol: Optional[ - EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol - ] = None - """The filtering protocol used by the DP. For private and public preview, SEG will only support TCP - filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be - set to TCP by default and hidden from the user. In the future, users may be able to select HTTP - filtering (i.e. SNI based filtering, filtering by FQDN).""" - - type: Optional[EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType] = None - - def as_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.destination is not None: - body["destination"] = self.destination - if self.protocol is not None: - body["protocol"] = self.protocol.value - if self.type is not None: - body["type"] = self.type.value - return body - - def as_shallow_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination into a shallow dictionary of its immediate attributes.""" - body = {} - if self.destination is not None: - body["destination"] = self.destination - if self.protocol is not None: - body["protocol"] = self.protocol - if self.type is not None: - body["type"] = self.type - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyInternetDestination: - """Deserializes the EgressNetworkPolicyInternetAccessPolicyInternetDestination from a dictionary.""" - return cls( - destination=d.get("destination", None), - protocol=_enum( - d, - "protocol", - EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol, - ), - type=_enum(d, "type", EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType), - ) - - -class EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol(Enum): - """The filtering protocol used by the DP. For private and public preview, SEG will only support TCP - filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be - set to TCP by default and hidden from the user. In the future, users may be able to select HTTP - filtering (i.e. SNI based filtering, filtering by FQDN).""" - - TCP = "TCP" - - -class EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType(Enum): - - FQDN = "FQDN" - - -@dataclass -class EgressNetworkPolicyInternetAccessPolicyLogOnlyMode: - log_only_mode_type: Optional[EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType] = None - - workloads: Optional[List[EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType]] = None - - def as_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.log_only_mode_type is not None: - body["log_only_mode_type"] = self.log_only_mode_type.value - if self.workloads: - body["workloads"] = [v.value for v in self.workloads] - return body - - def as_shallow_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode into a shallow dictionary of its immediate attributes.""" - body = {} - if self.log_only_mode_type is not None: - body["log_only_mode_type"] = self.log_only_mode_type - if self.workloads: - body["workloads"] = self.workloads - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyLogOnlyMode: - """Deserializes the EgressNetworkPolicyInternetAccessPolicyLogOnlyMode from a dictionary.""" - return cls( - log_only_mode_type=_enum( - d, "log_only_mode_type", EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType - ), - workloads=_repeated_enum(d, "workloads", EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType), - ) - - -class EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType(Enum): - - ALL_SERVICES = "ALL_SERVICES" - SELECTED_SERVICES = "SELECTED_SERVICES" - - -class EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType(Enum): - """The values should match the list of workloads used in networkconfig.proto""" - - DBSQL = "DBSQL" - ML_SERVING = "ML_SERVING" - - -class EgressNetworkPolicyInternetAccessPolicyRestrictionMode(Enum): - """At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: - Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can - only access explicitly allowed internet and storage destinations, as well as UC connections and - external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via - private link.""" - - FULL_ACCESS = "FULL_ACCESS" - PRIVATE_ACCESS_ONLY = "PRIVATE_ACCESS_ONLY" - RESTRICTED_ACCESS = "RESTRICTED_ACCESS" - - -@dataclass -class EgressNetworkPolicyInternetAccessPolicyStorageDestination: - """Users can specify accessible storage destinations.""" - - allowed_paths: Optional[List[str]] = None - - azure_container: Optional[str] = None - - azure_dns_zone: Optional[str] = None - - azure_storage_account: Optional[str] = None - - azure_storage_service: Optional[str] = None - - bucket_name: Optional[str] = None - - region: Optional[str] = None - - type: Optional[EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType] = None - - def as_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.allowed_paths: - body["allowed_paths"] = [v for v in self.allowed_paths] - if self.azure_container is not None: - body["azure_container"] = self.azure_container - if self.azure_dns_zone is not None: - body["azure_dns_zone"] = self.azure_dns_zone - if self.azure_storage_account is not None: - body["azure_storage_account"] = self.azure_storage_account - if self.azure_storage_service is not None: - body["azure_storage_service"] = self.azure_storage_service - if self.bucket_name is not None: - body["bucket_name"] = self.bucket_name - if self.region is not None: - body["region"] = self.region - if self.type is not None: - body["type"] = self.type.value - return body - - def as_shallow_dict(self) -> dict: - """Serializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination into a shallow dictionary of its immediate attributes.""" - body = {} - if self.allowed_paths: - body["allowed_paths"] = self.allowed_paths - if self.azure_container is not None: - body["azure_container"] = self.azure_container - if self.azure_dns_zone is not None: - body["azure_dns_zone"] = self.azure_dns_zone - if self.azure_storage_account is not None: - body["azure_storage_account"] = self.azure_storage_account - if self.azure_storage_service is not None: - body["azure_storage_service"] = self.azure_storage_service - if self.bucket_name is not None: - body["bucket_name"] = self.bucket_name - if self.region is not None: - body["region"] = self.region - if self.type is not None: - body["type"] = self.type - return body - - @classmethod - def from_dict(cls, d: Dict[str, Any]) -> EgressNetworkPolicyInternetAccessPolicyStorageDestination: - """Deserializes the EgressNetworkPolicyInternetAccessPolicyStorageDestination from a dictionary.""" - return cls( - allowed_paths=d.get("allowed_paths", None), - azure_container=d.get("azure_container", None), - azure_dns_zone=d.get("azure_dns_zone", None), - azure_storage_account=d.get("azure_storage_account", None), - azure_storage_service=d.get("azure_storage_service", None), - bucket_name=d.get("bucket_name", None), - region=d.get("region", None), - type=_enum(d, "type", EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType), - ) - - -class EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType(Enum): - - AWS_S3 = "AWS_S3" - AZURE_STORAGE = "AZURE_STORAGE" - CLOUDFLARE_R2 = "CLOUDFLARE_R2" - GOOGLE_CLOUD_STORAGE = "GOOGLE_CLOUD_STORAGE" - - @dataclass class EmailConfig: addresses: Optional[List[str]] = None @@ -2004,6 +1705,40 @@ def from_dict(cls, d: Dict[str, Any]) -> Empty: return cls() +@dataclass +class EnableResultsDownloading: + boolean_val: Optional[BooleanMessage] = None + + setting_name: Optional[str] = None + """Name of the corresponding setting. This field is populated in the response, but it will not be + respected even if it's set in the request body. The setting name in the path parameter will be + respected instead. Setting name is required to be 'default' if the setting only has one instance + per workspace.""" + + def as_dict(self) -> dict: + """Serializes the EnableResultsDownloading into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.boolean_val: + body["boolean_val"] = self.boolean_val.as_dict() + if self.setting_name is not None: + body["setting_name"] = self.setting_name + return body + + def as_shallow_dict(self) -> dict: + """Serializes the EnableResultsDownloading into a shallow dictionary of its immediate attributes.""" + body = {} + if self.boolean_val: + body["boolean_val"] = self.boolean_val + if self.setting_name is not None: + body["setting_name"] = self.setting_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> EnableResultsDownloading: + """Deserializes the EnableResultsDownloading from a dictionary.""" + return cls(boolean_val=_from_dict(d, "boolean_val", BooleanMessage), setting_name=d.get("setting_name", None)) + + @dataclass class EnhancedSecurityMonitoring: """SHIELD feature: ESM""" @@ -4146,6 +3881,10 @@ class TokenType(Enum): ARCLIGHT_AZURE_EXCHANGE_TOKEN = "ARCLIGHT_AZURE_EXCHANGE_TOKEN" ARCLIGHT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY = "ARCLIGHT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY" + ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN = "ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN" + ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY = ( + "ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY" + ) AZURE_ACTIVE_DIRECTORY_TOKEN = "AZURE_ACTIVE_DIRECTORY_TOKEN" @@ -4676,6 +4415,58 @@ def from_dict(cls, d: Dict[str, Any]) -> UpdateDisableLegacyFeaturesRequest: ) +@dataclass +class UpdateEnableResultsDownloadingRequest: + """Details required to update a setting.""" + + allow_missing: bool + """This should always be set to true for Settings API. Added for AIP compliance.""" + + setting: EnableResultsDownloading + + field_mask: str + """The field mask must be a single string, with multiple fields separated by commas (no spaces). + The field path is relative to the resource object, using a dot (`.`) to navigate sub-fields + (e.g., `author.given_name`). Specification of elements in sequence or map fields is not allowed, + as only the entire collection field can be specified. Field names must exactly match the + resource field names. + + A field mask of `*` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using `*` wildcards, as it can lead to unintended results if the + API changes in the future.""" + + def as_dict(self) -> dict: + """Serializes the UpdateEnableResultsDownloadingRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.allow_missing is not None: + body["allow_missing"] = self.allow_missing + if self.field_mask is not None: + body["field_mask"] = self.field_mask + if self.setting: + body["setting"] = self.setting.as_dict() + return body + + def as_shallow_dict(self) -> dict: + """Serializes the UpdateEnableResultsDownloadingRequest into a shallow dictionary of its immediate attributes.""" + body = {} + if self.allow_missing is not None: + body["allow_missing"] = self.allow_missing + if self.field_mask is not None: + body["field_mask"] = self.field_mask + if self.setting: + body["setting"] = self.setting + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> UpdateEnableResultsDownloadingRequest: + """Deserializes the UpdateEnableResultsDownloadingRequest from a dictionary.""" + return cls( + allow_missing=d.get("allow_missing", None), + field_mask=d.get("field_mask", None), + setting=_from_dict(d, "setting", EnableResultsDownloading), + ) + + @dataclass class UpdateEnhancedSecurityMonitoringSettingRequest: """Details required to update a setting.""" @@ -5255,37 +5046,6 @@ class AccountSettingsAPI: def __init__(self, api_client): self._api = api_client - self._csp_enablement_account = CspEnablementAccountAPI(self._api) - self._disable_legacy_features = DisableLegacyFeaturesAPI(self._api) - self._enable_ip_access_lists = EnableIpAccessListsAPI(self._api) - self._esm_enablement_account = EsmEnablementAccountAPI(self._api) - self._personal_compute = PersonalComputeAPI(self._api) - - @property - def csp_enablement_account(self) -> CspEnablementAccountAPI: - """The compliance security profile settings at the account level control whether to enable it for new workspaces.""" - return self._csp_enablement_account - - @property - def disable_legacy_features(self) -> DisableLegacyFeaturesAPI: - """Disable legacy features for new Databricks workspaces.""" - return self._disable_legacy_features - - @property - def enable_ip_access_lists(self) -> EnableIpAccessListsAPI: - """Controls the enforcement of IP access lists for accessing the account console.""" - return self._enable_ip_access_lists - - @property - def esm_enablement_account(self) -> EsmEnablementAccountAPI: - """The enhanced security monitoring setting at the account level controls whether to enable the feature on new workspaces.""" - return self._esm_enablement_account - - @property - def personal_compute(self) -> PersonalComputeAPI: - """The Personal Compute enablement setting lets you control which users can use the Personal Compute default policy to create compute resources.""" - return self._personal_compute - class AibiDashboardEmbeddingAccessPolicyAPI: """Controls whether AI/BI published dashboard embedding is enabled, conditionally enabled, or disabled at the @@ -5915,10 +5675,9 @@ def update(self, allow_missing: bool, setting: DefaultNamespaceSetting, field_ma class DisableLegacyAccessAPI: """'Disabling legacy access' has the following impacts: - 1. Disables direct access to the Hive Metastore. However, you can still access Hive Metastore through HMS - Federation. 2. Disables Fallback Mode (docs link) on any External Location access from the workspace. 3. - Alters DBFS path access to use External Location permissions in place of legacy credentials. 4. Enforces - Unity Catalog access on all path based access.""" + 1. Disables direct access to Hive Metastores from the workspace. However, you can still access a Hive + Metastore through Hive Metastore federation. 2. Disables fallback mode on external location access from + the workspace. 3. Disables Databricks Runtime versions prior to 13.3LTS.""" def __init__(self, api_client): self._api = api_client @@ -6337,6 +6096,70 @@ def update(self, allow_missing: bool, setting: AccountIpAccessEnable, field_mask return AccountIpAccessEnable.from_dict(res) +class EnableResultsDownloadingAPI: + """The Enable Results Downloading API controls the workspace level conf for the enablement of downloading + results.""" + + def __init__(self, api_client): + self._api = api_client + + def get_enable_results_downloading(self) -> EnableResultsDownloading: + """Get the Enable Results Downloading setting. + + Gets the Enable Results Downloading setting. + + :returns: :class:`EnableResultsDownloading` + """ + + headers = { + "Accept": "application/json", + } + + res = self._api.do("GET", "/api/2.0/settings/types/enable-results-downloading/names/default", headers=headers) + return EnableResultsDownloading.from_dict(res) + + def patch_enable_results_downloading( + self, allow_missing: bool, setting: EnableResultsDownloading, field_mask: str + ) -> EnableResultsDownloading: + """Update the Enable Results Downloading setting. + + Updates the Enable Results Downloading setting. The model follows eventual consistency, which means + the get after the update operation might receive stale values for some time. + + :param allow_missing: bool + This should always be set to true for Settings API. Added for AIP compliance. + :param setting: :class:`EnableResultsDownloading` + :param field_mask: str + The field mask must be a single string, with multiple fields separated by commas (no spaces). The + field path is relative to the resource object, using a dot (`.`) to navigate sub-fields (e.g., + `author.given_name`). Specification of elements in sequence or map fields is not allowed, as only + the entire collection field can be specified. Field names must exactly match the resource field + names. + + A field mask of `*` indicates full replacement. It’s recommended to always explicitly list the + fields being updated and avoid using `*` wildcards, as it can lead to unintended results if the API + changes in the future. + + :returns: :class:`EnableResultsDownloading` + """ + body = {} + if allow_missing is not None: + body["allow_missing"] = allow_missing + if field_mask is not None: + body["field_mask"] = field_mask + if setting is not None: + body["setting"] = setting.as_dict() + headers = { + "Accept": "application/json", + "Content-Type": "application/json", + } + + res = self._api.do( + "PATCH", "/api/2.0/settings/types/enable-results-downloading/names/default", body=body, headers=headers + ) + return EnableResultsDownloading.from_dict(res) + + class EnhancedSecurityMonitoringAPI: """Controls whether enhanced security monitoring is enabled for the current workspace. If the compliance security profile is enabled, this is automatically enabled. By default, it is disabled. However, if the @@ -7340,61 +7163,6 @@ class SettingsAPI: def __init__(self, api_client): self._api = api_client - self._aibi_dashboard_embedding_access_policy = AibiDashboardEmbeddingAccessPolicyAPI(self._api) - self._aibi_dashboard_embedding_approved_domains = AibiDashboardEmbeddingApprovedDomainsAPI(self._api) - self._automatic_cluster_update = AutomaticClusterUpdateAPI(self._api) - self._compliance_security_profile = ComplianceSecurityProfileAPI(self._api) - self._default_namespace = DefaultNamespaceAPI(self._api) - self._disable_legacy_access = DisableLegacyAccessAPI(self._api) - self._disable_legacy_dbfs = DisableLegacyDbfsAPI(self._api) - self._enhanced_security_monitoring = EnhancedSecurityMonitoringAPI(self._api) - self._restrict_workspace_admins = RestrictWorkspaceAdminsAPI(self._api) - - @property - def aibi_dashboard_embedding_access_policy(self) -> AibiDashboardEmbeddingAccessPolicyAPI: - """Controls whether AI/BI published dashboard embedding is enabled, conditionally enabled, or disabled at the workspace level.""" - return self._aibi_dashboard_embedding_access_policy - - @property - def aibi_dashboard_embedding_approved_domains(self) -> AibiDashboardEmbeddingApprovedDomainsAPI: - """Controls the list of domains approved to host the embedded AI/BI dashboards.""" - return self._aibi_dashboard_embedding_approved_domains - - @property - def automatic_cluster_update(self) -> AutomaticClusterUpdateAPI: - """Controls whether automatic cluster update is enabled for the current workspace.""" - return self._automatic_cluster_update - - @property - def compliance_security_profile(self) -> ComplianceSecurityProfileAPI: - """Controls whether to enable the compliance security profile for the current workspace.""" - return self._compliance_security_profile - - @property - def default_namespace(self) -> DefaultNamespaceAPI: - """The default namespace setting API allows users to configure the default namespace for a Databricks workspace.""" - return self._default_namespace - - @property - def disable_legacy_access(self) -> DisableLegacyAccessAPI: - """'Disabling legacy access' has the following impacts: 1.""" - return self._disable_legacy_access - - @property - def disable_legacy_dbfs(self) -> DisableLegacyDbfsAPI: - """When this setting is on, access to DBFS root and DBFS mounts is disallowed (as well as creation of new mounts).""" - return self._disable_legacy_dbfs - - @property - def enhanced_security_monitoring(self) -> EnhancedSecurityMonitoringAPI: - """Controls whether enhanced security monitoring is enabled for the current workspace.""" - return self._enhanced_security_monitoring - - @property - def restrict_workspace_admins(self) -> RestrictWorkspaceAdminsAPI: - """The Restrict Workspace Admins setting lets you control the capabilities of workspace admins.""" - return self._restrict_workspace_admins - class TokenManagementAPI: """Enables administrators to get all tokens and delete tokens for other users. Admins can either get every diff --git a/databricks/sdk/sharing/__init__.py b/databricks/sdk/sharing/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/sharing/v2/__init__.py b/databricks/sdk/sharing/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/sharing/v2/client.py b/databricks/sdk/sharing/v2/client.py new file mode 100755 index 000000000..5dd358003 --- /dev/null +++ b/databricks/sdk/sharing/v2/client.py @@ -0,0 +1,300 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .sharing import (ProvidersAPI, RecipientActivationAPI, RecipientsAPI, + SharesAPI) + +_LOG = logging.getLogger(__name__) + + +class ProvidersClient(ProvidersAPI): + """ + A data provider is an object representing the organization in the real world who shares the + data. A provider contains shares which further contain the shared data. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class RecipientActivationClient(RecipientActivationAPI): + """ + The Recipient Activation API is only applicable in the open sharing model where the recipient + object has the authentication type of `TOKEN`. The data recipient follows the activation link + shared by the data provider to download the credential file that includes the access token. The + recipient will then use the credential file to establish a secure connection with the provider + to receive the shared data. + + Note that you can download the credential file only once. Recipients should treat the downloaded + credential as a secret and must not share it outside of their organization. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class RecipientsClient(RecipientsAPI): + """ + A recipient is an object you create using :method:recipients/create to represent an organization + which you want to allow access shares. The way how sharing works differs depending on whether or + not your recipient has access to a Databricks workspace that is enabled for Unity Catalog: + + - For recipients with access to a Databricks workspace that is enabled for Unity Catalog, you + can create a recipient object along with a unique sharing identifier you get from the recipient. + The sharing identifier is the key identifier that enables the secure connection. This sharing + mode is called **Databricks-to-Databricks sharing**. + + - For recipients without access to a Databricks workspace that is enabled for Unity Catalog, + when you create a recipient object, Databricks generates an activation link you can send to the + recipient. The recipient follows the activation link to download the credential file, and then + uses the credential file to establish a secure connection to receive the shared data. This + sharing mode is called **open sharing**. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class SharesClient(SharesAPI): + """ + A share is a container instantiated with :method:shares/create. Once created you can iteratively + register a collection of existing data assets defined within the metastore using + :method:shares/update. You can register data assets under their original name, qualified by + their original schema, or provide alternate exposed names. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/sharing.py b/databricks/sdk/sharing/v2/sharing.py similarity index 98% rename from databricks/sdk/service/sharing.py rename to databricks/sdk/sharing/v2/sharing.py index 7325e5fdd..c8135ceb9 100755 --- a/databricks/sdk/service/sharing.py +++ b/databricks/sdk/sharing/v2/sharing.py @@ -7,13 +7,11 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - -from databricks.sdk.service import catalog - # all definitions in this file are in alphabetical order @@ -370,7 +368,7 @@ class DeltaSharingFunction: storage_location: Optional[str] = None """The storage location of the function.""" - tags: Optional[List[catalog.TagKeyValue]] = None + tags: Optional[List[TagKeyValue]] = None """The tags of the function.""" def as_dict(self) -> dict: @@ -466,7 +464,7 @@ def from_dict(cls, d: Dict[str, Any]) -> DeltaSharingFunction: share=d.get("share", None), share_id=d.get("share_id", None), storage_location=d.get("storage_location", None), - tags=_repeated_dict(d, "tags", catalog.TagKeyValue), + tags=_repeated_dict(d, "tags", TagKeyValue), ) @@ -1011,7 +1009,7 @@ class NotebookFile: share_id: Optional[str] = None """The id of the share that the notebook file belongs to.""" - tags: Optional[List[catalog.TagKeyValue]] = None + tags: Optional[List[TagKeyValue]] = None """The tags of the notebook file.""" def as_dict(self) -> dict: @@ -1057,7 +1055,7 @@ def from_dict(cls, d: Dict[str, Any]) -> NotebookFile: name=d.get("name", None), share=d.get("share", None), share_id=d.get("share_id", None), - tags=_repeated_dict(d, "tags", catalog.TagKeyValue), + tags=_repeated_dict(d, "tags", TagKeyValue), ) @@ -2247,7 +2245,7 @@ class Table: share_id: Optional[str] = None """The id of the share that the table belongs to.""" - tags: Optional[List[catalog.TagKeyValue]] = None + tags: Optional[List[TagKeyValue]] = None """The Tags of the table.""" def as_dict(self) -> dict: @@ -2308,7 +2306,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Table: schema=d.get("schema", None), share=d.get("share", None), share_id=d.get("share_id", None), - tags=_repeated_dict(d, "tags", catalog.TagKeyValue), + tags=_repeated_dict(d, "tags", TagKeyValue), ) @@ -2380,6 +2378,38 @@ class TableInternalAttributesSharedTableType(Enum): VIEW = "VIEW" +@dataclass +class TagKeyValue: + key: Optional[str] = None + """name of the tag""" + + value: Optional[str] = None + """value of the tag associated with the key, could be optional""" + + def as_dict(self) -> dict: + """Serializes the TagKeyValue into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.key is not None: + body["key"] = self.key + if self.value is not None: + body["value"] = self.value + return body + + def as_shallow_dict(self) -> dict: + """Serializes the TagKeyValue into a shallow dictionary of its immediate attributes.""" + body = {} + if self.key is not None: + body["key"] = self.key + if self.value is not None: + body["value"] = self.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> TagKeyValue: + """Deserializes the TagKeyValue from a dictionary.""" + return cls(key=d.get("key", None), value=d.get("value", None)) + + @dataclass class UpdateProvider: comment: Optional[str] = None @@ -2665,7 +2695,7 @@ class Volume: share_id: Optional[str] = None """/ The id of the share that the volume belongs to.""" - tags: Optional[List[catalog.TagKeyValue]] = None + tags: Optional[List[TagKeyValue]] = None """The tags of the volume.""" def as_dict(self) -> dict: @@ -2721,7 +2751,7 @@ def from_dict(cls, d: Dict[str, Any]) -> Volume: schema=d.get("schema", None), share=d.get("share", None), share_id=d.get("share_id", None), - tags=_repeated_dict(d, "tags", catalog.TagKeyValue), + tags=_repeated_dict(d, "tags", TagKeyValue), ) diff --git a/databricks/sdk/sql/__init__.py b/databricks/sdk/sql/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/sql/v2/__init__.py b/databricks/sdk/sql/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/sql/v2/client.py b/databricks/sdk/sql/v2/client.py new file mode 100755 index 000000000..f3c53b12b --- /dev/null +++ b/databricks/sdk/sql/v2/client.py @@ -0,0 +1,1089 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .sql import (AlertsAPI, AlertsLegacyAPI, DashboardsAPI, + DashboardWidgetsAPI, DataSourcesAPI, DbsqlPermissionsAPI, + QueriesAPI, QueriesLegacyAPI, QueryHistoryAPI, + QueryVisualizationsAPI, QueryVisualizationsLegacyAPI, + RedashConfigAPI, StatementExecutionAPI, WarehousesAPI) + +_LOG = logging.getLogger(__name__) + + +class AlertsClient(AlertsAPI): + """ + The alerts API can be used to perform CRUD operations on alerts. An alert is a Databricks SQL + object that periodically runs a query, evaluates a condition of its result, and notifies one or + more users and/or notification destinations if the condition was met. Alerts can be scheduled + using the `sql_task` type of the Jobs API, e.g. :method:jobs/create. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class AlertsLegacyClient(AlertsLegacyAPI): + """ + The alerts API can be used to perform CRUD operations on alerts. An alert is a Databricks SQL + object that periodically runs a query, evaluates a condition of its result, and notifies one or + more users and/or notification destinations if the condition was met. Alerts can be scheduled + using the `sql_task` type of the Jobs API, e.g. :method:jobs/create. + + **Note**: A new version of the Databricks SQL API is now available. Please see the latest + version. [Learn more] + + [Learn more]: https://docs.databricks.com/en/sql/dbsql-api-latest.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DashboardWidgetsClient(DashboardWidgetsAPI): + """ + This is an evolving API that facilitates the addition and removal of widgets from existing + dashboards within the Databricks Workspace. Data structures may change over time. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DashboardsClient(DashboardsAPI): + """ + In general, there is little need to modify dashboards using the API. However, it can be useful + to use dashboard objects to look-up a collection of related query IDs. The API can also be used + to duplicate multiple dashboards at once since you can get a dashboard definition with a GET + request and then POST it to create a new one. Dashboards can be scheduled using the `sql_task` + type of the Jobs API, e.g. :method:jobs/create. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DataSourcesClient(DataSourcesAPI): + """ + This API is provided to assist you in making new query objects. When creating a query object, + you may optionally specify a `data_source_id` for the SQL warehouse against which it will run. + If you don't already know the `data_source_id` for your desired SQL warehouse, this API will + help you find it. + + This API does not support searches. It returns the full list of SQL warehouses in your + workspace. We advise you to use any text editor, REST client, or `grep` to search the response + from this API for the name of your SQL warehouse as it appears in Databricks SQL. + + **Note**: A new version of the Databricks SQL API is now available. [Learn more] + + [Learn more]: https://docs.databricks.com/en/sql/dbsql-api-latest.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class DbsqlPermissionsClient(DbsqlPermissionsAPI): + """ + The SQL Permissions API is similar to the endpoints of the :method:permissions/set. However, + this exposes only one endpoint, which gets the Access Control List for a given object. You + cannot modify any permissions using this API. + + There are three levels of permission: + + - `CAN_VIEW`: Allows read-only access + + - `CAN_RUN`: Allows read access and run access (superset of `CAN_VIEW`) + + - `CAN_MANAGE`: Allows all actions: read, run, edit, delete, modify permissions (superset of + `CAN_RUN`) + + **Note**: A new version of the Databricks SQL API is now available. [Learn more] + + [Learn more]: https://docs.databricks.com/en/sql/dbsql-api-latest.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueriesClient(QueriesAPI): + """ + The queries API can be used to perform CRUD operations on queries. A query is a Databricks SQL + object that includes the target SQL warehouse, query text, name, description, tags, and + parameters. Queries can be scheduled using the `sql_task` type of the Jobs API, e.g. + :method:jobs/create. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueriesLegacyClient(QueriesLegacyAPI): + """ + These endpoints are used for CRUD operations on query definitions. Query definitions include the + target SQL warehouse, query text, name, description, tags, parameters, and visualizations. + Queries can be scheduled using the `sql_task` type of the Jobs API, e.g. :method:jobs/create. + + **Note**: A new version of the Databricks SQL API is now available. Please see the latest + version. [Learn more] + + [Learn more]: https://docs.databricks.com/en/sql/dbsql-api-latest.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueryHistoryClient(QueryHistoryAPI): + """ + A service responsible for storing and retrieving the list of queries run against SQL endpoints + and serverless compute. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueryVisualizationsClient(QueryVisualizationsAPI): + """ + This is an evolving API that facilitates the addition and removal of visualizations from + existing queries in the Databricks Workspace. Data structures can change over time. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class QueryVisualizationsLegacyClient(QueryVisualizationsLegacyAPI): + """ + This is an evolving API that facilitates the addition and removal of vizualisations from + existing queries within the Databricks Workspace. Data structures may change over time. + + **Note**: A new version of the Databricks SQL API is now available. Please see the latest + version. [Learn more] + + [Learn more]: https://docs.databricks.com/en/sql/dbsql-api-latest.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class RedashConfigClient(RedashConfigAPI): + """ + Redash V2 service for workspace configurations (internal) + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class StatementExecutionClient(StatementExecutionAPI): + """ + The Databricks SQL Statement Execution API can be used to execute SQL statements on a SQL + warehouse and fetch the result. + + **Getting started** + + We suggest beginning with the [Databricks SQL Statement Execution API tutorial]. + + **Overview of statement execution and result fetching** + + Statement execution begins by issuing a :method:statementexecution/executeStatement request with + a valid SQL statement and warehouse ID, along with optional parameters such as the data catalog + and output format. If no other parameters are specified, the server will wait for up to 10s + before returning a response. If the statement has completed within this timespan, the response + will include the result data as a JSON array and metadata. Otherwise, if no result is available + after the 10s timeout expired, the response will provide the statement ID that can be used to + poll for results by using a :method:statementexecution/getStatement request. + + You can specify whether the call should behave synchronously, asynchronously or start + synchronously with a fallback to asynchronous execution. This is controlled with the + `wait_timeout` and `on_wait_timeout` settings. If `wait_timeout` is set between 5-50 seconds + (default: 10s), the call waits for results up to the specified timeout; when set to `0s`, the + call is asynchronous and responds immediately with a statement ID. The `on_wait_timeout` setting + specifies what should happen when the timeout is reached while the statement execution has not + yet finished. This can be set to either `CONTINUE`, to fallback to asynchronous mode, or it can + be set to `CANCEL`, which cancels the statement. + + In summary: - Synchronous mode - `wait_timeout=30s` and `on_wait_timeout=CANCEL` - The call + waits up to 30 seconds; if the statement execution finishes within this time, the result data is + returned directly in the response. If the execution takes longer than 30 seconds, the execution + is canceled and the call returns with a `CANCELED` state. - Asynchronous mode - + `wait_timeout=0s` (`on_wait_timeout` is ignored) - The call doesn't wait for the statement to + finish but returns directly with a statement ID. The status of the statement execution can be + polled by issuing :method:statementexecution/getStatement with the statement ID. Once the + execution has succeeded, this call also returns the result and metadata in the response. - + Hybrid mode (default) - `wait_timeout=10s` and `on_wait_timeout=CONTINUE` - The call waits for + up to 10 seconds; if the statement execution finishes within this time, the result data is + returned directly in the response. If the execution takes longer than 10 seconds, a statement ID + is returned. The statement ID can be used to fetch status and results in the same way as in the + asynchronous mode. + + Depending on the size, the result can be split into multiple chunks. If the statement execution + is successful, the statement response contains a manifest and the first chunk of the result. The + manifest contains schema information and provides metadata for each chunk in the result. Result + chunks can be retrieved by index with :method:statementexecution/getStatementResultChunkN which + may be called in any order and in parallel. For sequential fetching, each chunk, apart from the + last, also contains a `next_chunk_index` and `next_chunk_internal_link` that point to the next + chunk. + + A statement can be canceled with :method:statementexecution/cancelExecution. + + **Fetching result data: format and disposition** + + To specify the format of the result data, use the `format` field, which can be set to one of the + following options: `JSON_ARRAY` (JSON), `ARROW_STREAM` ([Apache Arrow Columnar]), or `CSV`. + + There are two ways to receive statement results, controlled by the `disposition` setting, which + can be either `INLINE` or `EXTERNAL_LINKS`: + + - `INLINE`: In this mode, the result data is directly included in the response. It's best suited + for smaller results. This mode can only be used with the `JSON_ARRAY` format. + + - `EXTERNAL_LINKS`: In this mode, the response provides links that can be used to download the + result data in chunks separately. This approach is ideal for larger results and offers higher + throughput. This mode can be used with all the formats: `JSON_ARRAY`, `ARROW_STREAM`, and `CSV`. + + By default, the API uses `format=JSON_ARRAY` and `disposition=INLINE`. + + **Limits and limitations** + + Note: The byte limit for INLINE disposition is based on internal storage metrics and will not + exactly match the byte count of the actual payload. + + - Statements with `disposition=INLINE` are limited to 25 MiB and will fail when this limit is + exceeded. - Statements with `disposition=EXTERNAL_LINKS` are limited to 100 GiB. Result sets + larger than this limit will be truncated. Truncation is indicated by the `truncated` field in + the result manifest. - The maximum query text size is 16 MiB. - Cancelation might silently fail. + A successful response from a cancel request indicates that the cancel request was successfully + received and sent to the processing engine. However, an outstanding statement might have already + completed execution when the cancel request arrives. Polling for status until a terminal state + is reached is a reliable way to determine the final state. - Wait timeouts are approximate, + occur server-side, and cannot account for things such as caller delays and network latency from + caller to service. - To guarantee that the statement is kept alive, you must poll at least once + every 15 minutes. - The results are only available for one hour after success; polling does not + extend this. - The SQL Execution API must be used for the entire lifecycle of the statement. For + example, you cannot use the Jobs API to execute the command, and then the SQL Execution API to + cancel it. + + [Apache Arrow Columnar]: https://arrow.apache.org/overview/ + [Databricks SQL Statement Execution API tutorial]: https://docs.databricks.com/sql/api/sql-execution-tutorial.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WarehousesClient(WarehousesAPI): + """ + A SQL warehouse is a compute resource that lets you run SQL commands on data objects within + Databricks SQL. Compute resources are infrastructure resources that provide processing + capabilities in the cloud. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/sql/v2/sql.py similarity index 98% rename from databricks/sdk/service/sql.py rename to databricks/sdk/sql/v2/sql.py index 926d137a9..89bfffd79 100755 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/sql/v2/sql.py @@ -3,19 +3,15 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict, _repeated_enum +from ...service._internal import (_enum, _from_dict, _repeated_dict, + _repeated_enum) _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -3161,6 +3157,123 @@ def from_dict(cls, d: Dict[str, Any]) -> ExternalLink: ) +@dataclass +class ExternalQuerySource: + alert_id: Optional[str] = None + """The canonical identifier for this SQL alert""" + + dashboard_id: Optional[str] = None + """The canonical identifier for this Lakeview dashboard""" + + genie_space_id: Optional[str] = None + """The canonical identifier for this Genie space""" + + job_info: Optional[ExternalQuerySourceJobInfo] = None + + legacy_dashboard_id: Optional[str] = None + """The canonical identifier for this legacy dashboard""" + + notebook_id: Optional[str] = None + """The canonical identifier for this notebook""" + + sql_query_id: Optional[str] = None + """The canonical identifier for this SQL query""" + + def as_dict(self) -> dict: + """Serializes the ExternalQuerySource into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.alert_id is not None: + body["alert_id"] = self.alert_id + if self.dashboard_id is not None: + body["dashboard_id"] = self.dashboard_id + if self.genie_space_id is not None: + body["genie_space_id"] = self.genie_space_id + if self.job_info: + body["job_info"] = self.job_info.as_dict() + if self.legacy_dashboard_id is not None: + body["legacy_dashboard_id"] = self.legacy_dashboard_id + if self.notebook_id is not None: + body["notebook_id"] = self.notebook_id + if self.sql_query_id is not None: + body["sql_query_id"] = self.sql_query_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ExternalQuerySource into a shallow dictionary of its immediate attributes.""" + body = {} + if self.alert_id is not None: + body["alert_id"] = self.alert_id + if self.dashboard_id is not None: + body["dashboard_id"] = self.dashboard_id + if self.genie_space_id is not None: + body["genie_space_id"] = self.genie_space_id + if self.job_info: + body["job_info"] = self.job_info + if self.legacy_dashboard_id is not None: + body["legacy_dashboard_id"] = self.legacy_dashboard_id + if self.notebook_id is not None: + body["notebook_id"] = self.notebook_id + if self.sql_query_id is not None: + body["sql_query_id"] = self.sql_query_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ExternalQuerySource: + """Deserializes the ExternalQuerySource from a dictionary.""" + return cls( + alert_id=d.get("alert_id", None), + dashboard_id=d.get("dashboard_id", None), + genie_space_id=d.get("genie_space_id", None), + job_info=_from_dict(d, "job_info", ExternalQuerySourceJobInfo), + legacy_dashboard_id=d.get("legacy_dashboard_id", None), + notebook_id=d.get("notebook_id", None), + sql_query_id=d.get("sql_query_id", None), + ) + + +@dataclass +class ExternalQuerySourceJobInfo: + job_id: Optional[str] = None + """The canonical identifier for this job.""" + + job_run_id: Optional[str] = None + """The canonical identifier of the run. This ID is unique across all runs of all jobs.""" + + job_task_run_id: Optional[str] = None + """The canonical identifier of the task run.""" + + def as_dict(self) -> dict: + """Serializes the ExternalQuerySourceJobInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.job_id is not None: + body["job_id"] = self.job_id + if self.job_run_id is not None: + body["job_run_id"] = self.job_run_id + if self.job_task_run_id is not None: + body["job_task_run_id"] = self.job_task_run_id + return body + + def as_shallow_dict(self) -> dict: + """Serializes the ExternalQuerySourceJobInfo into a shallow dictionary of its immediate attributes.""" + body = {} + if self.job_id is not None: + body["job_id"] = self.job_id + if self.job_run_id is not None: + body["job_run_id"] = self.job_run_id + if self.job_task_run_id is not None: + body["job_task_run_id"] = self.job_task_run_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, Any]) -> ExternalQuerySourceJobInfo: + """Deserializes the ExternalQuerySourceJobInfo from a dictionary.""" + return cls( + job_id=d.get("job_id", None), + job_run_id=d.get("job_run_id", None), + job_task_run_id=d.get("job_task_run_id", None), + ) + + class Format(Enum): ARROW_STREAM = "ARROW_STREAM" @@ -5077,6 +5190,11 @@ class QueryInfo: channel_used: Optional[ChannelInfo] = None """SQL Warehouse channel information at the time of query execution""" + client_application: Optional[str] = None + """Client application that ran the statement. For example: Databricks SQL Editor, Tableau, and + Power BI. This field is derived from information provided by client applications. While values + are expected to remain static over time, this cannot be guaranteed.""" + duration: Optional[int] = None """Total execution time of the statement ( excluding result fetch time ).""" @@ -5113,6 +5231,11 @@ class QueryInfo: query_id: Optional[str] = None """The query ID.""" + query_source: Optional[ExternalQuerySource] = None + """A struct that contains key-value pairs representing Databricks entities that were involved in + the execution of this statement, such as jobs, notebooks, or dashboards. This field only records + Databricks entities.""" + query_start_time_ms: Optional[int] = None """The time the query started.""" @@ -5149,6 +5272,8 @@ def as_dict(self) -> dict: body = {} if self.channel_used: body["channel_used"] = self.channel_used.as_dict() + if self.client_application is not None: + body["client_application"] = self.client_application if self.duration is not None: body["duration"] = self.duration if self.endpoint_id is not None: @@ -5173,6 +5298,8 @@ def as_dict(self) -> dict: body["query_end_time_ms"] = self.query_end_time_ms if self.query_id is not None: body["query_id"] = self.query_id + if self.query_source: + body["query_source"] = self.query_source.as_dict() if self.query_start_time_ms is not None: body["query_start_time_ms"] = self.query_start_time_ms if self.query_text is not None: @@ -5198,6 +5325,8 @@ def as_shallow_dict(self) -> dict: body = {} if self.channel_used: body["channel_used"] = self.channel_used + if self.client_application is not None: + body["client_application"] = self.client_application if self.duration is not None: body["duration"] = self.duration if self.endpoint_id is not None: @@ -5222,6 +5351,8 @@ def as_shallow_dict(self) -> dict: body["query_end_time_ms"] = self.query_end_time_ms if self.query_id is not None: body["query_id"] = self.query_id + if self.query_source: + body["query_source"] = self.query_source if self.query_start_time_ms is not None: body["query_start_time_ms"] = self.query_start_time_ms if self.query_text is not None: @@ -5247,6 +5378,7 @@ def from_dict(cls, d: Dict[str, Any]) -> QueryInfo: """Deserializes the QueryInfo from a dictionary.""" return cls( channel_used=_from_dict(d, "channel_used", ChannelInfo), + client_application=d.get("client_application", None), duration=d.get("duration", None), endpoint_id=d.get("endpoint_id", None), error_message=d.get("error_message", None), @@ -5259,6 +5391,7 @@ def from_dict(cls, d: Dict[str, Any]) -> QueryInfo: plans_state=_enum(d, "plans_state", PlansState), query_end_time_ms=d.get("query_end_time_ms", None), query_id=d.get("query_id", None), + query_source=_from_dict(d, "query_source", ExternalQuerySource), query_start_time_ms=d.get("query_start_time_ms", None), query_text=d.get("query_text", None), rows_produced=d.get("rows_produced", None), @@ -7409,6 +7542,7 @@ class WarehousePermissionLevel(Enum): CAN_MANAGE = "CAN_MANAGE" CAN_MONITOR = "CAN_MONITOR" CAN_USE = "CAN_USE" + CAN_VIEW = "CAN_VIEW" IS_OWNER = "IS_OWNER" @@ -9608,67 +9742,6 @@ class WarehousesAPI: def __init__(self, api_client): self._api = api_client - def wait_get_warehouse_running( - self, id: str, timeout=timedelta(minutes=20), callback: Optional[Callable[[GetWarehouseResponse], None]] = None - ) -> GetWarehouseResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (State.RUNNING,) - failure_states = ( - State.STOPPED, - State.DELETED, - ) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(id=id) - status = poll.state - status_message = f"current status: {status}" - if poll.health: - status_message = poll.health.summary - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach RUNNING, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"id={id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def wait_get_warehouse_stopped( - self, id: str, timeout=timedelta(minutes=20), callback: Optional[Callable[[GetWarehouseResponse], None]] = None - ) -> GetWarehouseResponse: - deadline = time.time() + timeout.total_seconds() - target_states = (State.STOPPED,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get(id=id) - status = poll.state - status_message = f"current status: {status}" - if poll.health: - status_message = poll.health.summary - if status in target_states: - return poll - if callback: - callback(poll) - prefix = f"id={id}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - def create( self, *, @@ -9685,7 +9758,7 @@ def create( spot_instance_policy: Optional[SpotInstancePolicy] = None, tags: Optional[EndpointTags] = None, warehouse_type: Optional[CreateWarehouseRequestWarehouseType] = None, - ) -> Wait[GetWarehouseResponse]: + ) -> CreateWarehouseResponse: """Create a warehouse. Creates a new SQL warehouse. @@ -9748,7 +9821,7 @@ def create( :returns: Long-running operation waiter for :class:`GetWarehouseResponse`. - See :method:wait_get_warehouse_running for more details. + See :method:WaitGetWarehouseRunning for more details. """ body = {} if auto_stop_mins is not None: @@ -9782,46 +9855,8 @@ def create( "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/sql/warehouses", body=body, headers=headers) - return Wait( - self.wait_get_warehouse_running, - response=CreateWarehouseResponse.from_dict(op_response), - id=op_response["id"], - ) - - def create_and_wait( - self, - *, - auto_stop_mins: Optional[int] = None, - channel: Optional[Channel] = None, - cluster_size: Optional[str] = None, - creator_name: Optional[str] = None, - enable_photon: Optional[bool] = None, - enable_serverless_compute: Optional[bool] = None, - instance_profile_arn: Optional[str] = None, - max_num_clusters: Optional[int] = None, - min_num_clusters: Optional[int] = None, - name: Optional[str] = None, - spot_instance_policy: Optional[SpotInstancePolicy] = None, - tags: Optional[EndpointTags] = None, - warehouse_type: Optional[CreateWarehouseRequestWarehouseType] = None, - timeout=timedelta(minutes=20), - ) -> GetWarehouseResponse: - return self.create( - auto_stop_mins=auto_stop_mins, - channel=channel, - cluster_size=cluster_size, - creator_name=creator_name, - enable_photon=enable_photon, - enable_serverless_compute=enable_serverless_compute, - instance_profile_arn=instance_profile_arn, - max_num_clusters=max_num_clusters, - min_num_clusters=min_num_clusters, - name=name, - spot_instance_policy=spot_instance_policy, - tags=tags, - warehouse_type=warehouse_type, - ).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/sql/warehouses", body=body, headers=headers) + return CreateWarehouseResponse.from_dict(res) def delete(self, id: str): """Delete a warehouse. @@ -9857,7 +9892,7 @@ def edit( spot_instance_policy: Optional[SpotInstancePolicy] = None, tags: Optional[EndpointTags] = None, warehouse_type: Optional[EditWarehouseRequestWarehouseType] = None, - ) -> Wait[GetWarehouseResponse]: + ): """Update a warehouse. Updates the configuration for a SQL warehouse. @@ -9921,7 +9956,7 @@ def edit( :returns: Long-running operation waiter for :class:`GetWarehouseResponse`. - See :method:wait_get_warehouse_running for more details. + See :method:WaitGetWarehouseRunning for more details. """ body = {} if auto_stop_mins is not None: @@ -9955,44 +9990,7 @@ def edit( "Content-Type": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/edit", body=body, headers=headers) - return Wait(self.wait_get_warehouse_running, response=EditWarehouseResponse.from_dict(op_response), id=id) - - def edit_and_wait( - self, - id: str, - *, - auto_stop_mins: Optional[int] = None, - channel: Optional[Channel] = None, - cluster_size: Optional[str] = None, - creator_name: Optional[str] = None, - enable_photon: Optional[bool] = None, - enable_serverless_compute: Optional[bool] = None, - instance_profile_arn: Optional[str] = None, - max_num_clusters: Optional[int] = None, - min_num_clusters: Optional[int] = None, - name: Optional[str] = None, - spot_instance_policy: Optional[SpotInstancePolicy] = None, - tags: Optional[EndpointTags] = None, - warehouse_type: Optional[EditWarehouseRequestWarehouseType] = None, - timeout=timedelta(minutes=20), - ) -> GetWarehouseResponse: - return self.edit( - auto_stop_mins=auto_stop_mins, - channel=channel, - cluster_size=cluster_size, - creator_name=creator_name, - enable_photon=enable_photon, - enable_serverless_compute=enable_serverless_compute, - id=id, - instance_profile_arn=instance_profile_arn, - max_num_clusters=max_num_clusters, - min_num_clusters=min_num_clusters, - name=name, - spot_instance_policy=spot_instance_policy, - tags=tags, - warehouse_type=warehouse_type, - ).result(timeout=timeout) + self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/edit", body=body, headers=headers) def get(self, id: str) -> GetWarehouseResponse: """Get warehouse info. @@ -10180,7 +10178,7 @@ def set_workspace_warehouse_config( self._api.do("PUT", "/api/2.0/sql/config/warehouses", body=body, headers=headers) - def start(self, id: str) -> Wait[GetWarehouseResponse]: + def start(self, id: str): """Start a warehouse. Starts a SQL warehouse. @@ -10190,20 +10188,16 @@ def start(self, id: str) -> Wait[GetWarehouseResponse]: :returns: Long-running operation waiter for :class:`GetWarehouseResponse`. - See :method:wait_get_warehouse_running for more details. + See :method:WaitGetWarehouseRunning for more details. """ headers = { "Accept": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/start", headers=headers) - return Wait(self.wait_get_warehouse_running, response=StartWarehouseResponse.from_dict(op_response), id=id) + self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/start", headers=headers) - def start_and_wait(self, id: str, timeout=timedelta(minutes=20)) -> GetWarehouseResponse: - return self.start(id=id).result(timeout=timeout) - - def stop(self, id: str) -> Wait[GetWarehouseResponse]: + def stop(self, id: str): """Stop a warehouse. Stops a SQL warehouse. @@ -10213,18 +10207,14 @@ def stop(self, id: str) -> Wait[GetWarehouseResponse]: :returns: Long-running operation waiter for :class:`GetWarehouseResponse`. - See :method:wait_get_warehouse_stopped for more details. + See :method:WaitGetWarehouseStopped for more details. """ headers = { "Accept": "application/json", } - op_response = self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/stop", headers=headers) - return Wait(self.wait_get_warehouse_stopped, response=StopWarehouseResponse.from_dict(op_response), id=id) - - def stop_and_wait(self, id: str, timeout=timedelta(minutes=20)) -> GetWarehouseResponse: - return self.stop(id=id).result(timeout=timeout) + self._api.do("POST", f"/api/2.0/sql/warehouses/{id}/stop", headers=headers) def update_permissions( self, warehouse_id: str, *, access_control_list: Optional[List[WarehouseAccessControlRequest]] = None diff --git a/databricks/sdk/vectorsearch/__init__.py b/databricks/sdk/vectorsearch/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/vectorsearch/v2/__init__.py b/databricks/sdk/vectorsearch/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/vectorsearch/v2/client.py b/databricks/sdk/vectorsearch/v2/client.py new file mode 100755 index 000000000..d78b24681 --- /dev/null +++ b/databricks/sdk/vectorsearch/v2/client.py @@ -0,0 +1,150 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .vectorsearch import VectorSearchEndpointsAPI, VectorSearchIndexesAPI + +_LOG = logging.getLogger(__name__) + + +class VectorSearchEndpointsClient(VectorSearchEndpointsAPI): + """ + **Endpoint**: Represents the compute resources to host vector search indexes. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class VectorSearchIndexesClient(VectorSearchIndexesAPI): + """ + **Index**: An efficient representation of your embedding vectors that supports real-time and + efficient approximate nearest neighbor (ANN) search queries. + + There are 2 types of Vector Search indexes: * **Delta Sync Index**: An index that automatically + syncs with a source Delta Table, automatically and incrementally updating the index as the + underlying data in the Delta Table changes. * **Direct Vector Access Index**: An index that + supports direct read and write of vectors and metadata through our REST and SDK APIs. With this + model, the user manages index updates. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/service/vectorsearch.py b/databricks/sdk/vectorsearch/v2/vectorsearch.py similarity index 97% rename from databricks/sdk/service/vectorsearch.py rename to databricks/sdk/vectorsearch/v2/vectorsearch.py index f8b8429c6..a0145cbf7 100755 --- a/databricks/sdk/service/vectorsearch.py +++ b/databricks/sdk/vectorsearch/v2/vectorsearch.py @@ -3,19 +3,14 @@ from __future__ import annotations import logging -import random -import time from dataclasses import dataclass -from datetime import timedelta from enum import Enum -from typing import Any, Callable, Dict, Iterator, List, Optional +from typing import Any, Dict, Iterator, List, Optional -from ..databricks.errors import OperationFailed -from ._internal import Wait, _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order @@ -1627,41 +1622,7 @@ class VectorSearchEndpointsAPI: def __init__(self, api_client): self._api = api_client - def wait_get_endpoint_vector_search_endpoint_online( - self, - endpoint_name: str, - timeout=timedelta(minutes=20), - callback: Optional[Callable[[EndpointInfo], None]] = None, - ) -> EndpointInfo: - deadline = time.time() + timeout.total_seconds() - target_states = (EndpointStatusState.ONLINE,) - failure_states = (EndpointStatusState.OFFLINE,) - status_message = "polling..." - attempt = 1 - while time.time() < deadline: - poll = self.get_endpoint(endpoint_name=endpoint_name) - status = poll.endpoint_status.state - status_message = f"current status: {status}" - if poll.endpoint_status: - status_message = poll.endpoint_status.message - if status in target_states: - return poll - if callback: - callback(poll) - if status in failure_states: - msg = f"failed to reach ONLINE, got {status}: {status_message}" - raise OperationFailed(msg) - prefix = f"endpoint_name={endpoint_name}" - sleep = attempt - if sleep > 10: - # sleep 10s max per attempt - sleep = 10 - _LOG.debug(f"{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)") - time.sleep(sleep + random.random()) - attempt += 1 - raise TimeoutError(f"timed out after {timeout}: {status_message}") - - def create_endpoint(self, name: str, endpoint_type: EndpointType) -> Wait[EndpointInfo]: + def create_endpoint(self, name: str, endpoint_type: EndpointType) -> EndpointInfo: """Create an endpoint. Create a new endpoint. @@ -1673,7 +1634,7 @@ def create_endpoint(self, name: str, endpoint_type: EndpointType) -> Wait[Endpoi :returns: Long-running operation waiter for :class:`EndpointInfo`. - See :method:wait_get_endpoint_vector_search_endpoint_online for more details. + See :method:WaitGetEndpointVectorSearchEndpointOnline for more details. """ body = {} if endpoint_type is not None: @@ -1685,17 +1646,8 @@ def create_endpoint(self, name: str, endpoint_type: EndpointType) -> Wait[Endpoi "Content-Type": "application/json", } - op_response = self._api.do("POST", "/api/2.0/vector-search/endpoints", body=body, headers=headers) - return Wait( - self.wait_get_endpoint_vector_search_endpoint_online, - response=EndpointInfo.from_dict(op_response), - endpoint_name=op_response["name"], - ) - - def create_endpoint_and_wait( - self, name: str, endpoint_type: EndpointType, timeout=timedelta(minutes=20) - ) -> EndpointInfo: - return self.create_endpoint(endpoint_type=endpoint_type, name=name).result(timeout=timeout) + res = self._api.do("POST", "/api/2.0/vector-search/endpoints", body=body, headers=headers) + return EndpointInfo.from_dict(res) def delete_endpoint(self, endpoint_name: str): """Delete an endpoint. diff --git a/databricks/sdk/workspace/__init__.py b/databricks/sdk/workspace/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/workspace/v2/__init__.py b/databricks/sdk/workspace/v2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/databricks/sdk/workspace/v2/client.py b/databricks/sdk/workspace/v2/client.py new file mode 100755 index 000000000..29394271b --- /dev/null +++ b/databricks/sdk/workspace/v2/client.py @@ -0,0 +1,300 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +import logging +from typing import Optional + +import databricks.sdk.databricks.core as client +from databricks.sdk.databricks.credentials_provider import CredentialsStrategy + +from .mixin import WorkspaceExt +from .workspace import GitCredentialsAPI, ReposAPI, SecretsAPI + +_LOG = logging.getLogger(__name__) + + +class GitCredentialsClient(GitCredentialsAPI): + """ + Registers personal access token for Databricks to do operations on behalf of the user. + + See [more info]. + + [more info]: https://docs.databricks.com/repos/get-access-tokens-from-git-provider.html + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class ReposClient(ReposAPI): + """ + The Repos API allows users to manage their git repos. Users can use the API to access all repos + that they have manage permissions on. + + Databricks Repos is a visual Git client in Databricks. It supports common Git operations such a + cloning a repository, committing and pushing, pulling, branch management, and visual comparison + of diffs when committing. + + Within Repos you can develop code in notebooks or other files and follow data science and + engineering code development best practices using Git for version control, collaboration, and + CI/CD. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class SecretsClient(SecretsAPI): + """ + The Secrets API allows you to manage secrets, secret scopes, and access permissions. + + Sometimes accessing data requires that you authenticate to external data sources through JDBC. + Instead of directly entering your credentials into a notebook, use Databricks secrets to store + your credentials and reference them in notebooks and jobs. + + Administrators, secret creators, and users granted permission can read Databricks secrets. While + Databricks makes an effort to redact secret values that might be displayed in notebooks, it is + not possible to prevent such users from reading secrets. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) + + +class WorkspaceClient(WorkspaceExt): + """ + The Workspace API allows you to list, import, export, and delete notebooks and folders. + + A notebook is a web-based interface to a document that contains runnable code, visualizations, + and explanatory text. + """ + + 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, + ) + self._config = config.copy() + super().__init__(client.ApiClient(config)) diff --git a/databricks/sdk/mixins/workspace.py b/databricks/sdk/workspace/v2/mixin.py similarity index 96% rename from databricks/sdk/mixins/workspace.py rename to databricks/sdk/workspace/v2/mixin.py index e77e117bc..c29db98ee 100644 --- a/databricks/sdk/mixins/workspace.py +++ b/databricks/sdk/workspace/v2/mixin.py @@ -1,8 +1,9 @@ from typing import Any, BinaryIO, Iterator, Optional, Union -from ..databricks.core import DatabricksError -from ..service.workspace import (ExportFormat, ImportFormat, Language, - ObjectInfo, ObjectType, WorkspaceAPI) +from databricks.sdk.databricks.core import DatabricksError + +from .workspace import (ExportFormat, ImportFormat, Language, ObjectInfo, + ObjectType, WorkspaceAPI) def _fqcn(x: Any) -> str: diff --git a/databricks/sdk/service/workspace.py b/databricks/sdk/workspace/v2/workspace.py similarity index 99% rename from databricks/sdk/service/workspace.py rename to databricks/sdk/workspace/v2/workspace.py index 6753ad880..502fb170e 100755 --- a/databricks/sdk/service/workspace.py +++ b/databricks/sdk/workspace/v2/workspace.py @@ -7,11 +7,10 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional -from ._internal import _enum, _from_dict, _repeated_dict +from ...service._internal import _enum, _from_dict, _repeated_dict _LOG = logging.getLogger("databricks.sdk") - # all definitions in this file are in alphabetical order diff --git a/docs/account/iam/service_principals.rst b/docs/account/iam/service_principals.rst index e0fd8577a..302cf5f79 100644 --- a/docs/account/iam/service_principals.rst +++ b/docs/account/iam/service_principals.rst @@ -23,10 +23,7 @@ a = AccountClient() - sp_create = a.service_principals.create(active=True, display_name=f"sdk-{time.time_ns()}") - - # cleanup - a.service_principals.delete(id=sp_create.id) + spn = a.service_principals.create(display_name=f"sdk-{time.time_ns()}") Create a service principal. diff --git a/docs/account/iam/workspace_assignment.rst b/docs/account/iam/workspace_assignment.rst index 6f4c66be5..745bd75da 100644 --- a/docs/account/iam/workspace_assignment.rst +++ b/docs/account/iam/workspace_assignment.rst @@ -47,9 +47,9 @@ a = AccountClient() - workspace_id = os.environ["TEST_WORKSPACE_ID"] + workspace_id = os.environ["DUMMY_WORKSPACE_ID"] - all = a.workspace_assignment.list(list=workspace_id) + all = a.workspace_assignment.list(workspace_id=workspace_id) Get permission assignments. @@ -80,9 +80,9 @@ spn_id = spn.id - workspace_id = os.environ["DUMMY_WORKSPACE_ID"] + workspace_id = os.environ["TEST_WORKSPACE_ID"] - _ = a.workspace_assignment.update( + a.workspace_assignment.update( workspace_id=workspace_id, principal_id=spn_id, permissions=[iam.WorkspacePermission.USER], diff --git a/docs/account/provisioning/credentials.rst b/docs/account/provisioning/credentials.rst index d023d4f1f..e307588f1 100644 --- a/docs/account/provisioning/credentials.rst +++ b/docs/account/provisioning/credentials.rst @@ -24,15 +24,15 @@ a = AccountClient() - role = a.credentials.create( + creds = a.credentials.create( credentials_name=f"sdk-{time.time_ns()}", aws_credentials=provisioning.CreateCredentialAwsCredentials( - sts_role=provisioning.CreateCredentialStsRole(role_arn=os.environ["TEST_CROSSACCOUNT_ARN"]) + sts_role=provisioning.CreateCredentialStsRole(role_arn=os.environ["TEST_LOGDELIVERY_ARN"]) ), ) # cleanup - a.credentials.delete(credentials_id=role.credentials_id) + a.credentials.delete(credentials_id=creds.credentials_id) Create credential configuration. diff --git a/docs/account/provisioning/storage.rst b/docs/account/provisioning/storage.rst index b8e144f8c..010795885 100644 --- a/docs/account/provisioning/storage.rst +++ b/docs/account/provisioning/storage.rst @@ -16,7 +16,6 @@ .. code-block:: - import os import time from databricks.sdk import AccountClient @@ -24,13 +23,13 @@ a = AccountClient() - storage = a.storage.create( + bucket = a.storage.create( storage_configuration_name=f"sdk-{time.time_ns()}", - root_bucket_info=provisioning.RootBucketInfo(bucket_name=os.environ["TEST_ROOT_BUCKET"]), + root_bucket_info=provisioning.RootBucketInfo(bucket_name=f"sdk-{time.time_ns()}"), ) # cleanup - a.storage.delete(storage_configuration_id=storage.storage_configuration_id) + a.storage.delete(storage_configuration_id=bucket.storage_configuration_id) Create new storage configuration. diff --git a/docs/dbdataclasses/billing.rst b/docs/dbdataclasses/billing.rst index 590fd693e..cb89f4719 100644 --- a/docs/dbdataclasses/billing.rst +++ b/docs/dbdataclasses/billing.rst @@ -97,6 +97,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: CustomPolicyTag + :members: + :undoc-members: + .. autoclass:: DeleteBudgetConfigurationResponse :members: :undoc-members: diff --git a/docs/dbdataclasses/catalog.rst b/docs/dbdataclasses/catalog.rst index e302b781c..25d382f9c 100644 --- a/docs/dbdataclasses/catalog.rst +++ b/docs/dbdataclasses/catalog.rst @@ -130,6 +130,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: DELTASHARING_CATALOG :value: "DELTASHARING_CATALOG" + .. py:attribute:: FOREIGN_CATALOG + :value: "FOREIGN_CATALOG" + .. py:attribute:: MANAGED_CATALOG :value: "MANAGED_CATALOG" @@ -1023,6 +1026,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: APPLY_TAG :value: "APPLY_TAG" + .. py:attribute:: BROWSE + :value: "BROWSE" + .. py:attribute:: CREATE :value: "CREATE" @@ -1365,10 +1371,6 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: VIEW :value: "VIEW" -.. autoclass:: TagKeyValue - :members: - :undoc-members: - .. autoclass:: TemporaryCredentials :members: :undoc-members: @@ -1548,6 +1550,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:class:: VolumeType + The type of the volume. An external volume is located in the specified external location. A managed volume is located in the default location which is specified by the parent schema, or the parent catalog, or the Metastore. [Learn more] + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external + .. py:attribute:: EXTERNAL :value: "EXTERNAL" diff --git a/docs/dbdataclasses/cleanrooms.rst b/docs/dbdataclasses/cleanrooms.rst index 85ec98250..1983106d1 100644 --- a/docs/dbdataclasses/cleanrooms.rst +++ b/docs/dbdataclasses/cleanrooms.rst @@ -121,14 +121,211 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: PROVISIONING :value: "PROVISIONING" +.. py:class:: CleanRoomTaskRunLifeCycleState + + Copied from elastic-spark-common/api/messages/runs.proto. Using the original definition to remove coupling with jobs API definition + + .. py:attribute:: BLOCKED + :value: "BLOCKED" + + .. py:attribute:: INTERNAL_ERROR + :value: "INTERNAL_ERROR" + + .. py:attribute:: PENDING + :value: "PENDING" + + .. py:attribute:: QUEUED + :value: "QUEUED" + + .. py:attribute:: RUNNING + :value: "RUNNING" + + .. py:attribute:: RUN_LIFE_CYCLE_STATE_UNSPECIFIED + :value: "RUN_LIFE_CYCLE_STATE_UNSPECIFIED" + + .. py:attribute:: SKIPPED + :value: "SKIPPED" + + .. py:attribute:: TERMINATED + :value: "TERMINATED" + + .. py:attribute:: TERMINATING + :value: "TERMINATING" + + .. py:attribute:: WAITING_FOR_RETRY + :value: "WAITING_FOR_RETRY" + +.. py:class:: CleanRoomTaskRunResultState + + Copied from elastic-spark-common/api/messages/runs.proto. Using the original definition to avoid cyclic dependency. + + .. py:attribute:: CANCELED + :value: "CANCELED" + + .. py:attribute:: DISABLED + :value: "DISABLED" + + .. py:attribute:: EVICTED + :value: "EVICTED" + + .. py:attribute:: EXCLUDED + :value: "EXCLUDED" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: MAXIMUM_CONCURRENT_RUNS_REACHED + :value: "MAXIMUM_CONCURRENT_RUNS_REACHED" + + .. py:attribute:: RUN_RESULT_STATE_UNSPECIFIED + :value: "RUN_RESULT_STATE_UNSPECIFIED" + + .. py:attribute:: SUCCESS + :value: "SUCCESS" + + .. py:attribute:: SUCCESS_WITH_FAILURES + :value: "SUCCESS_WITH_FAILURES" + + .. py:attribute:: TIMEDOUT + :value: "TIMEDOUT" + + .. py:attribute:: UPSTREAM_CANCELED + :value: "UPSTREAM_CANCELED" + + .. py:attribute:: UPSTREAM_EVICTED + :value: "UPSTREAM_EVICTED" + + .. py:attribute:: UPSTREAM_FAILED + :value: "UPSTREAM_FAILED" + +.. autoclass:: CleanRoomTaskRunState + :members: + :undoc-members: + .. autoclass:: CollaboratorJobRunInfo :members: :undoc-members: +.. autoclass:: ColumnInfo + :members: + :undoc-members: + +.. autoclass:: ColumnMask + :members: + :undoc-members: + +.. py:class:: ColumnTypeName + + .. py:attribute:: ARRAY + :value: "ARRAY" + + .. py:attribute:: BINARY + :value: "BINARY" + + .. py:attribute:: BOOLEAN + :value: "BOOLEAN" + + .. py:attribute:: BYTE + :value: "BYTE" + + .. py:attribute:: CHAR + :value: "CHAR" + + .. py:attribute:: DATE + :value: "DATE" + + .. py:attribute:: DECIMAL + :value: "DECIMAL" + + .. py:attribute:: DOUBLE + :value: "DOUBLE" + + .. py:attribute:: FLOAT + :value: "FLOAT" + + .. py:attribute:: INT + :value: "INT" + + .. py:attribute:: INTERVAL + :value: "INTERVAL" + + .. py:attribute:: LONG + :value: "LONG" + + .. py:attribute:: MAP + :value: "MAP" + + .. py:attribute:: NULL + :value: "NULL" + + .. py:attribute:: SHORT + :value: "SHORT" + + .. py:attribute:: STRING + :value: "STRING" + + .. py:attribute:: STRUCT + :value: "STRUCT" + + .. py:attribute:: TABLE_TYPE + :value: "TABLE_TYPE" + + .. py:attribute:: TIMESTAMP + :value: "TIMESTAMP" + + .. py:attribute:: TIMESTAMP_NTZ + :value: "TIMESTAMP_NTZ" + + .. py:attribute:: USER_DEFINED_TYPE + :value: "USER_DEFINED_TYPE" + + .. py:attribute:: VARIANT + :value: "VARIANT" + .. autoclass:: ComplianceSecurityProfile :members: :undoc-members: +.. py:class:: ComplianceStandard + + Compliance stardard for SHIELD customers + + .. py:attribute:: CANADA_PROTECTED_B + :value: "CANADA_PROTECTED_B" + + .. py:attribute:: CYBER_ESSENTIAL_PLUS + :value: "CYBER_ESSENTIAL_PLUS" + + .. py:attribute:: FEDRAMP_HIGH + :value: "FEDRAMP_HIGH" + + .. py:attribute:: FEDRAMP_IL5 + :value: "FEDRAMP_IL5" + + .. py:attribute:: FEDRAMP_MODERATE + :value: "FEDRAMP_MODERATE" + + .. py:attribute:: HIPAA + :value: "HIPAA" + + .. py:attribute:: HITRUST + :value: "HITRUST" + + .. py:attribute:: IRAP_PROTECTED + :value: "IRAP_PROTECTED" + + .. py:attribute:: ISMAP + :value: "ISMAP" + + .. py:attribute:: ITAR_EAR + :value: "ITAR_EAR" + + .. py:attribute:: NONE + :value: "NONE" + + .. py:attribute:: PCI_DSS + :value: "PCI_DSS" + .. autoclass:: CreateCleanRoomOutputCatalogResponse :members: :undoc-members: @@ -141,6 +338,83 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: EgressNetworkPolicy + :members: + :undoc-members: + +.. autoclass:: EgressNetworkPolicyInternetAccessPolicy + :members: + :undoc-members: + +.. autoclass:: EgressNetworkPolicyInternetAccessPolicyInternetDestination + :members: + :undoc-members: + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol + + The filtering protocol used by the DP. For private and public preview, SEG will only support TCP filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be set to TCP by default and hidden from the user. In the future, users may be able to select HTTP filtering (i.e. SNI based filtering, filtering by FQDN). + + .. py:attribute:: TCP + :value: "TCP" + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType + + .. py:attribute:: FQDN + :value: "FQDN" + +.. autoclass:: EgressNetworkPolicyInternetAccessPolicyLogOnlyMode + :members: + :undoc-members: + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType + + .. py:attribute:: ALL_SERVICES + :value: "ALL_SERVICES" + + .. py:attribute:: SELECTED_SERVICES + :value: "SELECTED_SERVICES" + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType + + The values should match the list of workloads used in networkconfig.proto + + .. py:attribute:: DBSQL + :value: "DBSQL" + + .. py:attribute:: ML_SERVING + :value: "ML_SERVING" + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyRestrictionMode + + At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can only access explicitly allowed internet and storage destinations, as well as UC connections and external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via private link. + + .. py:attribute:: FULL_ACCESS + :value: "FULL_ACCESS" + + .. py:attribute:: PRIVATE_ACCESS_ONLY + :value: "PRIVATE_ACCESS_ONLY" + + .. py:attribute:: RESTRICTED_ACCESS + :value: "RESTRICTED_ACCESS" + +.. autoclass:: EgressNetworkPolicyInternetAccessPolicyStorageDestination + :members: + :undoc-members: + +.. py:class:: EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType + + .. py:attribute:: AWS_S3 + :value: "AWS_S3" + + .. py:attribute:: AZURE_STORAGE + :value: "AZURE_STORAGE" + + .. py:attribute:: CLOUDFLARE_R2 + :value: "CLOUDFLARE_R2" + + .. py:attribute:: GOOGLE_CLOUD_STORAGE + :value: "GOOGLE_CLOUD_STORAGE" + .. autoclass:: ListCleanRoomAssetsResponse :members: :undoc-members: @@ -153,6 +427,22 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: Partition + :members: + :undoc-members: + +.. autoclass:: PartitionValue + :members: + :undoc-members: + +.. py:class:: PartitionValueOp + + .. py:attribute:: EQUAL + :value: "EQUAL" + + .. py:attribute:: LIKE + :value: "LIKE" + .. autoclass:: UpdateCleanRoomRequest :members: :undoc-members: diff --git a/docs/dbdataclasses/compute.rst b/docs/dbdataclasses/compute.rst index 2424cf4cf..387c8e23c 100644 --- a/docs/dbdataclasses/compute.rst +++ b/docs/dbdataclasses/compute.rst @@ -299,10 +299,6 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: -.. autoclass:: CustomPolicyTag - :members: - :undoc-members: - .. autoclass:: DataPlaneEventDetails :members: :undoc-members: @@ -475,10 +471,6 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: -.. autoclass:: Environment - :members: - :undoc-members: - .. autoclass:: EventDetails :members: :undoc-members: diff --git a/docs/dbdataclasses/dashboards.rst b/docs/dbdataclasses/dashboards.rst index 776aac603..7997c1796 100644 --- a/docs/dbdataclasses/dashboards.rst +++ b/docs/dbdataclasses/dashboards.rst @@ -4,6 +4,10 @@ Dashboards These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.dashboards`` module. .. py:currentmodule:: databricks.sdk.service.dashboards +.. autoclass:: BaseChunkInfo + :members: + :undoc-members: + .. autoclass:: CancelQueryExecutionResponse :members: :undoc-members: @@ -12,6 +16,71 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ColumnInfo + :members: + :undoc-members: + +.. py:class:: ColumnInfoTypeName + + The name of the base data type. This doesn't include details for complex types such as STRUCT, MAP or ARRAY. + + .. py:attribute:: ARRAY + :value: "ARRAY" + + .. py:attribute:: BINARY + :value: "BINARY" + + .. py:attribute:: BOOLEAN + :value: "BOOLEAN" + + .. py:attribute:: BYTE + :value: "BYTE" + + .. py:attribute:: CHAR + :value: "CHAR" + + .. py:attribute:: DATE + :value: "DATE" + + .. py:attribute:: DECIMAL + :value: "DECIMAL" + + .. py:attribute:: DOUBLE + :value: "DOUBLE" + + .. py:attribute:: FLOAT + :value: "FLOAT" + + .. py:attribute:: INT + :value: "INT" + + .. py:attribute:: INTERVAL + :value: "INTERVAL" + + .. py:attribute:: LONG + :value: "LONG" + + .. py:attribute:: MAP + :value: "MAP" + + .. py:attribute:: NULL + :value: "NULL" + + .. py:attribute:: SHORT + :value: "SHORT" + + .. py:attribute:: STRING + :value: "STRING" + + .. py:attribute:: STRUCT + :value: "STRUCT" + + .. py:attribute:: TIMESTAMP + :value: "TIMESTAMP" + + .. py:attribute:: USER_DEFINED_TYPE + :value: "USER_DEFINED_TYPE" + .. autoclass:: CronSchedule :members: :undoc-members: @@ -45,6 +114,21 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ExternalLink + :members: + :undoc-members: + +.. py:class:: Format + + .. py:attribute:: ARROW_STREAM + :value: "ARROW_STREAM" + + .. py:attribute:: CSV + :value: "CSV" + + .. py:attribute:: JSON_ARRAY + :value: "JSON_ARRAY" + .. autoclass:: GenieAttachment :members: :undoc-members: @@ -57,6 +141,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: GenieGenerateDownloadFullQueryResultResponse + :members: + :undoc-members: + +.. autoclass:: GenieGetDownloadFullQueryResultResponse + :members: + :undoc-members: + .. autoclass:: GenieGetMessageQueryResultResponse :members: :undoc-members: @@ -319,6 +411,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ResultData + :members: + :undoc-members: + +.. autoclass:: ResultManifest + :members: + :undoc-members: + +.. autoclass:: ResultSchema + :members: + :undoc-members: + .. autoclass:: Schedule :members: :undoc-members: @@ -331,6 +435,84 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: UNPAUSED :value: "UNPAUSED" +.. autoclass:: ServiceError + :members: + :undoc-members: + +.. py:class:: ServiceErrorCode + + .. py:attribute:: ABORTED + :value: "ABORTED" + + .. py:attribute:: ALREADY_EXISTS + :value: "ALREADY_EXISTS" + + .. py:attribute:: BAD_REQUEST + :value: "BAD_REQUEST" + + .. py:attribute:: CANCELLED + :value: "CANCELLED" + + .. py:attribute:: DEADLINE_EXCEEDED + :value: "DEADLINE_EXCEEDED" + + .. py:attribute:: INTERNAL_ERROR + :value: "INTERNAL_ERROR" + + .. py:attribute:: IO_ERROR + :value: "IO_ERROR" + + .. py:attribute:: NOT_FOUND + :value: "NOT_FOUND" + + .. py:attribute:: RESOURCE_EXHAUSTED + :value: "RESOURCE_EXHAUSTED" + + .. py:attribute:: SERVICE_UNDER_MAINTENANCE + :value: "SERVICE_UNDER_MAINTENANCE" + + .. py:attribute:: TEMPORARILY_UNAVAILABLE + :value: "TEMPORARILY_UNAVAILABLE" + + .. py:attribute:: UNAUTHENTICATED + :value: "UNAUTHENTICATED" + + .. py:attribute:: UNKNOWN + :value: "UNKNOWN" + + .. py:attribute:: WORKSPACE_TEMPORARILY_UNAVAILABLE + :value: "WORKSPACE_TEMPORARILY_UNAVAILABLE" + +.. autoclass:: StatementResponse + :members: + :undoc-members: + +.. py:class:: StatementState + + Statement execution state: - `PENDING`: waiting for warehouse - `RUNNING`: running - `SUCCEEDED`: execution was successful, result data available for fetch - `FAILED`: execution failed; reason for failure described in accomanying error message - `CANCELED`: user canceled; can come from explicit cancel call, or timeout with `on_wait_timeout=CANCEL` - `CLOSED`: execution successful, and statement closed; result no longer available for fetch + + .. py:attribute:: CANCELED + :value: "CANCELED" + + .. py:attribute:: CLOSED + :value: "CLOSED" + + .. py:attribute:: FAILED + :value: "FAILED" + + .. py:attribute:: PENDING + :value: "PENDING" + + .. py:attribute:: RUNNING + :value: "RUNNING" + + .. py:attribute:: SUCCEEDED + :value: "SUCCEEDED" + +.. autoclass:: StatementStatus + :members: + :undoc-members: + .. autoclass:: Subscriber :members: :undoc-members: diff --git a/docs/dbdataclasses/jobs.rst b/docs/dbdataclasses/jobs.rst index fa5af4189..9b8b496f2 100644 --- a/docs/dbdataclasses/jobs.rst +++ b/docs/dbdataclasses/jobs.rst @@ -4,6 +4,49 @@ Jobs These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.jobs`` module. .. py:currentmodule:: databricks.sdk.service.jobs +.. autoclass:: Adlsgen2Info + :members: + :undoc-members: + +.. autoclass:: AutoScale + :members: + :undoc-members: + +.. autoclass:: AwsAttributes + :members: + :undoc-members: + +.. py:class:: AwsAvailability + + Availability type used for all subsequent nodes past the `first_on_demand` ones. + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster. + + .. py:attribute:: ON_DEMAND + :value: "ON_DEMAND" + + .. py:attribute:: SPOT + :value: "SPOT" + + .. py:attribute:: SPOT_WITH_FALLBACK + :value: "SPOT_WITH_FALLBACK" + +.. autoclass:: AzureAttributes + :members: + :undoc-members: + +.. py:class:: AzureAvailability + + Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster. + + .. py:attribute:: ON_DEMAND_AZURE + :value: "ON_DEMAND_AZURE" + + .. py:attribute:: SPOT_AZURE + :value: "SPOT_AZURE" + + .. py:attribute:: SPOT_WITH_FALLBACK_AZURE + :value: "SPOT_WITH_FALLBACK_AZURE" + .. autoclass:: BaseJob :members: :undoc-members: @@ -117,10 +160,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ClientsTypes + :members: + :undoc-members: + .. autoclass:: ClusterInstance :members: :undoc-members: +.. autoclass:: ClusterLogConf + :members: + :undoc-members: + .. autoclass:: ClusterSpec :members: :undoc-members: @@ -180,6 +231,48 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. py:class:: DataSecurityMode + + Data security mode decides what data governance model to use when accessing data from a cluster. + The following modes can only be used when `kind = CLASSIC_PREVIEW`. * `DATA_SECURITY_MODE_AUTO`: Databricks will choose the most appropriate access mode depending on your compute configuration. * `DATA_SECURITY_MODE_STANDARD`: Alias for `USER_ISOLATION`. * `DATA_SECURITY_MODE_DEDICATED`: Alias for `SINGLE_USER`. + The following modes can be used regardless of `kind`. * `NONE`: No security isolation for multiple users sharing the cluster. Data governance features are not available in this mode. * `SINGLE_USER`: A secure cluster that can only be exclusively used by a single user specified in `single_user_name`. Most programming languages, cluster features and data governance features are available in this mode. * `USER_ISOLATION`: A secure cluster that can be shared by multiple users. Cluster users are fully isolated so that they cannot see each other's data and credentials. Most data governance features are supported in this mode. But programming languages and cluster features might be limited. + The following modes are deprecated starting with Databricks Runtime 15.0 and will be removed for future Databricks Runtime versions: + * `LEGACY_TABLE_ACL`: This mode is for users migrating from legacy Table ACL clusters. * `LEGACY_PASSTHROUGH`: This mode is for users migrating from legacy Passthrough on high concurrency clusters. * `LEGACY_SINGLE_USER`: This mode is for users migrating from legacy Passthrough on standard clusters. * `LEGACY_SINGLE_USER_STANDARD`: This mode provides a way that doesn’t have UC nor passthrough enabled. + + .. py:attribute:: DATA_SECURITY_MODE_AUTO + :value: "DATA_SECURITY_MODE_AUTO" + + .. py:attribute:: DATA_SECURITY_MODE_DEDICATED + :value: "DATA_SECURITY_MODE_DEDICATED" + + .. py:attribute:: DATA_SECURITY_MODE_STANDARD + :value: "DATA_SECURITY_MODE_STANDARD" + + .. py:attribute:: LEGACY_PASSTHROUGH + :value: "LEGACY_PASSTHROUGH" + + .. py:attribute:: LEGACY_SINGLE_USER + :value: "LEGACY_SINGLE_USER" + + .. py:attribute:: LEGACY_SINGLE_USER_STANDARD + :value: "LEGACY_SINGLE_USER_STANDARD" + + .. py:attribute:: LEGACY_TABLE_ACL + :value: "LEGACY_TABLE_ACL" + + .. py:attribute:: NONE + :value: "NONE" + + .. py:attribute:: SINGLE_USER + :value: "SINGLE_USER" + + .. py:attribute:: USER_ISOLATION + :value: "USER_ISOLATION" + +.. autoclass:: DbfsStorageInfo + :members: + :undoc-members: + .. autoclass:: DbtOutput :members: :undoc-members: @@ -204,6 +297,24 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: DockerBasicAuth + :members: + :undoc-members: + +.. autoclass:: DockerImage + :members: + :undoc-members: + +.. py:class:: EbsVolumeType + + All EBS volume types that Databricks supports. See https://aws.amazon.com/ebs/details/ for details. + + .. py:attribute:: GENERAL_PURPOSE_SSD + :value: "GENERAL_PURPOSE_SSD" + + .. py:attribute:: THROUGHPUT_OPTIMIZED_HDD + :value: "THROUGHPUT_OPTIMIZED_HDD" + .. autoclass:: EnforcePolicyComplianceForJobResponseJobClusterSettingsChange :members: :undoc-members: @@ -216,6 +327,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: Environment + :members: + :undoc-members: + .. autoclass:: ExportRunOutput :members: :undoc-members: @@ -248,6 +363,27 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: SINGLE_TASK :value: "SINGLE_TASK" +.. autoclass:: GcpAttributes + :members: + :undoc-members: + +.. py:class:: GcpAvailability + + This field determines whether the instance pool will contain preemptible VMs, on-demand VMs, or preemptible VMs with a fallback to on-demand VMs if the former is unavailable. + + .. py:attribute:: ON_DEMAND_GCP + :value: "ON_DEMAND_GCP" + + .. py:attribute:: PREEMPTIBLE_GCP + :value: "PREEMPTIBLE_GCP" + + .. py:attribute:: PREEMPTIBLE_WITH_FALLBACK_GCP + :value: "PREEMPTIBLE_WITH_FALLBACK_GCP" + +.. autoclass:: GcsStorageInfo + :members: + :undoc-members: + .. autoclass:: GenAiComputeTask :members: :undoc-members: @@ -294,6 +430,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: InitScriptInfo + :members: + :undoc-members: + .. autoclass:: Job :members: :undoc-members: @@ -411,6 +551,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: NOT_SYNCED :value: "NOT_SYNCED" +.. autoclass:: JobsClusterSpec + :members: + :undoc-members: + .. py:class:: JobsHealthMetric Specifies the health metric that is being evaluated for a particular health rule. @@ -446,6 +590,21 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. py:class:: Kind + + The kind of compute described by this compute specification. + Depending on `kind`, different validations and default values will be applied. + Clusters with `kind = CLASSIC_PREVIEW` support the following fields, whereas clusters with no specified `kind` do not. * [is_single_node](/api/workspace/clusters/create#is_single_node) * [use_ml_runtime](/api/workspace/clusters/create#use_ml_runtime) * [data_security_mode](/api/workspace/clusters/create#data_security_mode) set to `DATA_SECURITY_MODE_AUTO`, `DATA_SECURITY_MODE_DEDICATED`, or `DATA_SECURITY_MODE_STANDARD` + By using the [simple form], your clusters are automatically using `kind = CLASSIC_PREVIEW`. + [simple form]: https://docs.databricks.com/compute/simple-form.html + + .. py:attribute:: CLASSIC_PREVIEW + :value: "CLASSIC_PREVIEW" + +.. autoclass:: Library + :members: + :undoc-members: + .. autoclass:: ListJobComplianceForPolicyResponse :members: :undoc-members: @@ -458,6 +617,18 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: LocalFileInfo + :members: + :undoc-members: + +.. autoclass:: LogAnalyticsInfo + :members: + :undoc-members: + +.. autoclass:: MavenLibrary + :members: + :undoc-members: + .. autoclass:: NotebookOutput :members: :undoc-members: @@ -514,6 +685,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: PythonPyPiLibrary + :members: + :undoc-members: + .. autoclass:: PythonWheelTask :members: :undoc-members: @@ -539,6 +714,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: RCranLibrary + :members: + :undoc-members: + .. autoclass:: RepairHistoryItem :members: :undoc-members: @@ -776,6 +955,21 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: WORKFLOW_RUN :value: "WORKFLOW_RUN" +.. py:class:: RuntimeEngine + + .. py:attribute:: NULL + :value: "NULL" + + .. py:attribute:: PHOTON + :value: "PHOTON" + + .. py:attribute:: STANDARD + :value: "STANDARD" + +.. autoclass:: S3StorageInfo + :members: + :undoc-members: + .. py:class:: Source Optional location type of the SQL file. When set to `WORKSPACE`, the SQL file will be retrieved from the local Databricks workspace. When set to `GIT`, the SQL file will be retrieved from a Git repository defined in `git_source`. If the value is empty, the task will use `GIT` if `git_source` is defined and `WORKSPACE` otherwise. @@ -1075,6 +1269,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: DASHBOARDS :value: "DASHBOARDS" +.. autoclass:: VolumesStorageInfo + :members: + :undoc-members: + .. autoclass:: Webhook :members: :undoc-members: @@ -1082,3 +1280,11 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. autoclass:: WebhookNotifications :members: :undoc-members: + +.. autoclass:: WorkloadType + :members: + :undoc-members: + +.. autoclass:: WorkspaceStorageInfo + :members: + :undoc-members: diff --git a/docs/dbdataclasses/ml.rst b/docs/dbdataclasses/ml.rst index 860a4ffbc..6045bfea0 100644 --- a/docs/dbdataclasses/ml.rst +++ b/docs/dbdataclasses/ml.rst @@ -408,6 +408,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: LogOutputsRequest + :members: + :undoc-members: + .. autoclass:: LogParam :members: :undoc-members: @@ -428,6 +432,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ModelInput + :members: + :undoc-members: + .. autoclass:: ModelTag :members: :undoc-members: diff --git a/docs/dbdataclasses/pipelines.rst b/docs/dbdataclasses/pipelines.rst index 903cb52ff..9d2d85a0e 100644 --- a/docs/dbdataclasses/pipelines.rst +++ b/docs/dbdataclasses/pipelines.rst @@ -4,6 +4,49 @@ Delta Live Tables These dataclasses are used in the SDK to represent API requests and responses for services in the ``databricks.sdk.service.pipelines`` module. .. py:currentmodule:: databricks.sdk.service.pipelines +.. autoclass:: Adlsgen2Info + :members: + :undoc-members: + +.. autoclass:: AwsAttributes + :members: + :undoc-members: + +.. py:class:: AwsAvailability + + Availability type used for all subsequent nodes past the `first_on_demand` ones. + Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster. + + .. py:attribute:: ON_DEMAND + :value: "ON_DEMAND" + + .. py:attribute:: SPOT + :value: "SPOT" + + .. py:attribute:: SPOT_WITH_FALLBACK + :value: "SPOT_WITH_FALLBACK" + +.. autoclass:: AzureAttributes + :members: + :undoc-members: + +.. py:class:: AzureAvailability + + Availability type used for all subsequent nodes past the `first_on_demand` ones. Note: If `first_on_demand` is zero, this availability type will be used for the entire cluster. + + .. py:attribute:: ON_DEMAND_AZURE + :value: "ON_DEMAND_AZURE" + + .. py:attribute:: SPOT_AZURE + :value: "SPOT_AZURE" + + .. py:attribute:: SPOT_WITH_FALLBACK_AZURE + :value: "SPOT_WITH_FALLBACK_AZURE" + +.. autoclass:: ClusterLogConf + :members: + :undoc-members: + .. autoclass:: CreatePipeline :members: :undoc-members: @@ -45,6 +88,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: WEDNESDAY :value: "WEDNESDAY" +.. autoclass:: DbfsStorageInfo + :members: + :undoc-members: + .. autoclass:: DeletePipelineResponse :members: :undoc-members: @@ -56,6 +103,16 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: BUNDLE :value: "BUNDLE" +.. py:class:: EbsVolumeType + + All EBS volume types that Databricks supports. See https://aws.amazon.com/ebs/details/ for details. + + .. py:attribute:: GENERAL_PURPOSE_SSD + :value: "GENERAL_PURPOSE_SSD" + + .. py:attribute:: THROUGHPUT_OPTIMIZED_HDD + :value: "THROUGHPUT_OPTIMIZED_HDD" + .. autoclass:: EditPipeline :members: :undoc-members: @@ -84,6 +141,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: WARN :value: "WARN" +.. autoclass:: EventLogSpec + :members: + :undoc-members: + .. autoclass:: FileLibrary :members: :undoc-members: @@ -92,6 +153,27 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: GcpAttributes + :members: + :undoc-members: + +.. py:class:: GcpAvailability + + This field determines whether the instance pool will contain preemptible VMs, on-demand VMs, or preemptible VMs with a fallback to on-demand VMs if the former is unavailable. + + .. py:attribute:: ON_DEMAND_GCP + :value: "ON_DEMAND_GCP" + + .. py:attribute:: PREEMPTIBLE_GCP + :value: "PREEMPTIBLE_GCP" + + .. py:attribute:: PREEMPTIBLE_WITH_FALLBACK_GCP + :value: "PREEMPTIBLE_WITH_FALLBACK_GCP" + +.. autoclass:: GcsStorageInfo + :members: + :undoc-members: + .. autoclass:: GetPipelinePermissionLevelsResponse :members: :undoc-members: @@ -126,6 +208,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: InitScriptInfo + :members: + :undoc-members: + .. autoclass:: ListPipelineEventsResponse :members: :undoc-members: @@ -138,6 +224,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: LocalFileInfo + :members: + :undoc-members: + +.. autoclass:: LogAnalyticsInfo + :members: + :undoc-members: + .. autoclass:: ManualTrigger :members: :undoc-members: @@ -155,6 +249,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: STABLE :value: "STABLE" +.. autoclass:: MavenLibrary + :members: + :undoc-members: + .. autoclass:: NotebookLibrary :members: :undoc-members: @@ -302,6 +400,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: S3StorageInfo + :members: + :undoc-members: + .. autoclass:: SchemaSpec :members: :undoc-members: @@ -469,3 +571,11 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: WAITING_FOR_RESOURCES :value: "WAITING_FOR_RESOURCES" + +.. autoclass:: VolumesStorageInfo + :members: + :undoc-members: + +.. autoclass:: WorkspaceStorageInfo + :members: + :undoc-members: diff --git a/docs/dbdataclasses/serving.rst b/docs/dbdataclasses/serving.rst index 367f41b90..ed6c73ffb 100644 --- a/docs/dbdataclasses/serving.rst +++ b/docs/dbdataclasses/serving.rst @@ -406,6 +406,8 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:class:: ServedModelInputWorkloadType + Please keep this in sync with with workload types in InferenceEndpointEntities.scala + .. py:attribute:: CPU :value: "CPU" @@ -512,6 +514,8 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:class:: ServingModelWorkloadType + Please keep this in sync with with workload types in InferenceEndpointEntities.scala + .. py:attribute:: CPU :value: "CPU" diff --git a/docs/dbdataclasses/settings.rst b/docs/dbdataclasses/settings.rst index 2325c4023..7582ab9df 100644 --- a/docs/dbdataclasses/settings.rst +++ b/docs/dbdataclasses/settings.rst @@ -298,83 +298,6 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: -.. autoclass:: EgressNetworkPolicy - :members: - :undoc-members: - -.. autoclass:: EgressNetworkPolicyInternetAccessPolicy - :members: - :undoc-members: - -.. autoclass:: EgressNetworkPolicyInternetAccessPolicyInternetDestination - :members: - :undoc-members: - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationFilteringProtocol - - The filtering protocol used by the DP. For private and public preview, SEG will only support TCP filtering (i.e. DNS based filtering, filtering by destination IP address), so protocol will be set to TCP by default and hidden from the user. In the future, users may be able to select HTTP filtering (i.e. SNI based filtering, filtering by FQDN). - - .. py:attribute:: TCP - :value: "TCP" - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyInternetDestinationInternetDestinationType - - .. py:attribute:: FQDN - :value: "FQDN" - -.. autoclass:: EgressNetworkPolicyInternetAccessPolicyLogOnlyMode - :members: - :undoc-members: - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyLogOnlyModeLogOnlyModeType - - .. py:attribute:: ALL_SERVICES - :value: "ALL_SERVICES" - - .. py:attribute:: SELECTED_SERVICES - :value: "SELECTED_SERVICES" - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyLogOnlyModeWorkloadType - - The values should match the list of workloads used in networkconfig.proto - - .. py:attribute:: DBSQL - :value: "DBSQL" - - .. py:attribute:: ML_SERVING - :value: "ML_SERVING" - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyRestrictionMode - - At which level can Databricks and Databricks managed compute access Internet. FULL_ACCESS: Databricks can access Internet. No blocking rules will apply. RESTRICTED_ACCESS: Databricks can only access explicitly allowed internet and storage destinations, as well as UC connections and external locations. PRIVATE_ACCESS_ONLY (not used): Databricks can only access destinations via private link. - - .. py:attribute:: FULL_ACCESS - :value: "FULL_ACCESS" - - .. py:attribute:: PRIVATE_ACCESS_ONLY - :value: "PRIVATE_ACCESS_ONLY" - - .. py:attribute:: RESTRICTED_ACCESS - :value: "RESTRICTED_ACCESS" - -.. autoclass:: EgressNetworkPolicyInternetAccessPolicyStorageDestination - :members: - :undoc-members: - -.. py:class:: EgressNetworkPolicyInternetAccessPolicyStorageDestinationStorageDestinationType - - .. py:attribute:: AWS_S3 - :value: "AWS_S3" - - .. py:attribute:: AZURE_STORAGE - :value: "AZURE_STORAGE" - - .. py:attribute:: CLOUDFLARE_R2 - :value: "CLOUDFLARE_R2" - - .. py:attribute:: GOOGLE_CLOUD_STORAGE - :value: "GOOGLE_CLOUD_STORAGE" - .. autoclass:: EmailConfig :members: :undoc-members: @@ -669,6 +592,12 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: ARCLIGHT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY :value: "ARCLIGHT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY" + .. py:attribute:: ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN + :value: "ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN" + + .. py:attribute:: ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY + :value: "ARCLIGHT_MULTI_TENANT_AZURE_EXCHANGE_TOKEN_WITH_USER_DELEGATION_KEY" + .. py:attribute:: AZURE_ACTIVE_DIRECTORY_TOKEN :value: "AZURE_ACTIVE_DIRECTORY_TOKEN" diff --git a/docs/dbdataclasses/sharing.rst b/docs/dbdataclasses/sharing.rst index f72c59b21..9135a9f47 100644 --- a/docs/dbdataclasses/sharing.rst +++ b/docs/dbdataclasses/sharing.rst @@ -500,6 +500,10 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: VIEW :value: "VIEW" +.. autoclass:: TagKeyValue + :members: + :undoc-members: + .. autoclass:: UpdateProvider :members: :undoc-members: diff --git a/docs/dbdataclasses/sql.rst b/docs/dbdataclasses/sql.rst index c63fe7cd2..fe84a3b89 100644 --- a/docs/dbdataclasses/sql.rst +++ b/docs/dbdataclasses/sql.rst @@ -443,6 +443,14 @@ These dataclasses are used in the SDK to represent API requests and responses fo :members: :undoc-members: +.. autoclass:: ExternalQuerySource + :members: + :undoc-members: + +.. autoclass:: ExternalQuerySourceJobInfo + :members: + :undoc-members: + .. py:class:: Format .. py:attribute:: ARROW_STREAM @@ -1370,6 +1378,9 @@ These dataclasses are used in the SDK to represent API requests and responses fo .. py:attribute:: CAN_USE :value: "CAN_USE" + .. py:attribute:: CAN_VIEW + :value: "CAN_VIEW" + .. py:attribute:: IS_OWNER :value: "IS_OWNER" diff --git a/docs/workspace/catalog/catalogs.rst b/docs/workspace/catalog/catalogs.rst index 60959cad4..2505551cd 100644 --- a/docs/workspace/catalog/catalogs.rst +++ b/docs/workspace/catalog/catalogs.rst @@ -24,10 +24,10 @@ w = WorkspaceClient() - created = w.catalogs.create(name=f"sdk-{time.time_ns()}") + created_catalog = w.catalogs.create(name=f"sdk-{time.time_ns()}") # cleanup - w.catalogs.delete(name=created.name, force=True) + w.catalogs.delete(name=created_catalog.name, force=True) Create a catalog. diff --git a/docs/workspace/catalog/external_locations.rst b/docs/workspace/catalog/external_locations.rst index 980467306..5828bf245 100644 --- a/docs/workspace/catalog/external_locations.rst +++ b/docs/workspace/catalog/external_locations.rst @@ -30,22 +30,20 @@ w = WorkspaceClient() - storage_credential = w.storage_credentials.create( + credential = w.storage_credentials.create( name=f"sdk-{time.time_ns()}", aws_iam_role=catalog.AwsIamRoleRequest(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), - comment="created via SDK", ) - external_location = w.external_locations.create( + created = w.external_locations.create( name=f"sdk-{time.time_ns()}", - credential_name=storage_credential.name, - comment="created via SDK", - url="s3://" + os.environ["TEST_BUCKET"] + "/" + f"sdk-{time.time_ns()}", + credential_name=credential.name, + url="s3://%s/%s" % (os.environ["TEST_BUCKET"], f"sdk-{time.time_ns()}"), ) # cleanup - w.storage_credentials.delete(name=storage_credential.name) - w.external_locations.delete(name=external_location.name) + w.storage_credentials.delete(name=credential.name) + w.external_locations.delete(name=created.name) Create an external location. @@ -109,20 +107,20 @@ credential = w.storage_credentials.create( name=f"sdk-{time.time_ns()}", - aws_iam_role=catalog.AwsIamRole(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), + aws_iam_role=catalog.AwsIamRoleRequest(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), ) created = w.external_locations.create( name=f"sdk-{time.time_ns()}", credential_name=credential.name, - url=f's3://{os.environ["TEST_BUCKET"]}/sdk-{time.time_ns()}', + url="s3://%s/%s" % (os.environ["TEST_BUCKET"], f"sdk-{time.time_ns()}"), ) - _ = w.external_locations.get(get=created.name) + _ = w.external_locations.get(name=created.name) # cleanup - w.storage_credentials.delete(delete=credential.name) - w.external_locations.delete(delete=created.name) + w.storage_credentials.delete(name=credential.name) + w.external_locations.delete(name=created.name) Get an external location. @@ -146,11 +144,10 @@ .. code-block:: from databricks.sdk import WorkspaceClient - from databricks.sdk.service import catalog w = WorkspaceClient() - all = w.external_locations.list(catalog.ListExternalLocationsRequest()) + all = w.external_locations.list() List external locations. diff --git a/docs/workspace/catalog/storage_credentials.rst b/docs/workspace/catalog/storage_credentials.rst index ea2aece71..9a5ed0a46 100644 --- a/docs/workspace/catalog/storage_credentials.rst +++ b/docs/workspace/catalog/storage_credentials.rst @@ -30,13 +30,13 @@ w = WorkspaceClient() - created = w.storage_credentials.create( + credential = w.storage_credentials.create( name=f"sdk-{time.time_ns()}", - aws_iam_role=catalog.AwsIamRole(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), + aws_iam_role=catalog.AwsIamRoleRequest(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), ) # cleanup - w.storage_credentials.delete(delete=created.name) + w.storage_credentials.delete(name=credential.name) Create a storage credential. @@ -96,13 +96,13 @@ created = w.storage_credentials.create( name=f"sdk-{time.time_ns()}", - aws_iam_role=catalog.AwsIamRoleRequest(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), + aws_iam_role=catalog.AwsIamRole(role_arn=os.environ["TEST_METASTORE_DATA_ACCESS_ARN"]), ) - by_name = w.storage_credentials.get(name=created.name) + by_name = w.storage_credentials.get(get=created.name) # cleanup - w.storage_credentials.delete(name=created.name) + w.storage_credentials.delete(delete=created.name) Get a credential. @@ -123,10 +123,11 @@ .. code-block:: from databricks.sdk import WorkspaceClient + from databricks.sdk.service import catalog w = WorkspaceClient() - all = w.storage_credentials.list() + all = w.storage_credentials.list(catalog.ListStorageCredentialsRequest()) List credentials. diff --git a/docs/workspace/catalog/volumes.rst b/docs/workspace/catalog/volumes.rst index 700659cc5..5b6662f48 100644 --- a/docs/workspace/catalog/volumes.rst +++ b/docs/workspace/catalog/volumes.rst @@ -83,6 +83,11 @@ :param name: str The name of the volume :param volume_type: :class:`VolumeType` + The type of the volume. An external volume is located in the specified external location. A managed + volume is located in the default location which is specified by the parent schema, or the parent + catalog, or the Metastore. [Learn more] + + [Learn more]: https://docs.databricks.com/aws/en/volumes/managed-vs-external :param comment: str (optional) The comment attached to the volume :param storage_location: str (optional) diff --git a/docs/workspace/compute/clusters.rst b/docs/workspace/compute/clusters.rst index e4423bc98..ca496d3a0 100644 --- a/docs/workspace/compute/clusters.rst +++ b/docs/workspace/compute/clusters.rst @@ -713,11 +713,10 @@ .. code-block:: from databricks.sdk import WorkspaceClient - from databricks.sdk.service import compute w = WorkspaceClient() - all = w.clusters.list(compute.ListClustersRequest()) + nodes = w.clusters.list_node_types() List clusters. diff --git a/docs/workspace/dashboards/genie.rst b/docs/workspace/dashboards/genie.rst index eb92d299f..360517812 100644 --- a/docs/workspace/dashboards/genie.rst +++ b/docs/workspace/dashboards/genie.rst @@ -66,6 +66,57 @@ :returns: :class:`GenieGetMessageQueryResultResponse` + .. py:method:: generate_download_full_query_result(space_id: str, conversation_id: str, message_id: str, attachment_id: str) -> GenieGenerateDownloadFullQueryResultResponse + + Generate full query result download. + + Initiate full SQL query result download and obtain a transient ID for tracking the download progress. + This call initiates a new SQL execution to generate the query result. The result is stored in an + external link can be retrieved using the [Get Download Full Query + Result](:method:genie/getdownloadfullqueryresult) API. Warning: Databricks strongly recommends that + you protect the URLs that are returned by the `EXTERNAL_LINKS` disposition. See [Execute + Statement](:method:statementexecution/executestatement) for more details. + + :param space_id: str + Space ID + :param conversation_id: str + Conversation ID + :param message_id: str + Message ID + :param attachment_id: str + Attachment ID + + :returns: :class:`GenieGenerateDownloadFullQueryResultResponse` + + + .. py:method:: get_download_full_query_result(space_id: str, conversation_id: str, message_id: str, attachment_id: str, transient_statement_id: str) -> GenieGetDownloadFullQueryResultResponse + + Get download full query result status. + + Poll download progress and retrieve the SQL query result external link(s) upon completion. Warning: + Databricks strongly recommends that you protect the URLs that are returned by the `EXTERNAL_LINKS` + disposition. When you use the `EXTERNAL_LINKS` disposition, a short-lived, presigned URL is generated, + which can be used to download the results directly from Amazon S3. As a short-lived access credential + is embedded in this presigned URL, you should protect the URL. Because presigned URLs are already + generated with embedded temporary access credentials, you must not set an Authorization header in the + download requests. See [Execute Statement](:method:statementexecution/executestatement) for more + details. + + :param space_id: str + Space ID + :param conversation_id: str + Conversation ID + :param message_id: str + Message ID + :param attachment_id: str + Attachment ID + :param transient_statement_id: str + Transient Statement ID. This ID is provided by the [Start Download + endpoint](:method:genie/startdownloadfullqueryresult) + + :returns: :class:`GenieGetDownloadFullQueryResultResponse` + + .. py:method:: get_message(space_id: str, conversation_id: str, message_id: str) -> GenieMessage Get conversation message. diff --git a/docs/workspace/iam/current_user.rst b/docs/workspace/iam/current_user.rst index 1df3adf9f..bf739025c 100644 --- a/docs/workspace/iam/current_user.rst +++ b/docs/workspace/iam/current_user.rst @@ -17,7 +17,7 @@ w = WorkspaceClient() - me2 = w.current_user.me() + me = w.current_user.me() Get current user info. diff --git a/docs/workspace/iam/groups.rst b/docs/workspace/iam/groups.rst index 0b62b675a..fe0187cd6 100644 --- a/docs/workspace/iam/groups.rst +++ b/docs/workspace/iam/groups.rst @@ -71,6 +71,9 @@ group = w.groups.create(display_name=f"sdk-{time.time_ns()}") w.groups.delete(id=group.id) + + # cleanup + w.groups.delete(id=group.id) Delete a group. diff --git a/docs/workspace/iam/permissions.rst b/docs/workspace/iam/permissions.rst index 8d504eb37..6cd5b269b 100644 --- a/docs/workspace/iam/permissions.rst +++ b/docs/workspace/iam/permissions.rst @@ -71,7 +71,7 @@ obj = w.workspace.get_status(path=notebook_path) - levels = w.permissions.get_permission_levels(request_object_type="notebooks", request_object_id="%d" % (obj.object_id)) + _ = w.permissions.get(request_object_type="notebooks", request_object_id="%d" % (obj.object_id)) Get object permissions. diff --git a/docs/workspace/jobs/jobs.rst b/docs/workspace/jobs/jobs.rst index 923b05901..184ff9bd9 100644 --- a/docs/workspace/jobs/jobs.rst +++ b/docs/workspace/jobs/jobs.rst @@ -362,21 +362,23 @@ w.clusters.ensure_cluster_is_running(os.environ["DATABRICKS_CLUSTER_ID"]) and os.environ["DATABRICKS_CLUSTER_ID"] ) - run = w.jobs.submit( - run_name=f"sdk-{time.time_ns()}", + created_job = w.jobs.create( + name=f"sdk-{time.time_ns()}", tasks=[ - jobs.SubmitTask( + jobs.Task( + description="test", existing_cluster_id=cluster_id, notebook_task=jobs.NotebookTask(notebook_path=notebook_path), - task_key=f"sdk-{time.time_ns()}", + task_key="test", + timeout_seconds=0, ) ], - ).result() + ) - output = w.jobs.get_run_output(run_id=run.tasks[0].run_id) + by_id = w.jobs.get(job_id=created_job.job_id) # cleanup - w.jobs.delete_run(run_id=run.run_id) + w.jobs.delete(job_id=created_job.job_id) Get a single job. diff --git a/docs/workspace/ml/experiments.rst b/docs/workspace/ml/experiments.rst index a8863edcd..33d366301 100644 --- a/docs/workspace/ml/experiments.rst +++ b/docs/workspace/ml/experiments.rst @@ -280,10 +280,11 @@ API](/api/workspace/files/listdirectorycontents). :param page_token: str (optional) - Token indicating the page of artifact results to fetch. `page_token` is not supported when listing - artifacts in UC Volumes. A maximum of 1000 artifacts will be retrieved for UC Volumes. Please call - `/api/2.0/fs/directories{directory_path}` for listing artifacts in UC Volumes, which supports - pagination. See [List directory contents | Files API](/api/workspace/files/listdirectorycontents). + The token indicating the page of artifact results to fetch. `page_token` is not supported when + listing artifacts in UC Volumes. A maximum of 1000 artifacts will be retrieved for UC Volumes. + Please call `/api/2.0/fs/directories{directory_path}` for listing artifacts in UC Volumes, which + supports pagination. See [List directory contents | Files + API](/api/workspace/files/listdirectorycontents). :param path: str (optional) Filter artifacts matching this path (a relative path from the root artifact directory). :param run_id: str (optional) @@ -385,7 +386,7 @@ - .. py:method:: log_inputs(run_id: str [, datasets: Optional[List[DatasetInput]]]) + .. py:method:: log_inputs(run_id: str [, datasets: Optional[List[DatasetInput]], models: Optional[List[ModelInput]]]) Log inputs to a run. @@ -397,11 +398,13 @@ ID of the run to log under :param datasets: List[:class:`DatasetInput`] (optional) Dataset inputs + :param models: List[:class:`ModelInput`] (optional) + Model inputs - .. py:method:: log_metric(key: str, value: float, timestamp: int [, run_id: Optional[str], run_uuid: Optional[str], step: Optional[int]]) + .. py:method:: log_metric(key: str, value: float, timestamp: int [, dataset_digest: Optional[str], dataset_name: Optional[str], model_id: Optional[str], run_id: Optional[str], run_uuid: Optional[str], step: Optional[int]]) Log a metric for a run. @@ -415,6 +418,14 @@ Double value of the metric being logged. :param timestamp: int Unix timestamp in milliseconds at the time metric was logged. + :param dataset_digest: str (optional) + Dataset digest of the dataset associated with the metric, e.g. an md5 hash of the dataset that + uniquely identifies it within datasets of the same name. + :param dataset_name: str (optional) + The name of the dataset associated with the metric. E.g. “my.uc.table@2” “nyc-taxi-dataset”, + “fantastic-elk-3” + :param model_id: str (optional) + ID of the logged model associated with the metric, if applicable :param run_id: str (optional) ID of the run under which to log the metric. Must be provided. :param run_uuid: str (optional) diff --git a/docs/workspace/ml/forecasting.rst b/docs/workspace/ml/forecasting.rst index bb667b3fc..a8f5d2899 100644 --- a/docs/workspace/ml/forecasting.rst +++ b/docs/workspace/ml/forecasting.rst @@ -13,49 +13,45 @@ Creates a serverless forecasting experiment. Returns the experiment ID. :param train_data_path: str - The three-level (fully qualified) name of a unity catalog table. This table serves as the training - data for the forecasting model. + The fully qualified name of a Unity Catalog table, formatted as catalog_name.schema_name.table_name, + used as training data for the forecasting model. :param target_column: str - Name of the column in the input training table that serves as the prediction target. The values in - this column will be used as the ground truth for model training. + The column in the input training table used as the prediction target for model training. The values + in this column are used as the ground truth for model training. :param time_column: str - Name of the column in the input training table that represents the timestamp of each row. + The column in the input training table that represents each row's timestamp. :param forecast_granularity: str - The granularity of the forecast. This defines the time interval between consecutive rows in the time - series data. Possible values: '1 second', '1 minute', '5 minutes', '10 minutes', '15 minutes', '30 - minutes', 'Hourly', 'Daily', 'Weekly', 'Monthly', 'Quarterly', 'Yearly'. + The time interval between consecutive rows in the time series data. Possible values include: '1 + second', '1 minute', '5 minutes', '10 minutes', '15 minutes', '30 minutes', 'Hourly', 'Daily', + 'Weekly', 'Monthly', 'Quarterly', 'Yearly'. :param forecast_horizon: int - The number of time steps into the future for which predictions should be made. This value represents - a multiple of forecast_granularity determining how far ahead the model will forecast. + The number of time steps into the future to make predictions, calculated as a multiple of + forecast_granularity. This value represents how far ahead the model should forecast. :param custom_weights_column: str (optional) - Name of the column in the input training table used to customize the weight for each time series to - calculate weighted metrics. + The column in the training table used to customize weights for each time series. :param experiment_path: str (optional) - The path to the created experiment. This is the path where the experiment will be stored in the - workspace. + The path in the workspace to store the created experiment. :param holiday_regions: List[str] (optional) - Region code(s) to consider when automatically adding holiday features. When empty, no holiday - features are added. Only supports 1 holiday region for now. + The region code(s) to automatically add holiday features. Currently supports only one region. :param max_runtime: int (optional) - The maximum duration in minutes for which the experiment is allowed to run. If the experiment - exceeds this time limit it will be stopped automatically. + The maximum duration for the experiment in minutes. The experiment stops automatically if it exceeds + this limit. :param prediction_data_path: str (optional) - The three-level (fully qualified) path to a unity catalog table. This table path serves to store the - predictions. + The fully qualified path of a Unity Catalog table, formatted as catalog_name.schema_name.table_name, + used to store predictions. :param primary_metric: str (optional) The evaluation metric used to optimize the forecasting model. :param register_to: str (optional) - The three-level (fully qualified) path to a unity catalog model. This model path serves to store the - best model. + The fully qualified path of a Unity Catalog model, formatted as catalog_name.schema_name.model_name, + used to store the best model. :param split_column: str (optional) - Name of the column in the input training table used for custom data splits. The values in this - column must be "train", "validate", or "test" to indicate which split each row belongs to. + // The column in the training table used for custom data splits. Values must be 'train', 'validate', + or 'test'. :param timeseries_identifier_columns: List[str] (optional) - Name of the column in the input training table used to group the dataset to predict individual time - series + The column in the training table used to group the dataset for predicting individual time series. :param training_frameworks: List[str] (optional) - The list of frameworks to include for model tuning. Possible values: 'Prophet', 'ARIMA', 'DeepAR'. - An empty list will include all supported frameworks. + List of frameworks to include for model tuning. Possible values are 'Prophet', 'ARIMA', 'DeepAR'. An + empty list includes all supported frameworks. :returns: Long-running operation waiter for :class:`ForecastingExperiment`. diff --git a/docs/workspace/ml/model_registry.rst b/docs/workspace/ml/model_registry.rst index 0f6b49ac8..b56fba95f 100644 --- a/docs/workspace/ml/model_registry.rst +++ b/docs/workspace/ml/model_registry.rst @@ -94,7 +94,7 @@ w = WorkspaceClient() - model = w.model_registry.create_model(name=f"sdk-{time.time_ns()}") + created = w.model_registry.create_model(name=f"sdk-{time.time_ns()}") Create a model. @@ -127,7 +127,7 @@ model = w.model_registry.create_model(name=f"sdk-{time.time_ns()}") - mv = w.model_registry.create_model_version(name=model.registered_model.name, source="dbfs:/tmp") + created = w.model_registry.create_model_version(name=model.registered_model.name, source="dbfs:/tmp") Create a model version. @@ -773,14 +773,13 @@ w = WorkspaceClient() - model = w.model_registry.create_model(name=f"sdk-{time.time_ns()}") + created = w.model_registry.create_model(name=f"sdk-{time.time_ns()}") - created = w.model_registry.create_model_version(name=model.registered_model.name, source="dbfs:/tmp") + model = w.model_registry.get_model(name=created.registered_model.name) - w.model_registry.update_model_version( + w.model_registry.update_model( + name=model.registered_model_databricks.name, description=f"sdk-{time.time_ns()}", - name=created.model_version.name, - version=created.model_version.version, ) Update model. diff --git a/docs/workspace/pipelines/pipelines.rst b/docs/workspace/pipelines/pipelines.rst index 935724d82..71f0cdff4 100644 --- a/docs/workspace/pipelines/pipelines.rst +++ b/docs/workspace/pipelines/pipelines.rst @@ -15,7 +15,7 @@ also enforce data quality with Delta Live Tables expectations. Expectations allow you to define expected data quality and specify how to handle records that fail those expectations. - .. py:method:: create( [, allow_duplicate_names: Optional[bool], budget_policy_id: Optional[str], catalog: Optional[str], channel: Optional[str], clusters: Optional[List[PipelineCluster]], configuration: Optional[Dict[str, str]], continuous: Optional[bool], deployment: Optional[PipelineDeployment], development: Optional[bool], dry_run: Optional[bool], edition: Optional[str], filters: Optional[Filters], gateway_definition: Optional[IngestionGatewayPipelineDefinition], id: Optional[str], ingestion_definition: Optional[IngestionPipelineDefinition], libraries: Optional[List[PipelineLibrary]], name: Optional[str], notifications: Optional[List[Notifications]], photon: Optional[bool], restart_window: Optional[RestartWindow], run_as: Optional[RunAs], schema: Optional[str], serverless: Optional[bool], storage: Optional[str], target: Optional[str], trigger: Optional[PipelineTrigger]]) -> CreatePipelineResponse + .. py:method:: create( [, allow_duplicate_names: Optional[bool], budget_policy_id: Optional[str], catalog: Optional[str], channel: Optional[str], clusters: Optional[List[PipelineCluster]], configuration: Optional[Dict[str, str]], continuous: Optional[bool], deployment: Optional[PipelineDeployment], development: Optional[bool], dry_run: Optional[bool], edition: Optional[str], event_log: Optional[EventLogSpec], filters: Optional[Filters], gateway_definition: Optional[IngestionGatewayPipelineDefinition], id: Optional[str], ingestion_definition: Optional[IngestionPipelineDefinition], libraries: Optional[List[PipelineLibrary]], name: Optional[str], notifications: Optional[List[Notifications]], photon: Optional[bool], restart_window: Optional[RestartWindow], run_as: Optional[RunAs], schema: Optional[str], serverless: Optional[bool], storage: Optional[str], target: Optional[str], trigger: Optional[PipelineTrigger]]) -> CreatePipelineResponse Usage: @@ -79,6 +79,8 @@ :param dry_run: bool (optional) :param edition: str (optional) Pipeline product edition. + :param event_log: :class:`EventLogSpec` (optional) + Event log configuration for this pipeline :param filters: :class:`Filters` (optional) Filters on which Pipeline packages to include in the deployed graph. :param gateway_definition: :class:`IngestionGatewayPipelineDefinition` (optional) @@ -394,7 +396,7 @@ .. py:method:: stop_and_wait(pipeline_id: str, timeout: datetime.timedelta = 0:20:00) -> GetPipelineResponse - .. py:method:: update(pipeline_id: str [, allow_duplicate_names: Optional[bool], budget_policy_id: Optional[str], catalog: Optional[str], channel: Optional[str], clusters: Optional[List[PipelineCluster]], configuration: Optional[Dict[str, str]], continuous: Optional[bool], deployment: Optional[PipelineDeployment], development: Optional[bool], edition: Optional[str], expected_last_modified: Optional[int], filters: Optional[Filters], gateway_definition: Optional[IngestionGatewayPipelineDefinition], id: Optional[str], ingestion_definition: Optional[IngestionPipelineDefinition], libraries: Optional[List[PipelineLibrary]], name: Optional[str], notifications: Optional[List[Notifications]], photon: Optional[bool], restart_window: Optional[RestartWindow], run_as: Optional[RunAs], schema: Optional[str], serverless: Optional[bool], storage: Optional[str], target: Optional[str], trigger: Optional[PipelineTrigger]]) + .. py:method:: update(pipeline_id: str [, allow_duplicate_names: Optional[bool], budget_policy_id: Optional[str], catalog: Optional[str], channel: Optional[str], clusters: Optional[List[PipelineCluster]], configuration: Optional[Dict[str, str]], continuous: Optional[bool], deployment: Optional[PipelineDeployment], development: Optional[bool], edition: Optional[str], event_log: Optional[EventLogSpec], expected_last_modified: Optional[int], filters: Optional[Filters], gateway_definition: Optional[IngestionGatewayPipelineDefinition], id: Optional[str], ingestion_definition: Optional[IngestionPipelineDefinition], libraries: Optional[List[PipelineLibrary]], name: Optional[str], notifications: Optional[List[Notifications]], photon: Optional[bool], restart_window: Optional[RestartWindow], run_as: Optional[RunAs], schema: Optional[str], serverless: Optional[bool], storage: Optional[str], target: Optional[str], trigger: Optional[PipelineTrigger]]) Usage: @@ -474,6 +476,8 @@ Whether the pipeline is in Development mode. Defaults to false. :param edition: str (optional) Pipeline product edition. + :param event_log: :class:`EventLogSpec` (optional) + Event log configuration for this pipeline :param expected_last_modified: int (optional) If present, the last-modified time of the pipeline settings before the edit. If the settings were modified after that time, then the request will fail with a conflict. diff --git a/docs/workspace/sharing/providers.rst b/docs/workspace/sharing/providers.rst index 263545400..d78dd62a0 100644 --- a/docs/workspace/sharing/providers.rst +++ b/docs/workspace/sharing/providers.rst @@ -108,12 +108,25 @@ .. code-block:: + import time + from databricks.sdk import WorkspaceClient - from databricks.sdk.service import sharing w = WorkspaceClient() - all = w.providers.list(sharing.ListProvidersRequest()) + public_share_recipient = """{ + "shareCredentialsVersion":1, + "bearerToken":"dapiabcdefghijklmonpqrstuvwxyz", + "endpoint":"https://sharing.delta.io/delta-sharing/" + } + """ + + created = w.providers.create(name=f"sdk-{time.time_ns()}", recipient_profile_str=public_share_recipient) + + shares = w.providers.list_shares(name=created.name) + + # cleanup + w.providers.delete(name=created.name) List providers. diff --git a/docs/workspace/sql/queries.rst b/docs/workspace/sql/queries.rst index 4cc9b5b52..fe5442485 100644 --- a/docs/workspace/sql/queries.rst +++ b/docs/workspace/sql/queries.rst @@ -29,7 +29,7 @@ display_name=f"sdk-{time.time_ns()}", warehouse_id=srcs[0].warehouse_id, description="test query from Go SDK", - query_text="SHOW TABLES", + query_text="SELECT 1", ) ) diff --git a/docs/workspace/workspace/workspace.rst b/docs/workspace/workspace/workspace.rst index 3eae91432..abfc30860 100644 --- a/docs/workspace/workspace/workspace.rst +++ b/docs/workspace/workspace/workspace.rst @@ -188,7 +188,7 @@ content=base64.b64encode(("CREATE LIVE TABLE dlt_sample AS SELECT 1").encode()).decode(), format=workspace.ImportFormat.SOURCE, language=workspace.Language.SQL, - overwrite=True, + overwrite=true_, path=notebook_path, ) diff --git a/tests/conftest.py b/tests/conftest.py index d51d017b6..a0795acdb 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,9 +23,7 @@ def config(): @pytest.fixture def w(config): - from databricks.sdk import WorkspaceClient - - return WorkspaceClient(config=config) + return config __tests__ = os.path.dirname(__file__) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 55114bd84..62f57089f 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -7,8 +7,10 @@ import pytest -from databricks.sdk import AccountClient, FilesAPI, FilesExt, WorkspaceClient -from databricks.sdk.service.catalog import VolumeType +from databricks.sdk.catalog.v2.catalog import VolumeType +from databricks.sdk.databricks.config import Config +from databricks.sdk.files.v2.files import FilesAPI +from databricks.sdk.files.v2.mixin import FilesExt def pytest_addoption(parser): @@ -59,45 +61,45 @@ def inner(k=16) -> str: @pytest.fixture(scope="session") -def a(env_or_skip) -> AccountClient: +def a(env_or_skip) -> Config: _load_debug_env_if_runs_from_ide("account") env_or_skip("CLOUD_ENV") - account_client = AccountClient() - if not account_client.config.is_account_client: + cfg = Config() + if not cfg.is_account_client: pytest.skip("not Databricks Account client") - return account_client + return cfg @pytest.fixture(scope="session") -def ucacct(env_or_skip) -> AccountClient: +def ucacct(env_or_skip) -> Config: _load_debug_env_if_runs_from_ide("ucacct") env_or_skip("CLOUD_ENV") - account_client = AccountClient() - if not account_client.config.is_account_client: + cfg = Config() + if not cfg.is_account_client: pytest.skip("not Databricks Account client") if "TEST_METASTORE_ID" not in os.environ: pytest.skip("not in Unity Catalog Workspace test env") - return account_client + return cfg @pytest.fixture(scope="session") -def w(env_or_skip) -> WorkspaceClient: +def w(env_or_skip) -> Config: _load_debug_env_if_runs_from_ide("workspace") env_or_skip("CLOUD_ENV") if "DATABRICKS_ACCOUNT_ID" in os.environ: pytest.skip("Skipping workspace test on account level") - return WorkspaceClient() + return Config() @pytest.fixture(scope="session") -def ucws(env_or_skip) -> WorkspaceClient: +def ucws(env_or_skip) -> Config: _load_debug_env_if_runs_from_ide("ucws") env_or_skip("CLOUD_ENV") if "DATABRICKS_ACCOUNT_ID" in os.environ: pytest.skip("Skipping workspace test on account level") if "TEST_METASTORE_ID" not in os.environ: pytest.skip("not in Unity Catalog Workspace test env") - return WorkspaceClient() + return Config() @pytest.fixture(scope="session") @@ -113,36 +115,58 @@ def inner(var: str) -> str: @pytest.fixture(scope="session") def schema(ucws, random): - schema = ucws.schemas.create("dbfs-" + random(), "main") + from databricks.sdk.catalog.v2.client import SchemasClient + + sc = SchemasClient(config=ucws) + schema = sc.create("dbfs-" + random(), "main") yield schema - ucws.schemas.delete(schema.full_name) + sc.delete(schema.full_name) @pytest.fixture(scope="session") def volume(ucws, schema): - volume = ucws.volumes.create("main", schema.name, "dbfs-test", VolumeType.MANAGED) + from databricks.sdk.catalog.v2.client import VolumesClient + + vc = VolumesClient(config=ucws) + volume = vc.create("main", schema.name, "dbfs-test", VolumeType.MANAGED) yield "/Volumes/" + volume.full_name.replace(".", "/") - ucws.volumes.delete(volume.full_name) + vc.delete(volume.full_name) @pytest.fixture(scope="session", params=[False, True]) def files_api(request, ucws) -> FilesAPI: + if request.param: # ensure new Files API client is used for files of any size - ucws.config.multipart_upload_min_stream_size = 0 + ucws.multipart_upload_min_stream_size = 0 # enable new Files API client - return FilesExt(ucws.api_client, ucws.config) + from databricks.sdk.databricks.core import ApiClient + + client = ApiClient(cfg=ucws) + + return FilesExt(client, ucws) else: # use the default client - return ucws.files + from databricks.sdk.databricks.core import ApiClient + from databricks.sdk.files.v2.files import FilesAPI + + client = ApiClient(cfg=ucws) + + api = FilesAPI(api_client=client) + return api @pytest.fixture() def workspace_dir(w, random): - directory = f"/Users/{w.current_user.me().user_name}/dir-{random(12)}" - w.workspace.mkdirs(directory) + from databricks.sdk.iam.v2.client import CurrentUserClient + from databricks.sdk.workspace.v2.client import WorkspaceClient + + cuc = CurrentUserClient(config=w) + wc = WorkspaceClient(config=w) + directory = f"/Users/{cuc.me().user_name}/dir-{random(12)}" + wc.mkdirs(directory) yield directory - w.workspace.delete(directory, recursive=True) + wc.delete(directory, recursive=True) def _load_debug_env_if_runs_from_ide(key) -> bool: @@ -162,6 +186,7 @@ def _is_in_debug() -> bool: return os.path.basename(sys.argv[0]) in [ "_jb_pytest_runner.py", "testlauncher.py", + "run_pytest_script.py", ] diff --git a/tests/integration/test_auth.py b/tests/integration/test_auth.py index b6b0d6289..7b1d2d51f 100644 --- a/tests/integration/test_auth.py +++ b/tests/integration/test_auth.py @@ -10,8 +10,8 @@ import pytest -from databricks.sdk.service.compute import SparkVersion -from databricks.sdk.service.jobs import ViewType +from databricks.sdk.compute.v2 import compute +from databricks.sdk.jobs.v2 import jobs @pytest.fixture @@ -101,7 +101,7 @@ def fresh_wheel_file(tmp_path) -> Path: # ucws.clusters.permanent_delete(interactive_cluster.cluster_id) -def _get_lts_versions(w) -> typing.List[SparkVersion]: +def _get_lts_versions(w) -> typing.List[compute.SparkVersion]: v = w.clusters.spark_versions() lts_runtimes = [ x @@ -186,7 +186,7 @@ def _task_outputs(w, run): output = "" run_output = w.jobs.export_run(task_run.run_id) for view in run_output.views: - if view.type != ViewType.NOTEBOOK: + if view.type != jobs.ViewType.NOTEBOOK: continue for b64 in notebook_model_re.findall(view.content): url_encoded: bytes = base64.b64decode(b64) diff --git a/tests/integration/test_client.py b/tests/integration/test_client.py index 4e13d6854..a2999cfa3 100644 --- a/tests/integration/test_client.py +++ b/tests/integration/test_client.py @@ -1,19 +1,4 @@ -def test_get_workspace_client(ucacct, env_or_skip): - # Need to switch to ucacct - workspace_id = env_or_skip("TEST_WORKSPACE_ID") - ws = ucacct.workspaces.get(workspace_id) - w = ucacct.get_workspace_client(ws) - assert w.current_user.me().active - - -def test_get_workspace_id(ucws, env_or_skip): - ws_id = int(env_or_skip("THIS_WORKSPACE_ID")) - assert ucws.get_workspace_id() == ws_id - - -def test_creating_ws_client_from_ac_client_does_not_override_config(ucacct, env_or_skip): - ws_id = env_or_skip("TEST_WORKSPACE_ID") - ws = ucacct.workspaces.get(ws_id) - w = ucacct.get_workspace_client(ws) - me = w.current_user.me() - assert me.user_name is not None +# TODO: need to decide if we want get_workspace_id on every client +# def test_get_workspace_id(ucws, env_or_skip): +# ws_id = int(env_or_skip("THIS_WORKSPACE_ID")) +# assert ucws.get_workspace_id() == ws_id diff --git a/tests/integration/test_clusters.py b/tests/integration/test_clusters.py index 8b3d2f014..f42303bce 100644 --- a/tests/integration/test_clusters.py +++ b/tests/integration/test_clusters.py @@ -1,49 +1,61 @@ -import logging -from datetime import timedelta - import pytest +from databricks.sdk.compute.v2 import compute +from databricks.sdk.compute.v2.client import ClustersClient from databricks.sdk.databricks.core import DatabricksError -from databricks.sdk.service.compute import EventType + +# from databricks.sdk.service.compute import EventType def test_smallest_node_type(w): - node_type_id = w.clusters.select_node_type(local_disk=True) + cc = ClustersClient(config=w) + node_type_id = cc.select_node_type(local_disk=True) assert node_type_id is not None def test_latest_runtime(w): - spark_version = w.clusters.select_spark_version(long_term_support=True) + cc = ClustersClient(config=w) + spark_version = cc.select_spark_version(long_term_support=True) assert spark_version is not None def test_cluster_events(w, env_or_skip): cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID") count = 0 - for e in w.clusters.events(cluster_id, event_types=[EventType.STARTING, EventType.TERMINATING]): + cc = ClustersClient(config=w) + for e in cc.events(cluster_id, event_types=[compute.EventType.STARTING, compute.EventType.TERMINATING]): count += 1 assert count > 0 -def test_ensure_cluster_is_running(w, env_or_skip): - cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID") - w.clusters.ensure_cluster_is_running(cluster_id) +# TODO: Re-enable this test after adding waiters to the SDK +# def test_ensure_cluster_is_running(w, env_or_skip): +# cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID") +# cc = ClustersClient(config=w) +# cc.ensure_cluster_is_running(cluster_id) -def test_create_cluster(w, env_or_skip, random): - info = w.clusters.create( - cluster_name=f"databricks-sdk-py-{random(8)}", - spark_version=w.clusters.select_spark_version(long_term_support=True), - instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"), - autotermination_minutes=10, - num_workers=1, - ).result(timeout=timedelta(minutes=20)) - logging.info(f"Created: {info}") +# TODO: Re-enable this test after adding LRO support to the SDK +# def test_create_cluster(w, env_or_skip, random): +# from databricks.sdk.compute.v2.client import ClustersClient +# cc = ClustersClient(config=w) + +# info = cc.create( +# cluster_name=f"databricks-sdk-py-{random(8)}", +# spark_version=cc.select_spark_version(long_term_support=True), +# instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"), +# autotermination_minutes=10, +# num_workers=1, +# ).result(timeout=timedelta(minutes=20)) +# logging.info(f"Created: {info}") def test_error_unmarshall(w, random): + from databricks.sdk.compute.v2.client import ClustersClient + + cc = ClustersClient(config=w) with pytest.raises(DatabricksError) as exc_info: - w.clusters.get("123__non_existing__") + cc.get("123__non_existing__") err = exc_info.value assert "Cluster 123__non_existing__ does not exist" in str(err) assert "INVALID_PARAMETER_VALUE" == err.error_code diff --git a/tests/integration/test_commands.py b/tests/integration/test_commands.py index 407f9ba20..bc8f023c7 100644 --- a/tests/integration/test_commands.py +++ b/tests/integration/test_commands.py @@ -4,8 +4,11 @@ def test_error_unmarshall(w, random): + from databricks.sdk.compute.v2.client import CommandExecutionClient + + cec = CommandExecutionClient(config=w) with pytest.raises(DatabricksError) as exc_info: - w.command_execution.execute(cluster_id="__non_existing__") + cec.execute(cluster_id="__non_existing__") err = exc_info.value assert "requirement failed: missing contextId" in str(err) assert err.error_code is None diff --git a/tests/integration/test_data_plane.py b/tests/integration/test_data_plane.py index 2558cb0bb..c40553dbf 100644 --- a/tests/integration/test_data_plane.py +++ b/tests/integration/test_data_plane.py @@ -1,22 +1,20 @@ -from databricks.sdk.databricks.data_plane import DataPlaneTokenSource +# TODO: Re-enable this after adding data plane services to the SDK +# def test_data_plane_token_source(ucws, env_or_skip): +# endpoint = env_or_skip("SERVING_ENDPOINT_NAME") +# serving_endpoint = ucws.serving_endpoints.get(endpoint) +# assert serving_endpoint.data_plane_info is not None +# assert serving_endpoint.data_plane_info.query_info is not None +# info = serving_endpoint.data_plane_info.query_info -def test_data_plane_token_source(ucws, env_or_skip): - endpoint = env_or_skip("SERVING_ENDPOINT_NAME") - serving_endpoint = ucws.serving_endpoints.get(endpoint) - assert serving_endpoint.data_plane_info is not None - assert serving_endpoint.data_plane_info.query_info is not None +# ts = DataPlaneTokenSource(ucws.config.host, ucws._config.oauth_token) +# dp_token = ts.token(info.endpoint_url, info.authorization_details) - info = serving_endpoint.data_plane_info.query_info +# assert dp_token.valid - ts = DataPlaneTokenSource(ucws.config.host, ucws._config.oauth_token) - dp_token = ts.token(info.endpoint_url, info.authorization_details) - assert dp_token.valid - - -def test_model_serving_data_plane(ucws, env_or_skip): - endpoint = env_or_skip("SERVING_ENDPOINT_NAME") - serving_endpoints = ucws.serving_endpoints_data_plane - response = serving_endpoints.query(name=endpoint, dataframe_records=[{"col": 1.0}]) - assert response is not None +# def test_model_serving_data_plane(ucws, env_or_skip): +# endpoint = env_or_skip("SERVING_ENDPOINT_NAME") +# serving_endpoints = ucws.serving_endpoints_data_plane +# response = serving_endpoints.query(name=endpoint, dataframe_records=[{"col": 1.0}]) +# assert response is not None diff --git a/tests/integration/test_dbutils.py b/tests/integration/test_dbutils.py index 0bf5ab4f5..eb33c57ac 100644 --- a/tests/integration/test_dbutils.py +++ b/tests/integration/test_dbutils.py @@ -5,6 +5,7 @@ import pytest from databricks.sdk.databricks.core import DatabricksError +from databricks.sdk.databricks.dbutils import RemoteDbUtils from databricks.sdk.databricks.errors import NotFound @@ -16,21 +17,25 @@ def test_rest_dbfs_ls(w, env_or_skip): assert len(x) > 1 -def test_proxy_dbfs_mounts(w, env_or_skip): - w.config.cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID") +# TODO: Re-enable this test after adding waiters to the SDK +# def test_proxy_dbfs_mounts(w, env_or_skip): - x = w.dbutils.fs.mounts() +# w.cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID") - assert len(x) > 1 +# dbu = RemoteDbUtils(config=w) +# x = dbu.fs.mounts() + +# assert len(x) > 1 @pytest.fixture(params=["dbfs", "volumes"]) def fs_and_base_path(request, ucws, volume): + dbu = RemoteDbUtils(config=ucws) if request.param == "dbfs": - fs = ucws.dbutils.fs + fs = dbu.fs base_path = "/tmp" else: - fs = ucws.dbutils.fs + fs = dbu.fs base_path = volume return fs, base_path @@ -54,8 +59,9 @@ def test_large_put(fs_and_base_path): def test_put_local_path(w, random, tmp_path): to_write = random(1024 * 1024 * 2) tmp_path = tmp_path / "tmp_file" - w.dbutils.fs.put(f"file:{tmp_path}", to_write, True) - assert w.dbutils.fs.head(f"file:{tmp_path}", 1024 * 1024 * 2) == to_write + dbu = RemoteDbUtils(config=w) + dbu.fs.put(f"file:{tmp_path}", to_write, True) + assert dbu.fs.head(f"file:{tmp_path}", 1024 * 1024 * 2) == to_write def test_cp_file(fs_and_base_path, random): @@ -184,9 +190,12 @@ def test_secrets(w, random): logger = logging.getLogger("foo") logger.info(f"Before loading secret: {random_value}") - w.secrets.create_scope(random_scope) - w.secrets.put_secret(random_scope, key_for_string, string_value=random_value) - w.secrets.put_secret( + from databricks.sdk.workspace.v2.client import SecretsClient + + sc = SecretsClient(config=w) + sc.create_scope(random_scope) + sc.put_secret(random_scope, key_for_string, string_value=random_value) + sc.put_secret( random_scope, key_for_bytes, bytes_value=base64.b64encode(random_value.encode()).decode(), diff --git a/tests/integration/test_deployment.py b/tests/integration/test_deployment.py index 2071645d2..9ec6e1c32 100644 --- a/tests/integration/test_deployment.py +++ b/tests/integration/test_deployment.py @@ -4,7 +4,10 @@ def test_workspaces(a): - if a.config.is_azure: + from databricks.sdk.provisioning.v2.client import WorkspacesClient + + wc = WorkspacesClient(config=a) + if a.is_azure: pytest.skip("not available on Azure") - for w in a.workspaces.list(): + for w in wc.list(): logging.info(f"Found workspace: {w.workspace_name}") diff --git a/tests/integration/test_external_browser.py b/tests/integration/test_external_browser.py index 883069217..acbc33cff 100644 --- a/tests/integration/test_external_browser.py +++ b/tests/integration/test_external_browser.py @@ -1,7 +1,5 @@ import pytest -from databricks.sdk import WorkspaceClient - from .conftest import _load_debug_env_if_runs_from_ide @@ -13,34 +11,40 @@ def env(env_or_skip): def test_pkce_app(env): - w = WorkspaceClient( + from databricks.sdk.compute.v2.client import ClustersClient + + cc = ClustersClient( host=env("DATABRICKS_HOST"), client_id=env("TEST_PKCE_APP_CLIENT_ID"), auth_type="external-browser", ) - clusters = w.clusters.list() + clusters = cc.list() for cl in clusters: print(f" - {cl.cluster_name} is {cl.state}") def test_public_app(env): - w = WorkspaceClient( + from databricks.sdk.compute.v2.client import ClustersClient + + cc = ClustersClient( host=env("DATABRICKS_HOST"), client_id=env("TEST_PUBLIC_APP_CLIENT_ID"), auth_type="external-browser", ) - clusters = w.clusters.list() + clusters = cc.list() for cl in clusters: print(f" - {cl.cluster_name} is {cl.state}") def test_private_app(env): - w = WorkspaceClient( + from databricks.sdk.compute.v2.client import ClustersClient + + cc = ClustersClient( host=env("DATABRICKS_HOST"), client_id=env("TEST_PRIVATE_APP_CLIENT_ID"), client_secret=env("TEST_PRIVATE_APP_CLIENT_SECRET"), auth_type="external-browser", ) - clusters = w.clusters.list() + clusters = cc.list() for cl in clusters: print(f" - {cl.cluster_name} is {cl.state}") diff --git a/tests/integration/test_files.py b/tests/integration/test_files.py index 4b3b9a649..3e4ba16f2 100644 --- a/tests/integration/test_files.py +++ b/tests/integration/test_files.py @@ -7,8 +7,9 @@ import pytest +from databricks.sdk.catalog.v2.client import SchemasClient from databricks.sdk.databricks.core import DatabricksError -from databricks.sdk.service.catalog import VolumeType +from databricks.sdk.files.v2.client import DbfsClient def test_local_io(random): @@ -29,11 +30,12 @@ def test_local_io(random): def test_dbfs_io(w, random): dummy_file = f"/tmp/{random()}" to_write = random(1024 * 1024 * 1.5).encode() - with w.dbfs.open(dummy_file, write=True) as f: + dc = DbfsClient(config=w) + with dc.open(dummy_file, write=True) as f: written = f.write(to_write) assert len(to_write) == written - f = w.dbfs.open(dummy_file, read=True) + f = dc.open(dummy_file, read=True) from_dbfs = f.read() assert from_dbfs == to_write f.close() @@ -41,10 +43,11 @@ def test_dbfs_io(w, random): @pytest.fixture def junk(w, random): + dc = DbfsClient(config=w) def inner(path: str, size=256) -> bytes: to_write = random(size).encode() - with w.dbfs.open(path, write=True) as f: + with dc.open(path, write=True) as f: written = f.write(to_write) assert len(to_write) == written return to_write @@ -54,9 +57,10 @@ def inner(path: str, size=256) -> bytes: @pytest.fixture def ls(w): + dc = DbfsClient(config=w) def inner(root: str, recursive=False) -> List[str]: - return [f.path.removeprefix(root) for f in w.dbfs.list(root, recursive=recursive)] + return [f.path.removeprefix(root) for f in dc.list(root, recursive=recursive)] return inner @@ -70,7 +74,10 @@ def test_recursive_listing(w, random, junk, ls): assert ["/01", "/a"] == ls(root) assert ["/01", "/a/02", "/a/b/03"] == ls(root, recursive=True) - w.dbfs.delete(root, recursive=True) + from databricks.sdk.files.v2.client import DbfsClient + + dc = DbfsClient(config=w) + dc.delete(root, recursive=True) def test_cp_dbfs_folder_to_folder_non_recursive(w, random, junk, ls): @@ -80,7 +87,9 @@ def test_cp_dbfs_folder_to_folder_non_recursive(w, random, junk, ls): junk(f"{root}/a/b/03") new_root = f"/tmp/{random()}" - w.dbfs.copy(root, new_root) + dc = DbfsClient(config=w) + + dc.copy(root, new_root) assert ["/01"] == ls(new_root, recursive=True) @@ -92,7 +101,9 @@ def test_cp_dbfs_folder_to_folder_recursive(w, random, junk, ls): junk(f"{root}/a/b/03") new_root = f"/tmp/{random()}" - w.dbfs.copy(root, new_root, recursive=True, overwrite=True) + dc = DbfsClient(config=w) + + dc.copy(root, new_root, recursive=True, overwrite=True) assert ["/01", "/a/02", "/a/b/03"] == ls(new_root, recursive=True) @@ -104,8 +115,10 @@ def test_cp_dbfs_folder_to_existing_folder_recursive(w, random, junk, ls): junk(f"{root}/a/b/03") new_root = f"/tmp/{random()}" - w.dbfs.mkdirs(new_root) - w.dbfs.copy(root, new_root, recursive=True, overwrite=True) + dc = DbfsClient(config=w) + + dc.mkdirs(new_root) + dc.copy(root, new_root, recursive=True, overwrite=True) base = root.split("/")[-1] assert [f"/{base}/01", f"/{base}/a/02", f"/{base}/a/b/03"] == ls(new_root, recursive=True) @@ -116,19 +129,24 @@ def test_cp_dbfs_file_to_non_existing_location(w, random, junk): payload = junk(f"{root}/01") copy_destination = f"{root}/{random()}" - w.dbfs.copy(f"{root}/01", copy_destination) + dc = DbfsClient(config=w) - with w.dbfs.open(copy_destination, read=True) as f: + dc.copy(f"{root}/01", copy_destination) + + with dc.open(copy_destination, read=True) as f: assert f.read() == payload def test_cp_dbfs_file_to_existing_folder(w, random, junk): root = f"/tmp/{random()}" payload = junk(f"{root}/01") - w.dbfs.mkdirs(f"{root}/02") - w.dbfs.copy(f"{root}/01", f"{root}/02") - with w.dbfs.open(f"{root}/02/01", read=True) as f: + dc = DbfsClient(config=w) + + dc.mkdirs(f"{root}/02") + dc.copy(f"{root}/01", f"{root}/02") + + with dc.open(f"{root}/02/01", read=True) as f: assert f.read() == payload @@ -136,8 +154,11 @@ def test_cp_dbfs_file_to_existing_location(w, random, junk): root = f"/tmp/{random()}" junk(f"{root}/01") junk(f"{root}/02") + + dc = DbfsClient(config=w) + with pytest.raises(DatabricksError) as ei: - w.dbfs.copy(f"{root}/01", f"{root}/02") + dc.copy(f"{root}/01", f"{root}/02") assert "A file or directory already exists" in str(ei.value) @@ -146,9 +167,11 @@ def test_cp_dbfs_file_to_existing_location_with_overwrite(w, random, junk): payload = junk(f"{root}/01") junk(f"{root}/02") - w.dbfs.copy(f"{root}/01", f"{root}/02", overwrite=True) + dc = DbfsClient(config=w) - with w.dbfs.open(f"{root}/02", read=True) as f: + dc.copy(f"{root}/01", f"{root}/02", overwrite=True) + + with dc.open(f"{root}/02", read=True) as f: assert f.read() == payload @@ -156,10 +179,12 @@ def test_move_within_dbfs(w, random, junk): root = f"/tmp/{random()}" payload = junk(f"{root}/01") - w.dbfs.move_(f"{root}/01", f"{root}/02") + dc = DbfsClient(config=w) + + dc.move_(f"{root}/01", f"{root}/02") - assert w.dbfs.exists(f"{root}/01") is False - with w.dbfs.open(f"{root}/02", read=True) as f: + assert dc.exists(f"{root}/01") is False + with dc.open(f"{root}/02", read=True) as f: assert f.read() == payload @@ -169,9 +194,11 @@ def test_move_from_dbfs_to_local(w, random, junk, tmp_path): payload_02 = junk(f"{root}/a/02") payload_03 = junk(f"{root}/a/b/03") - w.dbfs.move_(root, f"file:{tmp_path}", recursive=True) + dc = DbfsClient(config=w) - assert w.dbfs.exists(root) is False + dc.move_(root, f"file:{tmp_path}", recursive=True) + + assert dc.exists(root) is False with (tmp_path / root.name / "01").open("rb") as f: assert f.read() == payload_01 with (tmp_path / root.name / "a/02").open("rb") as f: @@ -184,9 +211,11 @@ def test_dbfs_upload_download(w, random, junk, tmp_path): root = pathlib.Path(f"/tmp/{random()}") f = io.BytesIO(b"some text data") - w.dbfs.upload(f"{root}/01", f) + dc = DbfsClient(config=w) + + dc.upload(f"{root}/01", f) - with w.dbfs.download(f"{root}/01") as f: + with dc.download(f"{root}/01") as f: assert f.read() == b"some text data" @@ -204,18 +233,23 @@ def __exit__(self, exc_type, exc_val, exc_tb): @staticmethod def create_schema(w, catalog, schema): - res = w.schemas.create(catalog_name=catalog, name=schema) - return ResourceWithCleanup(lambda: w.schemas.delete(res.full_name)) + sc = SchemasClient(config=w) + res = sc.create(catalog_name=catalog, name=schema) + return ResourceWithCleanup(lambda: sc.delete(res.full_name)) @staticmethod def create_volume(w, catalog, schema, volume): - res = w.volumes.create( + from databricks.sdk.catalog.v2.catalog import VolumeType + from databricks.sdk.catalog.v2.client import VolumesClient + + vc = VolumesClient(config=w) + res = vc.create( catalog_name=catalog, schema_name=schema, name=volume, volume_type=VolumeType.MANAGED, ) - return ResourceWithCleanup(lambda: w.volumes.delete(res.full_name)) + return ResourceWithCleanup(lambda: vc.delete(res.full_name)) def test_files_api_upload_download(ucws, files_api, random): diff --git a/tests/integration/test_iam.py b/tests/integration/test_iam.py index 36c2a0b87..b290d31d3 100644 --- a/tests/integration/test_iam.py +++ b/tests/integration/test_iam.py @@ -1,68 +1,70 @@ import pytest from databricks.sdk.databricks import errors -from databricks.sdk.databricks.core import DatabricksError +from databricks.sdk.databricks.core import ApiClient, DatabricksError +from databricks.sdk.iam.v2.client import (AccountGroupsClient, + AccountServicePrincipalsClient, + AccountUsersClient, GroupsClient, + ServicePrincipalsClient, UsersClient) def test_filtering_groups(w, random): - all = w.groups.list(filter=f"displayName eq any-{random(12)}") + gc = GroupsClient(config=w) + + all = gc.list(filter=f"displayName eq any-{random(12)}") found = len(list(all)) assert found == 0 def test_scim_error_unmarshall(w, random): + gc = GroupsClient(config=w) + with pytest.raises(DatabricksError) as exc_info: - groups = w.groups.list(filter=random(12)) + groups = gc.list(filter=random(12)) next(groups) assert isinstance(exc_info.value, errors.BadRequest) def test_scim_get_user_as_dict(w): - first_user = next(w.users.list()) - user = w.users.get(first_user.id) + uc = UsersClient(config=w) + + first_user = next(uc.list()) + user = uc.get(first_user.id) # should not throw user.as_dict() @pytest.mark.parametrize( - "path,call", + "client_class,path,count", [ - ("/api/2.0/preview/scim/v2/Users", lambda w: w.users.list(count=10)), - ("/api/2.0/preview/scim/v2/Groups", lambda w: w.groups.list(count=4)), - ( - "/api/2.0/preview/scim/v2/ServicePrincipals", - lambda w: w.service_principals.list(count=1), - ), + (UsersClient, "/api/2.0/preview/scim/v2/Users", 10), + (GroupsClient, "/api/2.0/preview/scim/v2/Groups", 40), + (ServicePrincipalsClient, "/api/2.0/preview/scim/v2/ServicePrincipals", 10), ], ) -def test_workspace_users_list_pagination(w, path, call): - raw = w.api_client.do("GET", path) +def test_workspace_users_list_pagination(w, client_class, path, count): + client = client_class(config=w) + api_client = ApiClient(cfg=w) + raw = api_client.do("GET", path) total = raw["totalResults"] - all = call(w) + all = client.list(count=count) found = len(list(all)) assert found == total @pytest.mark.parametrize( - "path,call", + "client_class,path,count", [ - ( - "/api/2.0/accounts/%s/scim/v2/Users", - lambda a: a.users.list(count=3000), - ), - ( - "/api/2.0/accounts/%s/scim/v2/Groups", - lambda a: a.groups.list(count=5), - ), - ( - "/api/2.0/accounts/%s/scim/v2/ServicePrincipals", - lambda a: a.service_principals.list(count=1000), - ), + (AccountUsersClient, "/api/2.0/accounts/%s/scim/v2/Users", 3000), + (AccountGroupsClient, "/api/2.0/accounts/%s/scim/v2/Groups", 50), + (AccountServicePrincipalsClient, "/api/2.0/accounts/%s/scim/v2/ServicePrincipals", 1000), ], ) -def test_account_users_list_pagination(a, path, call): - raw = a.api_client.do("GET", path.replace("%s", a.config.account_id)) +def test_account_users_list_pagination(a, client_class, path, count): + client = client_class(config=a) + api_client = ApiClient(cfg=a) + raw = api_client.do("GET", path.replace("%s", a.account_id)) total = raw["totalResults"] - all = call(a) + all = client.list(count=count) found = len(list(all)) assert found == total diff --git a/tests/integration/test_jobs.py b/tests/integration/test_jobs.py index cfc8de0b7..d6e835898 100644 --- a/tests/integration/test_jobs.py +++ b/tests/integration/test_jobs.py @@ -1,47 +1,57 @@ -import datetime import logging +from databricks.sdk.compute.v2.client import ClustersClient +from databricks.sdk.jobs.v2.client import JobsClient + def test_jobs(w): found = 0 - for job in w.jobs.list(): + jc = JobsClient(config=w) + + for job in jc.list(): logging.info(f"Looking at {job.settings.name}") found += 1 assert found > 0 -def test_submitting_jobs(w, random, env_or_skip): - from databricks.sdk.service import compute, jobs - - py_on_dbfs = f"/home/{w.current_user.me().user_name}/sample.py" - with w.dbfs.open(py_on_dbfs, write=True, overwrite=True) as f: - f.write(b'import time; time.sleep(10); print("Hello, World!")') - - waiter = w.jobs.submit( - run_name=f"py-sdk-{random(8)}", - tasks=[ - jobs.SubmitTask( - task_key="pi", - new_cluster=compute.ClusterSpec( - spark_version=w.clusters.select_spark_version(long_term_support=True), - # node_type_id=w.clusters.select_node_type(local_disk=True), - instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"), - num_workers=1, - ), - spark_python_task=jobs.SparkPythonTask(python_file=f"dbfs:{py_on_dbfs}"), - ) - ], - ) +# TODO: Re-enable this after adding waiters to the SDK +# def test_submitting_jobs(w, random, env_or_skip): +# from databricks.sdk.jobs.v2 import jobs +# from databricks.sdk.compute.v2 import compute + +# cuc = CurrentUserClient(config=w) +# jc = JobsClient(config=w) +# dc = DbfsClient(config=w) - logging.info(f"starting to poll: {waiter.run_id}") +# py_on_dbfs = f"/home/{cuc.me().user_name}/sample.py" +# with dc.open(py_on_dbfs, write=True, overwrite=True) as f: +# f.write(b'import time; time.sleep(10); print("Hello, World!")') - def print_status(run: jobs.Run): - statuses = [f"{t.task_key}: {t.state.life_cycle_state}" for t in run.tasks] - logging.info(f'workflow intermediate status: {", ".join(statuses)}') +# waiter = jc.submit( +# run_name=f"py-sdk-{random(8)}", +# tasks=[ +# jobs.SubmitTask( +# task_key="pi", +# new_cluster=jobs.JobsClusterSpec( +# spark_version=w.clusters.select_spark_version(long_term_support=True), +# # node_type_id=w.clusters.select_node_type(local_disk=True), +# instance_pool_id=env_or_skip("TEST_INSTANCE_POOL_ID"), +# num_workers=1, +# ), +# spark_python_task=jobs.SparkPythonTask(python_file=f"dbfs:{py_on_dbfs}"), +# ) +# ], +# ) - run = waiter.result(timeout=datetime.timedelta(minutes=15), callback=print_status) +# logging.info(f"starting to poll: {waiter.run_id}") - logging.info(f"job finished: {run.run_page_url}") +# def print_status(run: jobs.Run): +# statuses = [f"{t.task_key}: {t.state.life_cycle_state}" for t in run.tasks] +# logging.info(f'workflow intermediate status: {", ".join(statuses)}') + +# run = waiter.result(timeout=datetime.timedelta(minutes=15), callback=print_status) + +# logging.info(f"job finished: {run.run_page_url}") def test_last_job_runs(w): @@ -52,9 +62,10 @@ def test_last_job_runs(w): all_jobs = {} durations = defaultdict(list) - for job in w.jobs.list(): + jc = JobsClient(config=w) + for job in jc.list(): all_jobs[job.job_id] = job - for run in w.jobs.list_runs(job_id=job.job_id, expand_tasks=False): + for run in jc.list_runs(job_id=job.job_id, expand_tasks=False): durations[job.job_id].append(run.run_duration) if job.job_id not in latest_state: latest_state[job.job_id] = run @@ -79,14 +90,17 @@ def test_last_job_runs(w): def test_create_job(w): - from databricks.sdk.service import compute, jobs + from databricks.sdk.jobs.v2 import jobs + + jc = JobsClient(config=w) + cc = ClustersClient(config=w) cluster = jobs.JobCluster( job_cluster_key="cluster1", - new_cluster=compute.ClusterSpec( + new_cluster=jobs.JobsClusterSpec( num_workers=2, - spark_version=w.clusters.select_spark_version(), - node_type_id=w.clusters.select_node_type(local_disk=True), + spark_version=cc.select_spark_version(), + node_type_id=cc.select_node_type(local_disk=True), ), ) @@ -96,4 +110,4 @@ def test_create_job(w): python_wheel_task=jobs.PythonWheelTask(entry_point="test", package_name="deepspeed"), ) - w.jobs.create(job_clusters=[cluster], tasks=[task1]) + jc.create(job_clusters=[cluster], tasks=[task1]) diff --git a/tests/integration/test_repos.py b/tests/integration/test_repos.py index 1682d9b87..9c50995c5 100644 --- a/tests/integration/test_repos.py +++ b/tests/integration/test_repos.py @@ -1,6 +1,9 @@ import logging +from databricks.sdk.workspace.v2.client import ReposClient + def test_repos_list(w): - for repo in w.repos.list(): + rc = ReposClient(config=w) + for repo in rc.list(): logging.info(f"Found repo: {repo}") diff --git a/tests/integration/test_settings.py b/tests/integration/test_settings.py index a14466815..8b7c8fb97 100644 --- a/tests/integration/test_settings.py +++ b/tests/integration/test_settings.py @@ -1,7 +1,11 @@ +from databricks.sdk.settings.v2.client import WorkspaceConfClient + + def test_workspace_conf(w): - w.workspace_conf.set_status({"enableResultsDownloading": "false"}) - conf = w.workspace_conf.get_status(keys="enableResultsDownloading") + wcc = WorkspaceConfClient(config=w) + wcc.set_status({"enableResultsDownloading": "false"}) + conf = wcc.get_status(keys="enableResultsDownloading") assert conf["enableResultsDownloading"] == "false" - w.workspace_conf.set_status({"enableResultsDownloading": "true"}) - conf = w.workspace_conf.get_status(keys="enableResultsDownloading") + wcc.set_status({"enableResultsDownloading": "true"}) + conf = wcc.get_status(keys="enableResultsDownloading") assert conf["enableResultsDownloading"] == "true" diff --git a/tests/integration/test_sql.py b/tests/integration/test_sql.py index 4ab775a0a..7a42d6a39 100644 --- a/tests/integration/test_sql.py +++ b/tests/integration/test_sql.py @@ -1,6 +1,7 @@ from datetime import datetime -from databricks.sdk.service.sql import QueryFilter, TimeRange +from databricks.sdk.sql.v2 import sql +from databricks.sdk.sql.v2.client import QueryHistoryClient def test_query_history_list_with_filter(w): @@ -8,12 +9,15 @@ def test_query_history_list_with_filter(w): def date_to_ms(date): return int(datetime.strptime(date, "%Y-%m-%d").timestamp() * 1000) - filter = QueryFilter( - query_start_time_range=TimeRange( + qhc = QueryHistoryClient(config=w) + + filter = sql.QueryFilter( + query_start_time_range=sql.TimeRange( start_time_ms=date_to_ms("2023-01-01"), end_time_ms=date_to_ms("2023-01-02"), ) ) - queries = w.query_history.list(filter_by=filter) + queries = qhc.list(filter_by=filter) + assert queries.res is not None, "Query history result list (queries.res) is None" for q in queries.res: print(q) diff --git a/tests/integration/test_workspace.py b/tests/integration/test_workspace.py index a3b4fd9c5..de6cb071d 100644 --- a/tests/integration/test_workspace.py +++ b/tests/integration/test_workspace.py @@ -1,98 +1,119 @@ import io -from databricks.sdk.service.workspace import ImportFormat, Language +from databricks.sdk.iam.v2.client import CurrentUserClient +from databricks.sdk.workspace.v2.client import WorkspaceClient +from databricks.sdk.workspace.v2.workspace import ImportFormat, Language def test_workspace_recursive_list(w, workspace_dir, random): + wc = WorkspaceClient(config=w) # create a file in the directory file = f"{workspace_dir}/file-{random(12)}.py" - w.workspace.upload(file, io.BytesIO(b"print(1)")) + wc.upload(file, io.BytesIO(b"print(1)")) # create a subdirectory subdirectory = f"{workspace_dir}/subdir-{random(12)}" - w.workspace.mkdirs(subdirectory) + wc.mkdirs(subdirectory) # create a file in the subdirectory subfile = f"{subdirectory}/subfile-{random(12)}.py" - w.workspace.upload(subfile, io.BytesIO(b"print(2)")) + wc.upload(subfile, io.BytesIO(b"print(2)")) # list the directory recursively names = [] - for i in w.workspace.list(workspace_dir, recursive=True): + for i in wc.list(workspace_dir, recursive=True): names.append(i.path) assert len(names) == 2 def test_workspace_upload_download_notebooks(w, random): - notebook = f"/Users/{w.current_user.me().user_name}/notebook-{random(12)}.py" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) - w.workspace.upload(notebook, io.BytesIO(b"print(1)")) - with w.workspace.download(notebook) as f: + notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py" + + wc.upload(notebook, io.BytesIO(b"print(1)")) + with wc.download(notebook) as f: content = f.read() assert content == b"# Databricks notebook source\nprint(1)" - w.workspace.delete(notebook) + wc.delete(notebook) def test_workspace_unzip_notebooks(w, random): - notebook = f"/Users/{w.current_user.me().user_name}/notebook-{random(12)}.py" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) + + notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py" # Big notebooks can be gzipped during transfer by the API (out of our control) # Creating some large content to trigger this behaviour notebook_content = ("print(1)\n" * 1000).strip("\n") - w.workspace.upload(notebook, io.BytesIO(bytes(notebook_content, "utf-8"))) - with w.workspace.download(notebook) as f: + wc.upload(notebook, io.BytesIO(bytes(notebook_content, "utf-8"))) + with wc.download(notebook) as f: content = f.read() expected_content = bytes(f"# Databricks notebook source\n{notebook_content}", "utf-8") assert content == expected_content - w.workspace.delete(notebook) + wc.delete(notebook) def test_workspace_download_connection_closed(w, random): - notebook = f"/Users/{w.current_user.me().user_name}/notebook-{random(12)}.py" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) + + notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py" - w.workspace.upload(notebook, io.BytesIO(b"print(1)")) + wc.upload(notebook, io.BytesIO(b"print(1)")) for n in range(30): - with w.workspace.download(notebook) as f: + with wc.download(notebook) as f: content = f.read() assert content == b"# Databricks notebook source\nprint(1)" - w.workspace.delete(notebook) + wc.delete(notebook) def test_workspace_upload_download_files(w, random): - py_file = f"/Users/{w.current_user.me().user_name}/file-{random(12)}.py" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) - w.workspace.upload(py_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO) - with w.workspace.download(py_file) as f: + py_file = f"/Users/{cuc.me().user_name}/file-{random(12)}.py" + + wc.upload(py_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO) + with wc.download(py_file) as f: content = f.read() assert content == b"print(1)" - w.workspace.delete(py_file) + wc.delete(py_file) def test_workspace_upload_download_txt_files(w, random): - txt_file = f"/Users/{w.current_user.me().user_name}/txt-{random(12)}.txt" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) + + txt_file = f"/Users/{cuc.me().user_name}/txt-{random(12)}.txt" - w.workspace.upload(txt_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO) - with w.workspace.download(txt_file) as f: + wc.upload(txt_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO) + with wc.download(txt_file) as f: content = f.read() assert content == b"print(1)" - w.workspace.delete(txt_file) + wc.delete(txt_file) def test_workspace_upload_download_notebooks_no_extension(w, random): - nb = f"/Users/{w.current_user.me().user_name}/notebook-{random(12)}" + wc = WorkspaceClient(config=w) + cuc = CurrentUserClient(config=w) + + nb = f"/Users/{cuc.me().user_name}/notebook-{random(12)}" - w.workspace.upload( + wc.upload( nb, io.BytesIO(b"print(1)"), format=ImportFormat.SOURCE, language=Language.PYTHON, ) - with w.workspace.download(nb) as f: + with wc.download(nb) as f: content = f.read() assert content == b"# Databricks notebook source\nprint(1)" - w.workspace.delete(nb) + wc.delete(nb) diff --git a/tests/test_client.py b/tests/test_client.py index 7eaf308f1..2a1457955 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,17 +1,13 @@ -from unittest.mock import create_autospec +# from databricks.sdk import WorkspaceClient -import pytest -from databricks.sdk import WorkspaceClient +# def test_autospec_fails_on_unknown_service(): +# w = create_autospec(WorkspaceClient) +# with pytest.raises(AttributeError): +# w.foo.bar() -def test_autospec_fails_on_unknown_service(): - w = create_autospec(WorkspaceClient) - with pytest.raises(AttributeError): - w.foo.bar() - - -def test_autospec_fails_on_setting_unknown_property(): - w = create_autospec(WorkspaceClient, spec_set=True) - with pytest.raises(AttributeError): - w.bar = 1 +# def test_autospec_fails_on_setting_unknown_property(): +# w = create_autospec(WorkspaceClient, spec_set=True) +# with pytest.raises(AttributeError): +# w.bar = 1 diff --git a/tests/test_compute_mixins.py b/tests/test_compute_mixins.py index ec895b022..5b637a249 100644 --- a/tests/test_compute_mixins.py +++ b/tests/test_compute_mixins.py @@ -1,6 +1,6 @@ import pytest -from databricks.sdk.mixins.compute import SemVer +from databricks.sdk.compute.v2.mixin import SemVer @pytest.mark.parametrize( diff --git a/tests/test_core.py b/tests/test_core.py index c8008df37..2db75fc4c 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -8,7 +8,6 @@ import pytest -from databricks.sdk import WorkspaceClient from databricks.sdk.databricks import errors, useragent from databricks.sdk.databricks.core import ApiClient, Config, DatabricksError from databricks.sdk.databricks.credentials_provider import ( @@ -18,8 +17,6 @@ AzureEnvironment, Cloud, DatabricksEnvironment) from databricks.sdk.databricks.oauth import Token -from databricks.sdk.service.catalog import PermissionsChange -from databricks.sdk.service.jobs import JobAccessControlRequest from databricks.sdk.version import __version__ from .conftest import noop_credentials @@ -287,8 +284,12 @@ def test_access_control_list(config, requests_mock): "http://localhost/api/2.2/jobs/create", request_headers={"User-Agent": config.user_agent}, ) - w = WorkspaceClient(config=config) - res = w.jobs.create(access_control_list=[JobAccessControlRequest(group_name="group")]) + + from databricks.sdk.jobs.v2.client import JobsClient + from databricks.sdk.jobs.v2.jobs import JobAccessControlRequest + + jc = JobsClient(config=config) + res = jc.create(access_control_list=[JobAccessControlRequest(group_name="group")]) assert requests_mock.call_count == 1 assert requests_mock.called @@ -300,9 +301,11 @@ def test_shares(config, requests_mock): "http://localhost/api/2.1/unity-catalog/shares/jobId/permissions", request_headers={"User-Agent": config.user_agent}, ) + from databricks.sdk.sharing.v2.client import SharesClient + from databricks.sdk.sharing.v2.sharing import PermissionsChange - w = WorkspaceClient(config=config) - res = w.shares.update_permissions(name="jobId", changes=[PermissionsChange(principal="principal")]) + sc = SharesClient(config=config) + res = sc.update_permissions(name="jobId", changes=[PermissionsChange(principal="principal")]) assert requests_mock.call_count == 1 assert requests_mock.called @@ -315,9 +318,10 @@ def test_deletes(config, requests_mock): request_headers={"User-Agent": config.user_agent}, text="null", ) + from databricks.sdk.sql.v2.client import AlertsClient - w = WorkspaceClient(config=config) - res = w.alerts.delete(id="alertId") + ac = AlertsClient(config=config) + res = ac.delete(id="alertId") assert requests_mock.call_count == 1 assert requests_mock.called diff --git a/tests/test_dbfs_mixins.py b/tests/test_dbfs_mixins.py index 7184af843..b8fbea0b6 100644 --- a/tests/test_dbfs_mixins.py +++ b/tests/test_dbfs_mixins.py @@ -1,16 +1,16 @@ import pytest from databricks.sdk.databricks.errors import NotFound -from databricks.sdk.mixins.files import (DbfsExt, _DbfsPath, _LocalPath, - _VolumesPath) +from databricks.sdk.files.v2.mixin import (DbfsExt, _DbfsPath, _LocalPath, + _VolumesPath) def test_moving_dbfs_file_to_local_dir(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import FileInfo, ReadResponse + from databricks.sdk.files.v2.client import DbfsClient + from databricks.sdk.files.v2.files import FileInfo, ReadResponse get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", return_value=FileInfo(path="a", is_dir=False, file_size=4), ) @@ -21,11 +21,11 @@ def fake_read(path: str, *, length: int = None, offset: int = None): return ReadResponse(bytes_read=4, data="aGVsbG8=") return ReadResponse(bytes_read=0) - mocker.patch("databricks.sdk.service.files.DbfsAPI.read", wraps=fake_read) - delete = mocker.patch("databricks.sdk.service.files.DbfsAPI.delete") + mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.read", wraps=fake_read) + delete = mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.delete") - w = WorkspaceClient(config=config) - w.dbfs.move_("a", f"file:{tmp_path}", recursive=True) + jc = DbfsClient(config=config) + jc.move_("a", f"file:{tmp_path}", recursive=True) get_status.assert_called_with("a") delete.assert_called_with("a", recursive=True) @@ -35,26 +35,26 @@ def fake_read(path: str, *, length: int = None, offset: int = None): def test_moving_local_dir_to_dbfs(config, tmp_path, mocker): - from databricks.sdk import WorkspaceClient - from databricks.sdk.service.files import CreateResponse + from databricks.sdk.files.v2.client import DbfsClient + from databricks.sdk.files.v2.files import CreateResponse with (tmp_path / "a").open("wb") as f: f.write(b"hello") mocker.patch( - "databricks.sdk.service.files.DbfsAPI.create", + "databricks.sdk.files.v2.files.DbfsAPI.create", return_value=CreateResponse(123), ) get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", side_effect=NotFound(), ) - add_block = mocker.patch("databricks.sdk.service.files.DbfsAPI.add_block") - close = mocker.patch("databricks.sdk.service.files.DbfsAPI.close") + add_block = mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.add_block") + close = mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.close") - w = WorkspaceClient(config=config) - w.dbfs.move_(f"file:{tmp_path}", "a", recursive=True) + dc = DbfsClient(config=config) + dc.move_(f"file:{tmp_path}", "a", recursive=True) get_status.assert_called_with("a") close.assert_called_with(123) @@ -86,33 +86,33 @@ def test_fs_path_invalid(config): def test_dbfs_local_path_mkdir(config, tmp_path): - from databricks.sdk import WorkspaceClient + from databricks.sdk.files.v2.client import DbfsClient - w = WorkspaceClient(config=config) - w.dbfs._path(f"file:{tmp_path}/test_dir").mkdir() - assert w.dbfs.exists(f"file:{tmp_path}/test_dir") + dc = DbfsClient(config=config) + dc._path(f"file:{tmp_path}/test_dir").mkdir() + assert dc.exists(f"file:{tmp_path}/test_dir") def test_dbfs_exists(config, mocker): - from databricks.sdk import WorkspaceClient + from databricks.sdk.files.v2.client import DbfsClient get_status = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", side_effect=NotFound(), ) - client = WorkspaceClient(config=config) - client.dbfs.exists("/abc/def/ghi") + client = DbfsClient(config=config) + client.exists("/abc/def/ghi") get_status.assert_called_with("/abc/def/ghi") def test_volume_exists(config, mocker): - from databricks.sdk import WorkspaceClient + from databricks.sdk.files.v2.client import DbfsClient - get_metadata = mocker.patch("databricks.sdk.service.files.FilesAPI.get_metadata") + get_metadata = mocker.patch("databricks.sdk.files.v2.files.FilesAPI.get_metadata") - client = WorkspaceClient(config=config) - client.dbfs.exists("/Volumes/abc/def/ghi") + client = DbfsClient(config=config) + client.exists("/Volumes/abc/def/ghi") get_metadata.assert_called_with("/Volumes/abc/def/ghi") diff --git a/tests/test_dbutils.py b/tests/test_dbutils.py index 7c0933b2d..53faad3d4 100644 --- a/tests/test_dbutils.py +++ b/tests/test_dbutils.py @@ -1,7 +1,7 @@ import pytest as pytest from databricks.sdk.databricks.dbutils import FileInfo as DBUtilsFileInfo -from databricks.sdk.service.files import FileInfo, ReadResponse +from databricks.sdk.files.v2.files import FileInfo, ReadResponse from .conftest import raises @@ -14,7 +14,7 @@ def dbutils(config): def test_fs_cp(dbutils, mocker): - inner = mocker.patch("databricks.sdk.mixins.files.DbfsExt.copy") + inner = mocker.patch("databricks.sdk.files.v2.mixin.DbfsExt.copy") dbutils.fs.cp("a", "b", recurse=True) @@ -23,11 +23,11 @@ def test_fs_cp(dbutils, mocker): def test_fs_head(dbutils, mocker): inner = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.read", + "databricks.sdk.files.v2.files.DbfsAPI.read", return_value=ReadResponse(data="aGVsbG8=", bytes_read=5), ) inner2 = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", return_value=FileInfo(path="a", is_dir=False, file_size=5), ) @@ -40,14 +40,14 @@ def test_fs_head(dbutils, mocker): def test_fs_ls(dbutils, mocker): inner = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.list", + "databricks.sdk.files.v2.files.DbfsAPI.list", return_value=[ FileInfo(path="a/b", file_size=10, modification_time=20), FileInfo(path="a/c", file_size=30, modification_time=40), ], ) inner2 = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", side_effect=[ FileInfo(path="a", is_dir=True, file_size=5), FileInfo(path="a/b", is_dir=False, file_size=5), @@ -64,7 +64,7 @@ def test_fs_ls(dbutils, mocker): def test_fs_mkdirs(dbutils, mocker): - inner = mocker.patch("databricks.sdk.service.files.DbfsAPI.mkdirs") + inner = mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.mkdirs") dbutils.fs.mkdirs("a") @@ -72,7 +72,7 @@ def test_fs_mkdirs(dbutils, mocker): def test_fs_mv(dbutils, mocker): - inner = mocker.patch("databricks.sdk.mixins.files.DbfsExt.move_") + inner = mocker.patch("databricks.sdk.files.v2.mixin.DbfsExt.move_") dbutils.fs.mv("a", "b") @@ -94,7 +94,7 @@ def write(self, contents): self._written = contents mock_open = _MockOpen() - inner = mocker.patch("databricks.sdk.mixins.files.DbfsExt.open", return_value=mock_open) + inner = mocker.patch("databricks.sdk.files.v2.mixin.DbfsExt.open", return_value=mock_open) dbutils.fs.put("a", "b") @@ -103,9 +103,9 @@ def write(self, contents): def test_fs_rm(dbutils, mocker): - inner = mocker.patch("databricks.sdk.service.files.DbfsAPI.delete") + inner = mocker.patch("databricks.sdk.files.v2.files.DbfsAPI.delete") inner2 = mocker.patch( - "databricks.sdk.service.files.DbfsAPI.get_status", + "databricks.sdk.files.v2.files.DbfsAPI.get_status", return_value=FileInfo(path="a", is_dir=False, file_size=5), ) @@ -119,135 +119,136 @@ def test_fs_mount_without_cluster_fails(dbutils): dbutils.fs.mount("s3://foo", "bar") -@pytest.fixture -def dbutils_proxy(mocker): - from databricks.sdk.databricks.core import Config - from databricks.sdk.databricks.dbutils import RemoteDbUtils - from databricks.sdk.service._internal import Wait - from databricks.sdk.service.compute import (ClusterDetails, CommandStatus, - CommandStatusResponse, Created, - Language, Results, State) - - from .conftest import noop_credentials - - cluster_get = mocker.patch( - "databricks.sdk.service.compute.ClustersAPI.get", - return_value=ClusterDetails(state=State.RUNNING), - ) - context_create = mocker.patch( - "databricks.sdk.service.compute.CommandExecutionAPI.create", - return_value=Wait(lambda **kwargs: Created("y")), - ) - - def inner(results_data: any, expect_command: str): - import json - - command_execute = mocker.patch( - "databricks.sdk.service.compute.CommandExecutionAPI.execute", - return_value=Wait( - lambda **kwargs: CommandStatusResponse( - results=Results(data=json.dumps(results_data)), - status=CommandStatus.FINISHED, - ) - ), - ) - - def assertions(): - cluster_get.assert_called_with("x") - context_create.assert_called_with(cluster_id="x", language=Language.PYTHON) - command_execute.assert_called_with( - cluster_id="x", - context_id="y", - language=Language.PYTHON, - command=expect_command, - ) - - dbutils = RemoteDbUtils( - Config( - host="http://localhost", - cluster_id="x", - credentials_strategy=noop_credentials, - ) - ) - return dbutils, assertions - - return inner - - -def test_fs_mount(dbutils_proxy): - command = ( - "\n" - " import json\n" - ' (args, kwargs) = json.loads(\'[["s3://foo", "bar"], {}]\')\n' - " result = dbutils.fs.mount(*args, **kwargs)\n" - " dbutils.notebook.exit(json.dumps(result))\n" - " " - ) - dbutils, assertions = dbutils_proxy({}, command) - - dbutils.fs.mount("s3://foo", "bar") - - assertions() - - -def test_fs_update_mount(dbutils_proxy): - command = ( - "\n" - " import json\n" - ' (args, kwargs) = json.loads(\'[["s3://foo2", "bar"], {}]\')\n' - " result = dbutils.fs.updateMount(*args, **kwargs)\n" - " dbutils.notebook.exit(json.dumps(result))\n" - " " - ) - dbutils, assertions = dbutils_proxy({}, command) - - dbutils.fs.updateMount("s3://foo2", "bar") - - assertions() - - -def test_fs_mounts(dbutils_proxy): - command = ( - "\n" - " import json\n" - " (args, kwargs) = json.loads('[[], {}]')\n" - " result = dbutils.fs.mounts(*args, **kwargs)\n" - " dbutils.notebook.exit(json.dumps(result))\n" - " " - ) - dbutils, assertions = dbutils_proxy( - [ - ("a", "b", "c"), - ("d", "e", "f"), - ], - command, - ) - - mounts = dbutils.fs.mounts() - - assert len(mounts) == 2 - assert mounts[0].mountPoint == "a" - assert mounts[0].source == "b" - - assertions() - - -def test_any_proxy(dbutils_proxy): - command = ( - "\n" - " import json\n" - " (args, kwargs) = json.loads('[[\"a\"], {}]')\n" - " result = dbutils.notebook.exit(*args, **kwargs)\n" - " dbutils.notebook.exit(json.dumps(result))\n" - " " - ) - dbutils, assertions = dbutils_proxy("a", command) - - param = dbutils.notebook.exit("a") - - assert param == "a" - - assertions() +# TODO: Re-enable this test after adding waiters to the SDK +# @pytest.fixture +# def dbutils_proxy(mocker): +# from databricks.sdk.databricks.core import Config +# from databricks.sdk.databricks.dbutils import RemoteDbUtils +# from databricks.sdk.service._internal import Wait +# from databricks.sdk.service.compute import (ClusterDetails, CommandStatus, +# CommandStatusResponse, Created, +# Language, Results, State) + +# from .conftest import noop_credentials + +# cluster_get = mocker.patch( +# "databricks.sdk.service.compute.ClustersAPI.get", +# return_value=ClusterDetails(state=State.RUNNING), +# ) +# context_create = mocker.patch( +# "databricks.sdk.service.compute.CommandExecutionAPI.create", +# return_value=Wait(lambda **kwargs: Created("y")), +# ) + +# def inner(results_data: any, expect_command: str): +# import json + +# command_execute = mocker.patch( +# "databricks.sdk.service.compute.CommandExecutionAPI.execute", +# return_value=Wait( +# lambda **kwargs: CommandStatusResponse( +# results=Results(data=json.dumps(results_data)), +# status=CommandStatus.FINISHED, +# ) +# ), +# ) + +# def assertions(): +# cluster_get.assert_called_with("x") +# context_create.assert_called_with(cluster_id="x", language=Language.PYTHON) +# command_execute.assert_called_with( +# cluster_id="x", +# context_id="y", +# language=Language.PYTHON, +# command=expect_command, +# ) + +# dbutils = RemoteDbUtils( +# Config( +# host="http://localhost", +# cluster_id="x", +# credentials_strategy=noop_credentials, +# ) +# ) +# return dbutils, assertions + +# return inner + + +# def test_fs_mount(dbutils_proxy): +# command = ( +# "\n" +# " import json\n" +# ' (args, kwargs) = json.loads(\'[["s3://foo", "bar"], {}]\')\n' +# " result = dbutils.fs.mount(*args, **kwargs)\n" +# " dbutils.notebook.exit(json.dumps(result))\n" +# " " +# ) +# dbutils, assertions = dbutils_proxy({}, command) + +# dbutils.fs.mount("s3://foo", "bar") + +# assertions() + + +# def test_fs_update_mount(dbutils_proxy): +# command = ( +# "\n" +# " import json\n" +# ' (args, kwargs) = json.loads(\'[["s3://foo2", "bar"], {}]\')\n' +# " result = dbutils.fs.updateMount(*args, **kwargs)\n" +# " dbutils.notebook.exit(json.dumps(result))\n" +# " " +# ) +# dbutils, assertions = dbutils_proxy({}, command) + +# dbutils.fs.updateMount("s3://foo2", "bar") + +# assertions() + + +# def test_fs_mounts(dbutils_proxy): +# command = ( +# "\n" +# " import json\n" +# " (args, kwargs) = json.loads('[[], {}]')\n" +# " result = dbutils.fs.mounts(*args, **kwargs)\n" +# " dbutils.notebook.exit(json.dumps(result))\n" +# " " +# ) +# dbutils, assertions = dbutils_proxy( +# [ +# ("a", "b", "c"), +# ("d", "e", "f"), +# ], +# command, +# ) + +# mounts = dbutils.fs.mounts() + +# assert len(mounts) == 2 +# assert mounts[0].mountPoint == "a" +# assert mounts[0].source == "b" + +# assertions() + + +# def test_any_proxy(dbutils_proxy): +# command = ( +# "\n" +# " import json\n" +# " (args, kwargs) = json.loads('[[\"a\"], {}]')\n" +# " result = dbutils.notebook.exit(*args, **kwargs)\n" +# " dbutils.notebook.exit(json.dumps(result))\n" +# " " +# ) +# dbutils, assertions = dbutils_proxy("a", command) + +# param = dbutils.notebook.exit("a") + +# assert param == "a" + +# assertions() def test_secrets_get_and_redacting_logs(dbutils, mocker): diff --git a/tests/test_files.py b/tests/test_files.py index fad7c2215..3d60a7029 100644 --- a/tests/test_files.py +++ b/tests/test_files.py @@ -18,13 +18,13 @@ import requests_mock from requests import RequestException -from databricks.sdk import WorkspaceClient from databricks.sdk.databricks.core import Config from databricks.sdk.databricks.errors.platform import (AlreadyExists, BadRequest, InternalError, PermissionDenied, TooManyRequests) +from databricks.sdk.files.v2.client import FilesClient logger = logging.getLogger(__name__) @@ -68,12 +68,12 @@ def run(self, config: Config): config.files_api_client_download_max_total_recovers = self.max_recovers_total config.files_api_client_download_max_total_recovers_without_progressing = self.max_recovers_without_progressing - w = WorkspaceClient(config=config) + fc = FilesClient(config=config) session = MockSession(self) - w.files._api._api_client._session = session + fc._api._api_client._session = session - response = w.files.download("/test").contents + response = fc.download("/test").contents if self.expected_success: actual_content = response.read() assert len(actual_content) == len(session.content) @@ -862,12 +862,12 @@ def run(self, config: Config): upload_state = MultipartUploadServerState() try: - w = WorkspaceClient(config=config) + fc = FilesClient(config=config) with requests_mock.Mocker() as session_mock: self.setup_session_mock(session_mock, upload_state) def upload(): - w.files.upload("/test.txt", io.BytesIO(file_content), overwrite=True) + fc.upload("/test.txt", io.BytesIO(file_content), overwrite=True) if self.expected_exception_type is not None: with pytest.raises(self.expected_exception_type): @@ -1321,11 +1321,11 @@ def custom_matcher(request): session_mock.add_matcher(matcher=custom_matcher) - w = WorkspaceClient(config=config) - w.files._api._api_client._session = session + fc = FilesClient(config=config) + fc._api._api_client._session = session def upload(): - w.files.upload("/test.txt", io.BytesIO(file_content), overwrite=True) + fc.upload("/test.txt", io.BytesIO(file_content), overwrite=True) if self.expected_single_shot: upload() @@ -1633,10 +1633,10 @@ def run(self, config: Config): try: with requests_mock.Mocker() as session_mock: self.setup_session_mock(session_mock, upload_state) - w = WorkspaceClient(config=config) + fc = FilesClient(config=config) def upload(): - w.files.upload("/test.txt", io.BytesIO(file_content), overwrite=self.overwrite) + fc.upload("/test.txt", io.BytesIO(file_content), overwrite=self.overwrite) if self.expected_exception_type is not None: with pytest.raises(self.expected_exception_type): diff --git a/tests/test_jobs_mixin.py b/tests/test_jobs_mixin.py index a9039146c..cc15239c7 100644 --- a/tests/test_jobs_mixin.py +++ b/tests/test_jobs_mixin.py @@ -2,7 +2,7 @@ import re from typing import Optional, Pattern -from databricks.sdk import WorkspaceClient +from databricks.sdk.jobs.v2.client import JobsClient def make_getrun_path_pattern(run_id: int, page_token: Optional[str] = None) -> Pattern[str]: @@ -40,9 +40,9 @@ def test_get_run_with_no_pagination(config, requests_mock): "tasks": [{"run_id": 0}, {"run_id": 1}], } requests_mock.get(make_getrun_path_pattern(1337, "initialToken"), text=json.dumps(run1)) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - run = w.jobs.get_run(1337, page_token="initialToken") + run = jc.get_run(1337, page_token="initialToken") assert run.as_dict() == { "tasks": [{"run_id": 0}, {"run_id": 1}], @@ -50,9 +50,9 @@ def test_get_run_with_no_pagination(config, requests_mock): def test_get_run_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs + from databricks.sdk.jobs.v2 import jobs - cluster_spec = compute.ClusterSpec( + cluster_spec = jobs.JobsClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, num_workers=0, @@ -90,9 +90,9 @@ def test_get_run_pagination_with_tasks(config, requests_mock): make_getrun_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - run = w.jobs.get_run(1337, page_token="initialToken") + run = jc.get_run(1337, page_token="initialToken") assert run.as_dict() == { "tasks": [ @@ -139,9 +139,9 @@ def test_get_run_pagination_with_iterations(config, requests_mock): make_getrun_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(run3), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - run = w.jobs.get_run(1337, page_token="initialToken") + run = jc.get_run(1337, page_token="initialToken") assert run.as_dict() == { "tasks": [{"run_id": 1337}], @@ -162,9 +162,9 @@ def test_get_job_with_no_pagination(config, requests_mock): } } requests_mock.get(make_getjob_path_pattern(1337, "initialToken"), text=json.dumps(job1)) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - job = w.jobs.get(1337, page_token="initialToken") + job = jc.get(1337, page_token="initialToken") assert job.as_dict() == { "settings": { @@ -174,9 +174,9 @@ def test_get_job_with_no_pagination(config, requests_mock): def test_get_job_pagination_with_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs + from databricks.sdk.jobs.v2 import jobs - cluster_spec = compute.ClusterSpec( + cluster_spec = jobs.JobsClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, num_workers=0, @@ -223,9 +223,9 @@ def test_get_job_pagination_with_tasks(config, requests_mock): make_getjob_path_pattern(1337, "tokenToThirdPage"), text=json.dumps(job3), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - job = w.jobs.get(1337, page_token="initialToken") + job = jc.get(1337, page_token="initialToken") assert job.as_dict() == { "settings": { @@ -305,10 +305,10 @@ def test_list_jobs_without_task_expansion(config, requests_mock): make_listjobs_path_pattern("tokenToSecondPage"), text=json.dumps(listjobs_page2), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) # Converts the iterator to a list in order to compare the results - jobs_list = list(w.jobs.list(expand_tasks=False, page_token="initialToken")) + jobs_list = list(jc.list(expand_tasks=False, page_token="initialToken")) jobs_dict = [job.as_dict() for job in jobs_list] assert jobs_dict == [ @@ -349,9 +349,9 @@ def test_list_jobs_without_task_expansion(config, requests_mock): def test_list_jobs_with_many_tasks(config, requests_mock): - from databricks.sdk.service import compute, jobs + from databricks.sdk.jobs.v2 import jobs - cluster_spec = compute.ClusterSpec( + cluster_spec = jobs.JobsClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, num_workers=0, @@ -512,10 +512,10 @@ def test_list_jobs_with_many_tasks(config, requests_mock): make_getjob_path_pattern(400, "tokenToSecondPage_400"), text=json.dumps(getjob_400_page2), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) # Converts the iterator to a list in order to compare the results - jobs_list = list(w.jobs.list(expand_tasks=True, page_token="initialToken")) + jobs_list = list(jc.list(expand_tasks=True, page_token="initialToken")) jobs_dict = [job.as_dict() for job in jobs_list] assert jobs_dict == [ @@ -627,9 +627,9 @@ def test_list_runs_without_task_expansion(config, requests_mock): make_listruns_path_pattern("tokenToSecondPage"), text=json.dumps(listruns_page2), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - runs_list = list(w.jobs.list_runs(expand_tasks=False, page_token="initialToken")) + runs_list = list(jc.list_runs(expand_tasks=False, page_token="initialToken")) runs_dict = [run.as_dict() for run in runs_list] assert runs_dict == [ @@ -732,9 +732,9 @@ def test_list_runs(config, requests_mock): make_getrun_path_pattern(400, "tokenToSecondPage_400"), text=json.dumps(getrun_400_page2), ) - w = WorkspaceClient(config=config) + jc = JobsClient(config=config) - runs_list = list(w.jobs.list_runs(expand_tasks=True, page_token="initialToken")) + runs_list = list(jc.list_runs(expand_tasks=True, page_token="initialToken")) runs_dict = [run.as_dict() for run in runs_list] assert runs_dict == [ diff --git a/tests/test_misc.py b/tests/test_misc.py index e614da32b..47b2cb067 100644 --- a/tests/test_misc.py +++ b/tests/test_misc.py @@ -1,13 +1,12 @@ -from databricks.sdk.service import catalog +from databricks.sdk.catalog.v2 import catalog +from databricks.sdk.jobs.v2 import jobs # https://github.com/databricks/databricks-sdk-py/issues/135 def test_issue_135(): - from databricks.sdk.service.compute import Library, PythonPyPiLibrary - from databricks.sdk.service.jobs import Task - jts = Task( - libraries=[Library(pypi=PythonPyPiLibrary(package="databricks-sdk"))], + jts = jobs.Task( + libraries=[jobs.Library(pypi=jobs.PythonPyPiLibrary(package="databricks-sdk"))], task_key="abc", ) @@ -19,12 +18,9 @@ def test_issue_135(): # https://github.com/databricks/databricks-sdk-py/issues/103 def test_issue_103(): - from databricks.sdk.service.compute import ClusterSpec - from databricks.sdk.service.jobs import JobCluster - - jc = JobCluster( + jc = jobs.JobCluster( job_cluster_key="no_worker", - new_cluster=ClusterSpec( + new_cluster=jobs.JobsClusterSpec( spark_version="11.3.x-scala2.12", custom_tags={"ResourceClass": "SingleNode"}, num_workers=0, diff --git a/tests/test_open_ai_mixin.py b/tests/test_open_ai_mixin.py index 457479140..c685c682d 100644 --- a/tests/test_open_ai_mixin.py +++ b/tests/test_open_ai_mixin.py @@ -4,16 +4,15 @@ import pytest from databricks.sdk.databricks.core import Config -from databricks.sdk.service.serving import ExternalFunctionRequestHttpMethod +from databricks.sdk.serving.v2.client import ServingEndpointsClient +from databricks.sdk.serving.v2.serving import ExternalFunctionRequestHttpMethod def test_open_ai_client(monkeypatch): - from databricks.sdk import WorkspaceClient - monkeypatch.setenv("DATABRICKS_HOST", "test_host") monkeypatch.setenv("DATABRICKS_TOKEN", "test_token") - w = WorkspaceClient(config=Config()) - client = w.serving_endpoints.get_open_ai_client() + sec = ServingEndpointsClient(config=Config()) + client = sec.get_open_ai_client() assert client.base_url == "https://test_host/serving-endpoints/" assert client.api_key == "no-token" @@ -21,18 +20,17 @@ def test_open_ai_client(monkeypatch): @pytest.mark.skipif(sys.version_info < (3, 8), reason="Requires Python > 3.7") def test_langchain_open_ai_client(monkeypatch): - from databricks.sdk import WorkspaceClient - monkeypatch.setenv("DATABRICKS_HOST", "test_host") monkeypatch.setenv("DATABRICKS_TOKEN", "test_token") - w = WorkspaceClient(config=Config()) - client = w.serving_endpoints.get_langchain_chat_open_ai_client("databricks-meta-llama-3-1-70b-instruct") + sec = ServingEndpointsClient(config=Config()) + client = sec.get_langchain_chat_open_ai_client("databricks-meta-llama-3-1-70b-instruct") assert client.openai_api_base == "https://test_host/serving-endpoints" assert client.model_name == "databricks-meta-llama-3-1-70b-instruct" def test_http_request(w, requests_mock): + sec = ServingEndpointsClient(config=w) headers = { "Accept": "text/plain", "Content-Type": "application/json", @@ -46,7 +44,7 @@ def test_http_request(w, requests_mock): content=blob_response.getvalue(), status_code=200, ) - response = w.serving_endpoints.http_request( + response = sec.http_request( conn="test_conn", method=ExternalFunctionRequestHttpMethod.GET, path="test_path",