From e05c8c08679bae0d25fbccf4a0b0035edd4ed44e Mon Sep 17 00:00:00 2001 From: Miles Yucht Date: Mon, 28 Aug 2023 16:06:09 +0200 Subject: [PATCH] Update OpenAPI Spec & Merge Files API (#304) ## Changes Support for Files API. ## Tests - [x] Integration tests for FilesAPI updated to create their own schema/volume. - [ ] `make test` run locally - [ ] `make fmt` applied - [ ] relevant integration tests applied --- .codegen/__init__.py.tmpl | 3 +- .codegen/_openapi_sha | 2 +- .codegen/service.py.tmpl | 9 +- databricks/sdk/__init__.py | 15 +- databricks/sdk/core.py | 5 +- databricks/sdk/mixins/files.py | 17 +- databricks/sdk/service/catalog.py | 1169 ++++++++++++++++-------- databricks/sdk/service/compute.py | 143 +-- databricks/sdk/service/files.py | 79 +- databricks/sdk/service/ml.py | 48 +- databricks/sdk/service/provisioning.py | 33 +- databricks/sdk/service/sharing.py | 1 + databricks/sdk/service/sql.py | 58 ++ docs/account/metastores.rst | 2 +- docs/account/storage_credentials.rst | 2 +- docs/account/workspaces.rst | 17 +- docs/workspace/catalogs.rst | 8 +- docs/workspace/experiments.rst | 6 +- docs/workspace/instance_pools.rst | 8 +- docs/workspace/libraries.rst | 2 +- docs/workspace/model_versions.rst | 112 +++ docs/workspace/registered_models.rst | 181 ++++ docs/workspace/statement_execution.rst | 32 +- docs/workspace/workspace-catalog.rst | 4 +- tests/integration/test_files.py | 52 +- 25 files changed, 1430 insertions(+), 578 deletions(-) create mode 100644 docs/workspace/model_versions.rst create mode 100644 docs/workspace/registered_models.rst diff --git a/.codegen/__init__.py.tmpl b/.codegen/__init__.py.tmpl index 1f44121e3..7765fc7db 100644 --- a/.codegen/__init__.py.tmpl +++ b/.codegen/__init__.py.tmpl @@ -1,7 +1,7 @@ import databricks.sdk.core as client import databricks.sdk.dbutils as dbutils -from databricks.sdk.mixins.files import DbfsExt, FilesMixin +from databricks.sdk.mixins.files import DbfsExt from databricks.sdk.mixins.compute import ClustersExt from databricks.sdk.mixins.workspace import WorkspaceExt {{- range .Services}} @@ -51,7 +51,6 @@ class WorkspaceClient: self.config = config.copy() self.dbutils = _make_dbutils(self.config) self.api_client = client.ApiClient(self.config) - self.files = FilesMixin(self.api_client) {{- range .Services}}{{if not .IsAccounts}} self.{{.SnakeName}} = {{template "api" .}}(self.api_client){{end -}}{{end}} diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 8af74cf44..7d69db0e7 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -beff621d7b3e1d59244e2e34fc53a496f310e130 \ No newline at end of file +be3d4a799362f0d5ddbfeb0a0acdfb91f8736a3b \ No newline at end of file diff --git a/.codegen/service.py.tmpl b/.codegen/service.py.tmpl index da466b739..b5d175f18 100644 --- a/.codegen/service.py.tmpl +++ b/.codegen/service.py.tmpl @@ -200,8 +200,13 @@ class {{.Name}}API:{{if .Description}} {{if or .Request.HasJsonField .Request.HasQueryField -}} {{if .Request.HasJsonField}}body = {}{{end}}{{if .Request.HasQueryField}} query = {}{{end}} - {{- range .Request.Fields}}{{if not .IsPath}} - if {{.SnakeName}} is not None: {{if .IsQuery}}query{{else}}body{{end}}['{{.Name}}'] = {{template "method-param-bind" .}}{{end}}{{end}} + {{- range .Request.Fields}}{{ if not .IsPath }} + {{- if .IsQuery }} + if {{.SnakeName}} is not None: query['{{.Name}}'] = {{template "method-param-bind" .}}{{end}} + {{- if .IsJson }} + if {{.SnakeName}} is not None: body['{{.Name}}'] = {{template "method-param-bind" .}}{{end}} + {{- end}} + {{- end}} {{- end}} {{- end}} diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 1e3d2c7fe..6bb372257 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -1,7 +1,7 @@ import databricks.sdk.core as client import databricks.sdk.dbutils as dbutils from databricks.sdk.mixins.compute import ClustersExt -from databricks.sdk.mixins.files import DbfsExt, FilesMixin +from databricks.sdk.mixins.files import DbfsExt from databricks.sdk.mixins.workspace import WorkspaceExt from databricks.sdk.service.billing import (BillableUsageAPI, BudgetsAPI, LogDeliveryAPI) @@ -12,9 +12,10 @@ ConnectionsAPI, ExternalLocationsAPI, FunctionsAPI, GrantsAPI, MetastoresAPI, - SchemasAPI, SecurableTagsAPI, + ModelVersionsAPI, + RegisteredModelsAPI, SchemasAPI, StorageCredentialsAPI, - SubentityTagsAPI, SystemSchemasAPI, + SystemSchemasAPI, TableConstraintsAPI, TablesAPI, VolumesAPI, WorkspaceBindingsAPI) from databricks.sdk.service.compute import (ClusterPoliciesAPI, ClustersAPI, @@ -23,7 +24,7 @@ InstancePoolsAPI, InstanceProfilesAPI, LibrariesAPI, PolicyFamiliesAPI) -from databricks.sdk.service.files import DbfsAPI +from databricks.sdk.service.files import DbfsAPI, FilesAPI from databricks.sdk.service.iam import (AccountAccessControlAPI, AccountAccessControlProxyAPI, AccountGroupsAPI, @@ -129,7 +130,6 @@ def __init__(self, self.config = config.copy() self.dbutils = _make_dbutils(self.config) self.api_client = client.ApiClient(self.config) - self.files = FilesMixin(self.api_client) self.account_access_control_proxy = AccountAccessControlProxyAPI(self.api_client) self.alerts = AlertsAPI(self.api_client) self.artifact_allowlists = ArtifactAllowlistsAPI(self.api_client) @@ -146,6 +146,7 @@ def __init__(self, self.dbsql_permissions = DbsqlPermissionsAPI(self.api_client) self.experiments = ExperimentsAPI(self.api_client) self.external_locations = ExternalLocationsAPI(self.api_client) + self.files = FilesAPI(self.api_client) self.functions = FunctionsAPI(self.api_client) self.git_credentials = GitCredentialsAPI(self.api_client) self.global_init_scripts = GlobalInitScriptsAPI(self.api_client) @@ -158,6 +159,7 @@ def __init__(self, self.libraries = LibrariesAPI(self.api_client) self.metastores = MetastoresAPI(self.api_client) self.model_registry = ModelRegistryAPI(self.api_client) + self.model_versions = ModelVersionsAPI(self.api_client) self.permissions = PermissionsAPI(self.api_client) self.pipelines = PipelinesAPI(self.api_client) self.policy_families = PolicyFamiliesAPI(self.api_client) @@ -166,16 +168,15 @@ def __init__(self, self.query_history = QueryHistoryAPI(self.api_client) self.recipient_activation = RecipientActivationAPI(self.api_client) self.recipients = RecipientsAPI(self.api_client) + self.registered_models = RegisteredModelsAPI(self.api_client) self.repos = ReposAPI(self.api_client) self.schemas = SchemasAPI(self.api_client) self.secrets = SecretsAPI(self.api_client) - self.securable_tags = SecurableTagsAPI(self.api_client) self.service_principals = ServicePrincipalsAPI(self.api_client) self.serving_endpoints = ServingEndpointsAPI(self.api_client) self.shares = SharesAPI(self.api_client) self.statement_execution = StatementExecutionAPI(self.api_client) self.storage_credentials = StorageCredentialsAPI(self.api_client) - self.subentity_tags = SubentityTagsAPI(self.api_client) self.system_schemas = SystemSchemasAPI(self.api_client) self.table_constraints = TableConstraintsAPI(self.api_client) self.tables = TablesAPI(self.api_client) diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py index 3611f06e7..e27970563 100644 --- a/databricks/sdk/core.py +++ b/databricks/sdk/core.py @@ -988,11 +988,12 @@ def do(self, raw: bool = False, files=None, data=None) -> Union[dict, BinaryIO]: + # Remove extra `/` from path for Files API + # Once we've fixed the OpenAPI spec, we can remove this + path = re.sub('^/api/2.0/fs/files//', '/api/2.0/fs/files/', path) if headers is None: headers = {} headers['User-Agent'] = self._user_agent_base - # Replace // with / in path (for Files API where users specify filenames beginning with /) - path = re.sub(r'//', '/', path) response = self._session.request(method, f"{self._cfg.host}{path}", params=self._fix_query_string(query), diff --git a/databricks/sdk/mixins/files.py b/databricks/sdk/mixins/files.py index 03d59d948..351468ed3 100644 --- a/databricks/sdk/mixins/files.py +++ b/databricks/sdk/mixins/files.py @@ -8,7 +8,7 @@ from types import TracebackType from typing import TYPE_CHECKING, AnyStr, BinaryIO, Iterable, Iterator, Type -from databricks.sdk.core import ApiClient, DatabricksError +from databricks.sdk.core import DatabricksError from ..service import files @@ -397,18 +397,3 @@ def move_(self, src: str, dst: str, *, recursive=False, overwrite=False): # do cross-fs moving self.copy(src, dst, recursive=recursive, overwrite=overwrite) source.delete(recursive=recursive) - - -class FilesMixin: - - def __init__(self, api_client: ApiClient): - self._api = api_client - - def upload(self, path: str, src: BinaryIO): - self._api.do('PUT', f'/api/2.0/fs/files{path}', data=src) - - def download(self, path: str) -> BinaryIO: - return self._api.do('GET', f'/api/2.0/fs/files{path}', raw=True) - - def delete(self, path: str): - self._api.do('DELETE', f'/api/2.0/fs/files{path}') diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index f152a43a7..2fa684831 100755 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -271,6 +271,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'AzureServicePrincipal': @dataclass class CatalogInfo: + browse_only: Optional[bool] = None catalog_type: Optional['CatalogType'] = None comment: Optional[str] = None connection_name: Optional[str] = None @@ -278,6 +279,7 @@ class CatalogInfo: created_by: Optional[str] = None effective_predictive_optimization_flag: Optional['EffectivePredictiveOptimizationFlag'] = None enable_predictive_optimization: Optional['EnablePredictiveOptimization'] = None + full_name: Optional[str] = None isolation_mode: Optional['IsolationMode'] = None metastore_id: Optional[str] = None name: Optional[str] = None @@ -285,6 +287,9 @@ class CatalogInfo: owner: Optional[str] = None properties: Optional['Dict[str,str]'] = None provider_name: Optional[str] = None + provisioning_info: Optional['ProvisioningInfo'] = None + securable_kind: Optional['CatalogInfoSecurableKind'] = None + securable_type: Optional[str] = None share_name: Optional[str] = None storage_location: Optional[str] = None storage_root: Optional[str] = None @@ -293,6 +298,7 @@ class CatalogInfo: def as_dict(self) -> dict: body = {} + if self.browse_only is not None: body['browse_only'] = self.browse_only if self.catalog_type is not None: body['catalog_type'] = self.catalog_type.value if self.comment is not None: body['comment'] = self.comment if self.connection_name is not None: body['connection_name'] = self.connection_name @@ -304,6 +310,7 @@ def as_dict(self) -> dict: ) if self.enable_predictive_optimization is not None: body['enable_predictive_optimization'] = self.enable_predictive_optimization.value + if self.full_name is not None: body['full_name'] = self.full_name if self.isolation_mode is not None: body['isolation_mode'] = self.isolation_mode.value if self.metastore_id is not None: body['metastore_id'] = self.metastore_id if self.name is not None: body['name'] = self.name @@ -311,6 +318,9 @@ def as_dict(self) -> dict: if self.owner is not None: body['owner'] = self.owner if self.properties: body['properties'] = self.properties if self.provider_name is not None: body['provider_name'] = self.provider_name + if self.provisioning_info is not None: body['provisioning_info'] = self.provisioning_info.value + if self.securable_kind is not None: body['securable_kind'] = self.securable_kind.value + if self.securable_type is not None: body['securable_type'] = self.securable_type if self.share_name is not None: body['share_name'] = self.share_name if self.storage_location is not None: body['storage_location'] = self.storage_location if self.storage_root is not None: body['storage_root'] = self.storage_root @@ -320,7 +330,8 @@ def as_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, any]) -> 'CatalogInfo': - return cls(catalog_type=_enum(d, 'catalog_type', CatalogType), + return cls(browse_only=d.get('browse_only', None), + catalog_type=_enum(d, 'catalog_type', CatalogType), comment=d.get('comment', None), connection_name=d.get('connection_name', None), created_at=d.get('created_at', None), @@ -329,6 +340,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'CatalogInfo': d, 'effective_predictive_optimization_flag', EffectivePredictiveOptimizationFlag), enable_predictive_optimization=_enum(d, 'enable_predictive_optimization', EnablePredictiveOptimization), + full_name=d.get('full_name', None), isolation_mode=_enum(d, 'isolation_mode', IsolationMode), metastore_id=d.get('metastore_id', None), name=d.get('name', None), @@ -336,6 +348,9 @@ def from_dict(cls, d: Dict[str, any]) -> 'CatalogInfo': owner=d.get('owner', None), properties=d.get('properties', None), provider_name=d.get('provider_name', None), + provisioning_info=_enum(d, 'provisioning_info', ProvisioningInfo), + securable_kind=_enum(d, 'securable_kind', CatalogInfoSecurableKind), + securable_type=d.get('securable_type', None), share_name=d.get('share_name', None), storage_location=d.get('storage_location', None), storage_root=d.get('storage_root', None), @@ -343,6 +358,26 @@ def from_dict(cls, d: Dict[str, any]) -> 'CatalogInfo': updated_by=d.get('updated_by', None)) +class CatalogInfoSecurableKind(Enum): + """Kind of catalog securable.""" + + CATALOG_DELTASHARING = 'CATALOG_DELTASHARING' + CATALOG_FOREIGN_BIGQUERY = 'CATALOG_FOREIGN_BIGQUERY' + CATALOG_FOREIGN_DATABRICKS = 'CATALOG_FOREIGN_DATABRICKS' + CATALOG_FOREIGN_MYSQL = 'CATALOG_FOREIGN_MYSQL' + CATALOG_FOREIGN_POSTGRESQL = 'CATALOG_FOREIGN_POSTGRESQL' + CATALOG_FOREIGN_REDSHIFT = 'CATALOG_FOREIGN_REDSHIFT' + CATALOG_FOREIGN_SNOWFLAKE = 'CATALOG_FOREIGN_SNOWFLAKE' + CATALOG_FOREIGN_SQLDW = 'CATALOG_FOREIGN_SQLDW' + CATALOG_FOREIGN_SQLSERVER = 'CATALOG_FOREIGN_SQLSERVER' + CATALOG_INTERNAL = 'CATALOG_INTERNAL' + CATALOG_ONLINE = 'CATALOG_ONLINE' + CATALOG_ONLINE_INDEX = 'CATALOG_ONLINE_INDEX' + CATALOG_STANDARD = 'CATALOG_STANDARD' + CATALOG_SYSTEM = 'CATALOG_SYSTEM' + CATALOG_SYSTEM_DELTASHARING = 'CATALOG_SYSTEM_DELTASHARING' + + class CatalogType(Enum): """The type of the catalog.""" @@ -455,7 +490,7 @@ class ConnectionInfo: options: Optional['Dict[str,str]'] = None owner: Optional[str] = None properties: Optional['Dict[str,str]'] = None - provisioning_state: Optional['ProvisioningState'] = None + provisioning_info: Optional['ProvisioningInfo'] = None read_only: Optional[bool] = None securable_kind: Optional['ConnectionInfoSecurableKind'] = None securable_type: Optional[str] = None @@ -477,7 +512,7 @@ def as_dict(self) -> dict: if self.options: body['options'] = self.options if self.owner is not None: body['owner'] = self.owner if self.properties: body['properties'] = self.properties - if self.provisioning_state is not None: body['provisioning_state'] = self.provisioning_state.value + if self.provisioning_info is not None: body['provisioning_info'] = self.provisioning_info.value if self.read_only is not None: body['read_only'] = self.read_only if self.securable_kind is not None: body['securable_kind'] = self.securable_kind.value if self.securable_type is not None: body['securable_type'] = self.securable_type @@ -500,7 +535,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'ConnectionInfo': options=d.get('options', None), owner=d.get('owner', None), properties=d.get('properties', None), - provisioning_state=_enum(d, 'provisioning_state', ProvisioningState), + provisioning_info=_enum(d, 'provisioning_info', ProvisioningInfo), read_only=d.get('read_only', None), securable_kind=_enum(d, 'securable_kind', ConnectionInfoSecurableKind), securable_type=d.get('securable_type', None), @@ -540,6 +575,7 @@ class CreateCatalog: name: str comment: Optional[str] = None connection_name: Optional[str] = None + options: Optional['Dict[str,str]'] = None properties: Optional['Dict[str,str]'] = None provider_name: Optional[str] = None share_name: Optional[str] = None @@ -550,6 +586,7 @@ def as_dict(self) -> dict: if self.comment is not None: body['comment'] = self.comment if self.connection_name is not None: body['connection_name'] = self.connection_name if self.name is not None: body['name'] = self.name + if self.options: body['options'] = self.options if self.properties: body['properties'] = self.properties if self.provider_name is not None: body['provider_name'] = self.provider_name if self.share_name is not None: body['share_name'] = self.share_name @@ -561,6 +598,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'CreateCatalog': return cls(comment=d.get('comment', None), connection_name=d.get('connection_name', None), name=d.get('name', None), + options=d.get('options', None), properties=d.get('properties', None), provider_name=d.get('provider_name', None), share_name=d.get('share_name', None), @@ -779,6 +817,32 @@ def from_dict(cls, d: Dict[str, any]) -> 'CreateMetastoreAssignment': workspace_id=d.get('workspace_id', None)) +@dataclass +class CreateRegisteredModelRequest: + catalog_name: str + schema_name: str + name: str + comment: Optional[str] = None + storage_location: Optional[str] = None + + def as_dict(self) -> dict: + body = {} + if self.catalog_name is not None: body['catalog_name'] = self.catalog_name + if self.comment is not None: body['comment'] = self.comment + if self.name is not None: body['name'] = self.name + if self.schema_name is not None: body['schema_name'] = self.schema_name + if self.storage_location is not None: body['storage_location'] = self.storage_location + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'CreateRegisteredModelRequest': + return cls(catalog_name=d.get('catalog_name', None), + comment=d.get('comment', None), + name=d.get('name', None), + schema_name=d.get('schema_name', None), + storage_location=d.get('storage_location', None)) + + @dataclass class CreateSchema: name: str @@ -940,6 +1004,14 @@ def from_dict(cls, d: Dict[str, any]) -> 'DatabricksGcpServiceAccountResponse': return cls(credential_id=d.get('credential_id', None), email=d.get('email', None)) +@dataclass +class DeleteModelVersionRequest: + """Delete a Model Version""" + + full_name: str + version: int + + @dataclass class DeltaRuntimePropertiesKvPairs: """Properties pertaining to the current state of the delta table as given by the commit server. @@ -1379,6 +1451,14 @@ class FunctionParameterType(Enum): PARAM = 'PARAM' +@dataclass +class GetByAliasRequest: + """Get Model Version By Alias""" + + full_name: str + alias: str + + @dataclass class GetMetastoreSummaryResponse: cloud: Optional[str] = None @@ -1460,6 +1540,14 @@ class GetMetastoreSummaryResponseDeltaSharingScope(Enum): INTERNAL_AND_EXTERNAL = 'INTERNAL_AND_EXTERNAL' +@dataclass +class GetModelVersionRequest: + """Get a Model Version""" + + full_name: str + version: int + + class IsolationMode(Enum): """Whether the current securable is accessible from all workspaces or a specific set of workspaces.""" @@ -1538,6 +1626,49 @@ def from_dict(cls, d: Dict[str, any]) -> 'ListMetastoresResponse': return cls(metastores=_repeated(d, 'metastores', MetastoreInfo)) +@dataclass +class ListModelVersionsRequest: + """List Model Versions""" + + full_name: str + max_results: Optional[int] = None + page_token: Optional[str] = None + + +@dataclass +class ListModelVersionsResponse: + model_versions: Optional['List[ModelVersionInfo]'] = None + next_page_token: Optional[str] = None + + def as_dict(self) -> dict: + body = {} + if self.model_versions: body['model_versions'] = [v.as_dict() for v in self.model_versions] + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'ListModelVersionsResponse': + return cls(model_versions=_repeated(d, 'model_versions', ModelVersionInfo), + next_page_token=d.get('next_page_token', None)) + + +@dataclass +class ListRegisteredModelsResponse: + next_page_token: Optional[str] = None + registered_models: Optional['List[RegisteredModelInfo]'] = None + + def as_dict(self) -> dict: + body = {} + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token + if self.registered_models: body['registered_models'] = [v.as_dict() for v in self.registered_models] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'ListRegisteredModelsResponse': + return cls(next_page_token=d.get('next_page_token', None), + registered_models=_repeated(d, 'registered_models', RegisteredModelInfo)) + + @dataclass class ListSchemasResponse: schemas: Optional['List[SchemaInfo]'] = None @@ -1552,13 +1683,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'ListSchemasResponse': return cls(schemas=_repeated(d, 'schemas', SchemaInfo)) -class ListSecurableType(Enum): - - CATALOG = 'catalog' - SCHEMA = 'schema' - TABLE = 'table' - - @dataclass class ListStorageCredentialsResponse: storage_credentials: Optional['List[StorageCredentialInfo]'] = None @@ -1741,6 +1865,79 @@ class MetastoreInfoDeltaSharingScope(Enum): INTERNAL_AND_EXTERNAL = 'INTERNAL_AND_EXTERNAL' +@dataclass +class ModelVersionInfo: + catalog_name: Optional[str] = None + comment: Optional[str] = None + created_at: Optional[int] = None + created_by: Optional[str] = None + id: Optional[str] = None + metastore_id: Optional[str] = None + model_name: Optional[str] = None + model_version_dependencies: Optional['List[Dependency]'] = None + run_id: Optional[str] = None + run_workspace_id: Optional[int] = None + schema_name: Optional[str] = None + source: Optional[str] = None + status: Optional['ModelVersionInfoStatus'] = None + storage_location: Optional[str] = None + updated_at: Optional[int] = None + updated_by: Optional[str] = None + version: Optional[int] = None + + def as_dict(self) -> dict: + body = {} + if self.catalog_name is not None: body['catalog_name'] = self.catalog_name + if self.comment is not None: body['comment'] = self.comment + if self.created_at is not None: body['created_at'] = self.created_at + if self.created_by is not None: body['created_by'] = self.created_by + if self.id is not None: body['id'] = self.id + if self.metastore_id is not None: body['metastore_id'] = self.metastore_id + if self.model_name is not None: body['model_name'] = self.model_name + if self.model_version_dependencies: + body['model_version_dependencies'] = [v.as_dict() for v in self.model_version_dependencies] + if self.run_id is not None: body['run_id'] = self.run_id + if self.run_workspace_id is not None: body['run_workspace_id'] = self.run_workspace_id + if self.schema_name is not None: body['schema_name'] = self.schema_name + if self.source is not None: body['source'] = self.source + if self.status is not None: body['status'] = self.status.value + if self.storage_location is not None: body['storage_location'] = self.storage_location + if self.updated_at is not None: body['updated_at'] = self.updated_at + if self.updated_by is not None: body['updated_by'] = self.updated_by + if self.version is not None: body['version'] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'ModelVersionInfo': + return cls(catalog_name=d.get('catalog_name', None), + comment=d.get('comment', None), + created_at=d.get('created_at', None), + created_by=d.get('created_by', None), + id=d.get('id', None), + metastore_id=d.get('metastore_id', None), + model_name=d.get('model_name', None), + model_version_dependencies=_repeated(d, 'model_version_dependencies', Dependency), + run_id=d.get('run_id', None), + run_workspace_id=d.get('run_workspace_id', None), + schema_name=d.get('schema_name', None), + source=d.get('source', None), + status=_enum(d, 'status', ModelVersionInfoStatus), + storage_location=d.get('storage_location', None), + updated_at=d.get('updated_at', None), + updated_by=d.get('updated_by', None), + version=d.get('version', None)) + + +class ModelVersionInfoStatus(Enum): + """Current status of the model version. Newly created model versions start in PENDING_REGISTRATION + status, then move to READY status once the model version files are uploaded and the model + version is finalized. Only model versions in READY status can be loaded for inference or served.""" + + FAILED_REGISTRATION = 'FAILED_REGISTRATION' + PENDING_REGISTRATION = 'PENDING_REGISTRATION' + READY = 'READY' + + @dataclass class NamedTableConstraint: name: str @@ -1817,6 +2014,7 @@ class Privilege(Enum): CREATE_FUNCTION = 'CREATE_FUNCTION' CREATE_MANAGED_STORAGE = 'CREATE_MANAGED_STORAGE' CREATE_MATERIALIZED_VIEW = 'CREATE_MATERIALIZED_VIEW' + CREATE_MODEL = 'CREATE_MODEL' CREATE_PROVIDER = 'CREATE_PROVIDER' CREATE_RECIPIENT = 'CREATE_RECIPIENT' CREATE_SCHEMA = 'CREATE_SCHEMA' @@ -1862,7 +2060,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'PrivilegeAssignment': PropertiesKvPairs = Dict[str, str] -class ProvisioningState(Enum): +class ProvisioningInfo(Enum): """Status of an asynchronously provisioned resource.""" ACTIVE = 'ACTIVE' @@ -1872,6 +2070,74 @@ class ProvisioningState(Enum): STATE_UNSPECIFIED = 'STATE_UNSPECIFIED' +@dataclass +class RegisteredModelAlias: + """Registered model alias.""" + + alias_name: Optional[str] = None + version_num: Optional[int] = None + + def as_dict(self) -> dict: + body = {} + if self.alias_name is not None: body['alias_name'] = self.alias_name + if self.version_num is not None: body['version_num'] = self.version_num + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'RegisteredModelAlias': + return cls(alias_name=d.get('alias_name', None), version_num=d.get('version_num', None)) + + +@dataclass +class RegisteredModelInfo: + aliases: Optional['List[RegisteredModelAlias]'] = None + catalog_name: Optional[str] = None + comment: Optional[str] = None + created_at: Optional[int] = None + created_by: Optional[str] = None + full_name: Optional[str] = None + metastore_id: Optional[str] = None + name: Optional[str] = None + owner: Optional[str] = None + schema_name: Optional[str] = None + storage_location: Optional[str] = None + updated_at: Optional[int] = None + updated_by: Optional[str] = None + + def as_dict(self) -> dict: + body = {} + if self.aliases: body['aliases'] = [v.as_dict() for v in self.aliases] + if self.catalog_name is not None: body['catalog_name'] = self.catalog_name + if self.comment is not None: body['comment'] = self.comment + if self.created_at is not None: body['created_at'] = self.created_at + if self.created_by is not None: body['created_by'] = self.created_by + if self.full_name is not None: body['full_name'] = self.full_name + if self.metastore_id is not None: body['metastore_id'] = self.metastore_id + if self.name is not None: body['name'] = self.name + if self.owner is not None: body['owner'] = self.owner + if self.schema_name is not None: body['schema_name'] = self.schema_name + if self.storage_location is not None: body['storage_location'] = self.storage_location + if self.updated_at is not None: body['updated_at'] = self.updated_at + if self.updated_by is not None: body['updated_by'] = self.updated_by + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'RegisteredModelInfo': + return cls(aliases=_repeated(d, 'aliases', RegisteredModelAlias), + catalog_name=d.get('catalog_name', None), + comment=d.get('comment', None), + created_at=d.get('created_at', None), + created_by=d.get('created_by', None), + full_name=d.get('full_name', None), + metastore_id=d.get('metastore_id', None), + name=d.get('name', None), + owner=d.get('owner', None), + schema_name=d.get('schema_name', None), + storage_location=d.get('storage_location', None), + updated_at=d.get('updated_at', None), + updated_by=d.get('updated_by', None)) + + @dataclass class SchemaInfo: catalog_name: Optional[str] = None @@ -1976,6 +2242,26 @@ def from_dict(cls, d: Dict[str, any]) -> 'SetArtifactAllowlist': artifact_type=_enum(d, 'artifact_type', ArtifactType)) +@dataclass +class SetRegisteredModelAliasRequest: + full_name: str + alias: str + version_num: int + + def as_dict(self) -> dict: + body = {} + if self.alias is not None: body['alias'] = self.alias + if self.full_name is not None: body['full_name'] = self.full_name + if self.version_num is not None: body['version_num'] = self.version_num + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'SetRegisteredModelAliasRequest': + return cls(alias=d.get('alias', None), + full_name=d.get('full_name', None), + version_num=d.get('version_num', None)) + + @dataclass class SseEncryptionDetails: """Server-Side Encryption properties for clients communicating with AWS s3.""" @@ -2295,220 +2581,88 @@ class TableType(Enum): @dataclass -class TagChanges: - add_tags: Optional['List[TagKeyValuePair]'] = None - remove: Optional['List[str]'] = None +class UpdateCatalog: + comment: Optional[str] = None + isolation_mode: Optional['IsolationMode'] = None + name: Optional[str] = None + options: Optional['Dict[str,str]'] = None + owner: Optional[str] = None + properties: Optional['Dict[str,str]'] = None def as_dict(self) -> dict: body = {} - if self.add_tags: body['add_tags'] = [v.as_dict() for v in self.add_tags] - if self.remove: body['remove'] = [v for v in self.remove] + if self.comment is not None: body['comment'] = self.comment + if self.isolation_mode is not None: body['isolation_mode'] = self.isolation_mode.value + if self.name is not None: body['name'] = self.name + if self.options: body['options'] = self.options + if self.owner is not None: body['owner'] = self.owner + if self.properties: body['properties'] = self.properties return body @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagChanges': - return cls(add_tags=_repeated(d, 'add_tags', TagKeyValuePair), remove=d.get('remove', None)) + def from_dict(cls, d: Dict[str, any]) -> 'UpdateCatalog': + return cls(comment=d.get('comment', None), + isolation_mode=_enum(d, 'isolation_mode', IsolationMode), + name=d.get('name', None), + options=d.get('options', None), + owner=d.get('owner', None), + properties=d.get('properties', None)) @dataclass -class TagKeyValuePair: - key: str - value: str +class UpdateConnection: + name: str + options: 'Dict[str,str]' + name_arg: Optional[str] = None def as_dict(self) -> dict: body = {} - if self.key is not None: body['key'] = self.key - if self.value is not None: body['value'] = self.value + if self.name is not None: body['name'] = self.name + if self.name_arg is not None: body['name_arg'] = self.name_arg + if self.options: body['options'] = self.options return body @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagKeyValuePair': - return cls(key=d.get('key', None), value=d.get('value', None)) + def from_dict(cls, d: Dict[str, any]) -> 'UpdateConnection': + return cls(name=d.get('name', None), name_arg=d.get('name_arg', None), options=d.get('options', None)) @dataclass -class TagSecurable: - type: str - full_name: str +class UpdateExternalLocation: + access_point: Optional[str] = None + comment: Optional[str] = None + credential_name: Optional[str] = None + encryption_details: Optional['EncryptionDetails'] = None + force: Optional[bool] = None + name: Optional[str] = None + owner: Optional[str] = None + read_only: Optional[bool] = None + url: Optional[str] = None def as_dict(self) -> dict: body = {} - if self.full_name is not None: body['full_name'] = self.full_name - if self.type is not None: body['type'] = self.type + if self.access_point is not None: body['access_point'] = self.access_point + if self.comment is not None: body['comment'] = self.comment + if self.credential_name is not None: body['credential_name'] = self.credential_name + if self.encryption_details: body['encryption_details'] = self.encryption_details.as_dict() + if self.force is not None: body['force'] = self.force + if self.name is not None: body['name'] = self.name + if self.owner is not None: body['owner'] = self.owner + if self.read_only is not None: body['read_only'] = self.read_only + if self.url is not None: body['url'] = self.url return body @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagSecurable': - return cls(full_name=d.get('full_name', None), type=d.get('type', None)) - - -@dataclass -class TagSecurableAssignment: - securable: 'TagSecurable' - tag_key_value_pairs: 'List[TagKeyValuePair]' - - def as_dict(self) -> dict: - body = {} - if self.securable: body['securable'] = self.securable.as_dict() - if self.tag_key_value_pairs: - body['tag_key_value_pairs'] = [v.as_dict() for v in self.tag_key_value_pairs] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagSecurableAssignment': - return cls(securable=_from_dict(d, 'securable', TagSecurable), - tag_key_value_pairs=_repeated(d, 'tag_key_value_pairs', TagKeyValuePair)) - - -@dataclass -class TagSecurableAssignmentsList: - tag_assignments: 'List[TagSecurableAssignment]' - - def as_dict(self) -> dict: - body = {} - if self.tag_assignments: body['tag_assignments'] = [v.as_dict() for v in self.tag_assignments] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagSecurableAssignmentsList': - return cls(tag_assignments=_repeated(d, 'tag_assignments', TagSecurableAssignment)) - - -@dataclass -class TagSubentity: - type: str - full_name: str - subentity: str - - def as_dict(self) -> dict: - body = {} - if self.full_name is not None: body['full_name'] = self.full_name - if self.subentity is not None: body['subentity'] = self.subentity - if self.type is not None: body['type'] = self.type - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagSubentity': - return cls(full_name=d.get('full_name', None), - subentity=d.get('subentity', None), - type=d.get('type', None)) - - -@dataclass -class TagSubentityAssignmentsList: - tag_assignments: 'List[TagsSubentityAssignment]' - - def as_dict(self) -> dict: - body = {} - if self.tag_assignments: body['tag_assignments'] = [v.as_dict() for v in self.tag_assignments] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagSubentityAssignmentsList': - return cls(tag_assignments=_repeated(d, 'tag_assignments', TagsSubentityAssignment)) - - -@dataclass -class TagsSubentityAssignment: - securable: 'TagSubentity' - subentity: str - tag_key_value_pairs: 'List[TagKeyValuePair]' - - def as_dict(self) -> dict: - body = {} - if self.securable: body['securable'] = self.securable.as_dict() - if self.subentity is not None: body['subentity'] = self.subentity - if self.tag_key_value_pairs: - body['tag_key_value_pairs'] = [v.as_dict() for v in self.tag_key_value_pairs] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'TagsSubentityAssignment': - return cls(securable=_from_dict(d, 'securable', TagSubentity), - subentity=d.get('subentity', None), - tag_key_value_pairs=_repeated(d, 'tag_key_value_pairs', TagKeyValuePair)) - - -@dataclass -class UpdateCatalog: - comment: Optional[str] = None - isolation_mode: Optional['IsolationMode'] = None - name: Optional[str] = None - owner: Optional[str] = None - properties: Optional['Dict[str,str]'] = None - - def as_dict(self) -> dict: - body = {} - if self.comment is not None: body['comment'] = self.comment - if self.isolation_mode is not None: body['isolation_mode'] = self.isolation_mode.value - if self.name is not None: body['name'] = self.name - if self.owner is not None: body['owner'] = self.owner - if self.properties: body['properties'] = self.properties - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'UpdateCatalog': - return cls(comment=d.get('comment', None), - isolation_mode=_enum(d, 'isolation_mode', IsolationMode), - name=d.get('name', None), - owner=d.get('owner', None), - properties=d.get('properties', None)) - - -@dataclass -class UpdateConnection: - name: str - options: 'Dict[str,str]' - name_arg: Optional[str] = None - - def as_dict(self) -> dict: - body = {} - if self.name is not None: body['name'] = self.name - if self.name_arg is not None: body['name_arg'] = self.name_arg - if self.options: body['options'] = self.options - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'UpdateConnection': - return cls(name=d.get('name', None), name_arg=d.get('name_arg', None), options=d.get('options', None)) - - -@dataclass -class UpdateExternalLocation: - access_point: Optional[str] = None - comment: Optional[str] = None - credential_name: Optional[str] = None - encryption_details: Optional['EncryptionDetails'] = None - force: Optional[bool] = None - name: Optional[str] = None - owner: Optional[str] = None - read_only: Optional[bool] = None - url: Optional[str] = None - - def as_dict(self) -> dict: - body = {} - if self.access_point is not None: body['access_point'] = self.access_point - if self.comment is not None: body['comment'] = self.comment - if self.credential_name is not None: body['credential_name'] = self.credential_name - if self.encryption_details: body['encryption_details'] = self.encryption_details.as_dict() - if self.force is not None: body['force'] = self.force - if self.name is not None: body['name'] = self.name - if self.owner is not None: body['owner'] = self.owner - if self.read_only is not None: body['read_only'] = self.read_only - if self.url is not None: body['url'] = self.url - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'UpdateExternalLocation': - return cls(access_point=d.get('access_point', None), - comment=d.get('comment', None), - credential_name=d.get('credential_name', None), - encryption_details=_from_dict(d, 'encryption_details', EncryptionDetails), - force=d.get('force', None), - name=d.get('name', None), - owner=d.get('owner', None), - read_only=d.get('read_only', None), - url=d.get('url', None)) + def from_dict(cls, d: Dict[str, any]) -> 'UpdateExternalLocation': + return cls(access_point=d.get('access_point', None), + comment=d.get('comment', None), + credential_name=d.get('credential_name', None), + encryption_details=_from_dict(d, 'encryption_details', EncryptionDetails), + force=d.get('force', None), + name=d.get('name', None), + owner=d.get('owner', None), + read_only=d.get('read_only', None), + url=d.get('url', None)) @dataclass @@ -2595,6 +2749,26 @@ class UpdateMetastoreDeltaSharingScope(Enum): INTERNAL_AND_EXTERNAL = 'INTERNAL_AND_EXTERNAL' +@dataclass +class UpdateModelVersionRequest: + comment: Optional[str] = None + full_name: Optional[str] = None + version: Optional[int] = None + + def as_dict(self) -> dict: + body = {} + if self.comment is not None: body['comment'] = self.comment + if self.full_name is not None: body['full_name'] = self.full_name + if self.version is not None: body['version'] = self.version + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'UpdateModelVersionRequest': + return cls(comment=d.get('comment', None), + full_name=d.get('full_name', None), + version=d.get('version', None)) + + @dataclass class UpdatePermissions: changes: Optional['List[PermissionsChange]'] = None @@ -2651,6 +2825,29 @@ def from_dict(cls, d: Dict[str, any]) -> 'UpdatePredictiveOptimizationResponse': username=d.get('username', None)) +@dataclass +class UpdateRegisteredModelRequest: + comment: Optional[str] = None + full_name: Optional[str] = None + name: Optional[str] = None + owner: Optional[str] = None + + def as_dict(self) -> dict: + body = {} + if self.comment is not None: body['comment'] = self.comment + if self.full_name is not None: body['full_name'] = self.full_name + if self.name is not None: body['name'] = self.name + if self.owner is not None: body['owner'] = self.owner + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'UpdateRegisteredModelRequest': + return cls(comment=d.get('comment', None), + full_name=d.get('full_name', None), + name=d.get('name', None), + owner=d.get('owner', None)) + + @dataclass class UpdateSchema: comment: Optional[str] = None @@ -2677,13 +2874,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'UpdateSchema': properties=d.get('properties', None)) -class UpdateSecurableType(Enum): - - CATALOG = 'catalog' - SCHEMA = 'schema' - TABLE = 'table' - - @dataclass class UpdateStorageCredential: aws_iam_role: Optional['AwsIamRole'] = None @@ -2727,29 +2917,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'UpdateStorageCredential': skip_validation=d.get('skip_validation', None)) -@dataclass -class UpdateTags: - changes: 'TagChanges' - full_name: Optional[str] = None - securable_type: Optional['UpdateSecurableType'] = None - subentity_name: Optional[str] = None - - def as_dict(self) -> dict: - body = {} - if self.changes: body['changes'] = self.changes.as_dict() - if self.full_name is not None: body['full_name'] = self.full_name - if self.securable_type is not None: body['securable_type'] = self.securable_type.value - if self.subentity_name is not None: body['subentity_name'] = self.subentity_name - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'UpdateTags': - return cls(changes=_from_dict(d, 'changes', TagChanges), - full_name=d.get('full_name', None), - securable_type=_enum(d, 'securable_type', UpdateSecurableType), - subentity_name=d.get('subentity_name', None)) - - @dataclass class UpdateVolumeRequestContent: comment: Optional[str] = None @@ -3127,17 +3294,17 @@ def get(self, metastore_id: str) -> AccountsMetastoreInfo: headers=headers) return AccountsMetastoreInfo.from_dict(res) - def list(self) -> ListMetastoresResponse: + def list(self) -> Iterator[MetastoreInfo]: """Get all metastores associated with an account. Gets all Unity Catalog metastores associated with an account specified by ID. - :returns: :class:`ListMetastoresResponse` + :returns: Iterator over :class:`MetastoreInfo` """ headers = {'Accept': 'application/json', } - res = self._api.do('GET', f'/api/2.0/accounts/{self._api.account_id}/metastores', headers=headers) - return ListMetastoresResponse.from_dict(res) + json = self._api.do('GET', f'/api/2.0/accounts/{self._api.account_id}/metastores', headers=headers) + return [MetastoreInfo.from_dict(v) for v in json.get('metastores', [])] def update(self, metastore_id: str, @@ -3245,7 +3412,7 @@ def get(self, metastore_id: str, name: str) -> AccountsStorageCredentialInfo: headers=headers) return AccountsStorageCredentialInfo.from_dict(res) - def list(self, metastore_id: str) -> ListStorageCredentialsResponse: + def list(self, metastore_id: str) -> Iterator[StorageCredentialInfo]: """Get all storage credentials assigned to a metastore. Gets a list of all storage credentials that have been assigned to given metastore. @@ -3253,7 +3420,7 @@ def list(self, metastore_id: str) -> ListStorageCredentialsResponse: :param metastore_id: str Unity Catalog metastore ID - :returns: :class:`ListStorageCredentialsResponse` + :returns: Iterator over :class:`StorageCredentialInfo` """ headers = {'Accept': 'application/json', } @@ -3261,7 +3428,7 @@ def list(self, metastore_id: str) -> ListStorageCredentialsResponse: 'GET', f'/api/2.0/accounts/{self._api.account_id}/metastores/{metastore_id}/storage-credentials', headers=headers) - return ListStorageCredentialsResponse.from_dict(res) + return [StorageCredentialInfo.from_dict(v) for v in res] def update(self, metastore_id: str, @@ -3355,6 +3522,7 @@ def create(self, *, comment: Optional[str] = None, connection_name: Optional[str] = None, + options: Optional[Dict[str, str]] = None, properties: Optional[Dict[str, str]] = None, provider_name: Optional[str] = None, share_name: Optional[str] = None, @@ -3370,6 +3538,8 @@ def create(self, User-provided free-form text description. :param connection_name: str (optional) The name of the connection to an external data source. + :param options: Dict[str,str] (optional) + A map of key-value properties attached to the securable. :param properties: Dict[str,str] (optional) A map of key-value properties attached to the securable. :param provider_name: str (optional) @@ -3387,6 +3557,7 @@ def create(self, if comment is not None: body['comment'] = comment if connection_name is not None: body['connection_name'] = connection_name if name is not None: body['name'] = name + if options is not None: body['options'] = options if properties is not None: body['properties'] = properties if provider_name is not None: body['provider_name'] = provider_name if share_name is not None: body['share_name'] = share_name @@ -3450,6 +3621,7 @@ def update(self, *, comment: Optional[str] = None, isolation_mode: Optional[IsolationMode] = None, + options: Optional[Dict[str, str]] = None, owner: Optional[str] = None, properties: Optional[Dict[str, str]] = None) -> CatalogInfo: """Update a catalog. @@ -3463,6 +3635,8 @@ def update(self, User-provided free-form text description. :param isolation_mode: :class:`IsolationMode` (optional) Whether the current securable is accessible from all workspaces or a specific set of workspaces. + :param options: Dict[str,str] (optional) + A map of key-value properties attached to the securable. :param owner: str (optional) Username of current owner of catalog. :param properties: Dict[str,str] (optional) @@ -3473,6 +3647,7 @@ def update(self, body = {} if comment is not None: body['comment'] = comment if isolation_mode is not None: body['isolation_mode'] = isolation_mode.value + if options is not None: body['options'] = options if owner is not None: body['owner'] = owner if properties is not None: body['properties'] = properties headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } @@ -4339,6 +4514,405 @@ def update_assignment(self, headers=headers) +class 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, api_client): + self._api = api_client + + def delete(self, full_name: str, version: int): + """Delete a Model Version. + + Deletes a model version from the specified registered model. Any aliases assigned to the model version + will also be deleted. + + The caller must be a metastore admin or an owner of the parent registered model. For the latter case, + the caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + + + """ + + headers = {} + self._api.do('DELETE', + f'/api/2.1/unity-catalog/models/{full_name}/versions/{version}', + headers=headers) + + def get(self, full_name: str, version: int) -> RegisteredModelInfo: + """Get a Model Version. + + Get a model version. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the parent + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + + :returns: :class:`RegisteredModelInfo` + """ + + headers = {'Accept': 'application/json', } + res = self._api.do('GET', + f'/api/2.1/unity-catalog/models/{full_name}/versions/{version}', + headers=headers) + return RegisteredModelInfo.from_dict(res) + + def get_by_alias(self, full_name: str, alias: str) -> ModelVersionInfo: + """Get Model Version By Alias. + + Get a model version by alias. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param alias: str + The name of the alias + + :returns: :class:`ModelVersionInfo` + """ + + headers = {'Accept': 'application/json', } + res = self._api.do('GET', + f'/api/2.1/unity-catalog/models/{full_name}/aliases/{alias}', + headers=headers) + return ModelVersionInfo.from_dict(res) + + def list(self, + full_name: str, + *, + max_results: Optional[int] = None, + page_token: Optional[str] = None) -> Iterator[ModelVersionInfo]: + """List Model Versions. + + List model versions. You can list model versions under a particular schema, or list all model versions + in the current metastore. + + The returned models are filtered based on the privileges of the calling user. For example, the + metastore admin is able to list all the model versions. A regular user needs to be the owner or have + the **EXECUTE** privilege on the parent registered model to recieve the model versions in the + response. For the latter case, the caller must also be the owner or have the **USE_CATALOG** privilege + on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + There is no guarantee of a specific ordering of the elements in the response. + + :param full_name: str + The full three-level name of the registered model under which to list model versions + :param max_results: int (optional) + Max number of model versions to return + :param page_token: str (optional) + Opaque token to send for the next page of results (pagination). + + :returns: Iterator over :class:`ModelVersionInfo` + """ + + query = {} + if max_results is not None: query['max_results'] = max_results + if page_token is not None: query['page_token'] = page_token + headers = {'Accept': 'application/json', } + + while True: + json = self._api.do('GET', + f'/api/2.1/unity-catalog/models/{full_name}/versions', + query=query, + headers=headers) + if 'model_versions' not in json or not json['model_versions']: + return + for v in json['model_versions']: + yield ModelVersionInfo.from_dict(v) + if 'next_page_token' not in json or not json['next_page_token']: + return + query['page_token'] = json['next_page_token'] + + def update(self, full_name: str, version: int, *, comment: Optional[str] = None) -> ModelVersionInfo: + """Update a Model Version. + + Updates the specified model version. + + The caller must be a metastore admin or an owner of the parent registered model. For the latter case, + the caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + Currently only the comment of the model version can be updated. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + :param comment: str (optional) + The comment attached to the model version + + :returns: :class:`ModelVersionInfo` + """ + body = {} + if comment is not None: body['comment'] = comment + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('PATCH', + f'/api/2.1/unity-catalog/models/{full_name}/versions/{version}', + body=body, + headers=headers) + return ModelVersionInfo.from_dict(res) + + +class 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, api_client): + self._api = api_client + + def create(self, + catalog_name: str, + schema_name: str, + name: str, + *, + comment: Optional[str] = None, + storage_location: Optional[str] = None) -> RegisteredModelInfo: + """Create a Registered Model. + + Creates a new registered model in Unity Catalog. + + File storage for model versions in the registered model will be located in the default location which + is specified by the parent schema, or the parent catalog, or the Metastore. + + For registered model creation to succeed, the user must satisfy the following conditions: - The caller + must be a metastore admin, or be the owner of the parent catalog and schema, or have the + **USE_CATALOG** privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + - The caller must have the **CREATE MODEL** or **CREATE FUNCTION** privilege on the parent schema. + + :param catalog_name: str + The name of the catalog where the schema and the registered model reside + :param schema_name: str + The name of the schema where the registered model resides + :param name: str + The name of the registered model + :param comment: str (optional) + The comment attached to the registered model + :param storage_location: str (optional) + The storage location on the cloud under which model version data files are stored + + :returns: :class:`RegisteredModelInfo` + """ + body = {} + if catalog_name is not None: body['catalog_name'] = catalog_name + if comment is not None: body['comment'] = comment + if name is not None: body['name'] = name + if schema_name is not None: body['schema_name'] = schema_name + if storage_location is not None: body['storage_location'] = storage_location + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', '/api/2.1/unity-catalog/models', body=body, headers=headers) + return RegisteredModelInfo.from_dict(res) + + def delete(self, full_name: str): + """Delete a Registered Model. + + Deletes a registered model and all its model versions from the specified parent catalog and schema. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + + + """ + + headers = {} + self._api.do('DELETE', f'/api/2.1/unity-catalog/models/{full_name}', headers=headers) + + def delete_alias(self, full_name: str, alias: str): + """Delete a Registered Model Alias. + + Deletes a registered model alias. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param alias: str + The name of the alias + + + """ + + headers = {} + self._api.do('DELETE', f'/api/2.1/unity-catalog/models/{full_name}/aliases/{alias}', headers=headers) + + def get(self, full_name: str) -> RegisteredModelInfo: + """Get a Registered Model. + + Get a registered model. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + + :returns: :class:`RegisteredModelInfo` + """ + + headers = {'Accept': 'application/json', } + res = self._api.do('GET', f'/api/2.1/unity-catalog/models/{full_name}', headers=headers) + return RegisteredModelInfo.from_dict(res) + + def list(self, + *, + catalog_name: Optional[str] = None, + max_results: Optional[int] = None, + page_token: Optional[str] = None, + schema_name: Optional[str] = None) -> Iterator[RegisteredModelInfo]: + """List Registered Models. + + List registered models. You can list registered models under a particular schema, or list all + registered models in the current metastore. + + The returned models are filtered based on the privileges of the calling user. For example, the + metastore admin is able to list all the registered models. A regular user needs to be the owner or + have the **EXECUTE** privilege on the registered model to recieve the registered models in the + response. For the latter case, the caller must also be the owner or have the **USE_CATALOG** privilege + on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + There is no guarantee of a specific ordering of the elements in the response. + + :param catalog_name: str (optional) + The identifier of the catalog under which to list registered models. If specified, schema_name must + be specified. + :param max_results: int (optional) + Max number of registered models to return. If catalog and schema are unspecified, max_results must + be specified. If max_results is unspecified, we return all results, starting from the page specified + by page_token. + :param page_token: str (optional) + Opaque token to send for the next page of results (pagination). + :param schema_name: str (optional) + The identifier of the schema under which to list registered models. If specified, catalog_name must + be specified. + + :returns: Iterator over :class:`RegisteredModelInfo` + """ + + query = {} + if catalog_name is not None: query['catalog_name'] = catalog_name + if max_results is not None: query['max_results'] = max_results + if page_token is not None: query['page_token'] = page_token + if schema_name is not None: query['schema_name'] = schema_name + headers = {'Accept': 'application/json', } + + while True: + json = self._api.do('GET', '/api/2.1/unity-catalog/models', query=query, headers=headers) + if 'registered_models' not in json or not json['registered_models']: + return + for v in json['registered_models']: + yield RegisteredModelInfo.from_dict(v) + if 'next_page_token' not in json or not json['next_page_token']: + return + query['page_token'] = json['next_page_token'] + + def set_alias(self, full_name: str, alias: str, version_num: int) -> RegisteredModelAlias: + """Set a Registered Model Alias. + + Set an alias on the specified registered model. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + Full name of the registered model + :param alias: str + The name of the alias + :param version_num: int + The version number of the model version to which the alias points + + :returns: :class:`RegisteredModelAlias` + """ + body = {} + if version_num is not None: body['version_num'] = version_num + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('PUT', + f'/api/2.1/unity-catalog/models/{full_name}/aliases/{alias}', + body=body, + headers=headers) + return RegisteredModelAlias.from_dict(res) + + def update(self, + full_name: str, + *, + comment: Optional[str] = None, + name: Optional[str] = None, + owner: Optional[str] = None) -> RegisteredModelInfo: + """Update a Registered Model. + + Updates the specified registered model. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + Currently only the name, the owner or the comment of the registered model can be updated. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param comment: str (optional) + The comment attached to the registered model + :param name: str (optional) + The name of the registered model + :param owner: str (optional) + The identifier of the user who owns the registered model + + :returns: :class:`RegisteredModelInfo` + """ + body = {} + if comment is not None: body['comment'] = comment + if name is not None: body['name'] = name + if owner is not None: body['owner'] = owner + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('PATCH', f'/api/2.1/unity-catalog/models/{full_name}', body=body, headers=headers) + return RegisteredModelInfo.from_dict(res) + + class 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 @@ -4471,60 +5045,6 @@ def update(self, return SchemaInfo.from_dict(res) -class SecurableTagsAPI: - """Tags are attributes containing keys and values that can be applied to different entities in Unity Catalog. - Tags are useful for organizing and categorizing different entities within a metastore. SecurableTags are - attached to Unity Catalog securable entities.""" - - def __init__(self, api_client): - self._api = api_client - - def list(self, securable_type: ListSecurableType, full_name: str) -> Iterator[TagSecurableAssignment]: - """Get tags for a securable. - - Gets tag assignments for an entity. The caller must be either the owner of the securable, or a - metastore admin, or have at least USE / SELECT privilege on the associated securable. - - :param securable_type: :class:`ListSecurableType` - The type of the unity catalog securable entity. - :param full_name: str - The fully qualified name of the unity catalog securable entity. - - :returns: Iterator over :class:`TagSecurableAssignment` - """ - - headers = {'Accept': 'application/json', } - json = self._api.do('GET', - f'/api/2.1/unity-catalog/securable-tags/{securable_type.value}/{full_name}', - headers=headers) - return [TagSecurableAssignment.from_dict(v) for v in json.get('tag_assignments', [])] - - def update(self, changes: TagChanges, securable_type: UpdateSecurableType, - full_name: str) -> TagSecurableAssignmentsList: - """Update tags for a securable. - - Update tag assignments for an entity The caller must be either the owner of the securable, or a - metastore admin, or have at least USE / SELECT and APPLY_TAG privilege on the associated securable. - - :param changes: :class:`TagChanges` - Desired changes to be made to the tag assignments on the entity - :param securable_type: :class:`UpdateSecurableType` - The type of the unity catalog securable entity. - :param full_name: str - The fully qualified name of the unity catalog securable entity. - - :returns: :class:`TagSecurableAssignmentsList` - """ - body = {} - if changes is not None: body['changes'] = changes.as_dict() - headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } - res = self._api.do('PATCH', - f'/api/2.1/unity-catalog/securable-tags/{securable_type.value}/{full_name}', - body=body, - headers=headers) - return TagSecurableAssignmentsList.from_dict(res) - - class 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 @@ -4771,69 +5291,6 @@ def validate(self, return ValidateStorageCredentialResponse.from_dict(res) -class SubentityTagsAPI: - """Tags are attributes containing keys and values that can be applied to different entities in Unity Catalog. - Tags are useful for organizing and categorizing different entities within a metastore. SubentityTags are - attached to Unity Catalog subentities.""" - - def __init__(self, api_client): - self._api = api_client - - def list(self, securable_type: ListSecurableType, full_name: str, - subentity_name: str) -> Iterator[TagsSubentityAssignment]: - """Get tags for a subentity. - - Gets tag assignments for a subentity associated with a securable entity. Eg. column of a table The - caller must be either the owner of the securable, or a metastore admin, or have at least USE / SELECT - privilege on the associated securable. - - :param securable_type: :class:`ListSecurableType` - The type of the unity catalog securable entity. - :param full_name: str - The fully qualified name of the unity catalog securable entity. - :param subentity_name: str - The name of subentity associated with the securable entity - - :returns: Iterator over :class:`TagsSubentityAssignment` - """ - - headers = {'Accept': 'application/json', } - json = self._api.do( - 'GET', - f'/api/2.1/unity-catalog/subentity-tags/{securable_type.value}/{full_name}/{subentity_name}', - headers=headers) - return [TagsSubentityAssignment.from_dict(v) for v in json.get('tag_assignments', [])] - - def update(self, changes: TagChanges, securable_type: UpdateSecurableType, full_name: str, - subentity_name: str) -> TagSubentityAssignmentsList: - """Update tags for a subentity. - - Update tag assignments for a subentity associated with a securable entity. The caller must be either - the owner of the securable, or a metastore admin, or have at least USE / SELECT and APPLY_TAG - privilege on the associated securable. - - :param changes: :class:`TagChanges` - Desired changes to be made to the tag assignments on the entity - :param securable_type: :class:`UpdateSecurableType` - The type of the unity catalog securable entity. - :param full_name: str - The fully qualified name of the unity catalog securable entity. - :param subentity_name: str - The name of subentity associated with the securable entity - - :returns: :class:`TagSubentityAssignmentsList` - """ - body = {} - if changes is not None: body['changes'] = changes.as_dict() - headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } - res = self._api.do( - 'PATCH', - f'/api/2.1/unity-catalog/subentity-tags/{securable_type.value}/{full_name}/{subentity_name}', - body=body, - headers=headers) - return TagSubentityAssignmentsList.from_dict(res) - - class 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.""" diff --git a/databricks/sdk/service/compute.py b/databricks/sdk/service/compute.py index 584f0f709..b3556cb23 100755 --- a/databricks/sdk/service/compute.py +++ b/databricks/sdk/service/compute.py @@ -1137,7 +1137,6 @@ class CreateInstancePool: enable_elastic_disk: Optional[bool] = None gcp_attributes: Optional['InstancePoolGcpAttributes'] = None idle_instance_autotermination_minutes: Optional[int] = None - instance_pool_fleet_attributes: Optional['InstancePoolFleetAttributes'] = None max_capacity: Optional[int] = None min_idle_instances: Optional[int] = None preloaded_docker_images: Optional['List[DockerImage]'] = None @@ -1153,8 +1152,6 @@ def as_dict(self) -> dict: if self.gcp_attributes: body['gcp_attributes'] = self.gcp_attributes.as_dict() if self.idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = self.idle_instance_autotermination_minutes - if self.instance_pool_fleet_attributes: - body['instance_pool_fleet_attributes'] = self.instance_pool_fleet_attributes.as_dict() if self.instance_pool_name is not None: body['instance_pool_name'] = self.instance_pool_name if self.max_capacity is not None: body['max_capacity'] = self.max_capacity if self.min_idle_instances is not None: body['min_idle_instances'] = self.min_idle_instances @@ -1174,8 +1171,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'CreateInstancePool': enable_elastic_disk=d.get('enable_elastic_disk', None), gcp_attributes=_from_dict(d, 'gcp_attributes', InstancePoolGcpAttributes), idle_instance_autotermination_minutes=d.get('idle_instance_autotermination_minutes', None), - instance_pool_fleet_attributes=_from_dict(d, 'instance_pool_fleet_attributes', - InstancePoolFleetAttributes), instance_pool_name=d.get('instance_pool_name', None), max_capacity=d.get('max_capacity', None), min_idle_instances=d.get('min_idle_instances', None), @@ -1592,7 +1587,6 @@ class EditInstancePool: enable_elastic_disk: Optional[bool] = None gcp_attributes: Optional['InstancePoolGcpAttributes'] = None idle_instance_autotermination_minutes: Optional[int] = None - instance_pool_fleet_attributes: Optional['InstancePoolFleetAttributes'] = None max_capacity: Optional[int] = None min_idle_instances: Optional[int] = None preloaded_docker_images: Optional['List[DockerImage]'] = None @@ -1608,8 +1602,6 @@ def as_dict(self) -> dict: if self.gcp_attributes: body['gcp_attributes'] = self.gcp_attributes.as_dict() if self.idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = self.idle_instance_autotermination_minutes - if self.instance_pool_fleet_attributes: - body['instance_pool_fleet_attributes'] = self.instance_pool_fleet_attributes.as_dict() if self.instance_pool_id is not None: body['instance_pool_id'] = self.instance_pool_id if self.instance_pool_name is not None: body['instance_pool_name'] = self.instance_pool_name if self.max_capacity is not None: body['max_capacity'] = self.max_capacity @@ -1630,8 +1622,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'EditInstancePool': enable_elastic_disk=d.get('enable_elastic_disk', None), gcp_attributes=_from_dict(d, 'gcp_attributes', InstancePoolGcpAttributes), idle_instance_autotermination_minutes=d.get('idle_instance_autotermination_minutes', None), - instance_pool_fleet_attributes=_from_dict(d, 'instance_pool_fleet_attributes', - InstancePoolFleetAttributes), instance_pool_id=d.get('instance_pool_id', None), instance_pool_name=d.get('instance_pool_name', None), max_capacity=d.get('max_capacity', None), @@ -1782,89 +1772,6 @@ class EventType(Enum): UPSIZE_COMPLETED = 'UPSIZE_COMPLETED' -@dataclass -class FleetLaunchTemplateOverride: - availability_zone: str - instance_type: str - max_price: Optional[float] = None - priority: Optional[float] = None - - def as_dict(self) -> dict: - body = {} - if self.availability_zone is not None: body['availability_zone'] = self.availability_zone - if self.instance_type is not None: body['instance_type'] = self.instance_type - if self.max_price is not None: body['max_price'] = self.max_price - if self.priority is not None: body['priority'] = self.priority - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'FleetLaunchTemplateOverride': - return cls(availability_zone=d.get('availability_zone', None), - instance_type=d.get('instance_type', None), - max_price=d.get('max_price', None), - priority=d.get('priority', None)) - - -@dataclass -class FleetOnDemandOption: - allocation_strategy: Optional['FleetOnDemandOptionAllocationStrategy'] = None - max_total_price: Optional[float] = None - use_capacity_reservations_first: Optional[bool] = None - - def as_dict(self) -> dict: - body = {} - if self.allocation_strategy is not None: body['allocation_strategy'] = self.allocation_strategy.value - if self.max_total_price is not None: body['max_total_price'] = self.max_total_price - if self.use_capacity_reservations_first is not None: - body['use_capacity_reservations_first'] = self.use_capacity_reservations_first - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'FleetOnDemandOption': - return cls(allocation_strategy=_enum(d, 'allocation_strategy', FleetOnDemandOptionAllocationStrategy), - max_total_price=d.get('max_total_price', None), - use_capacity_reservations_first=d.get('use_capacity_reservations_first', None)) - - -class FleetOnDemandOptionAllocationStrategy(Enum): - """Only lowest-price and prioritized are allowed""" - - CAPACITY_OPTIMIZED = 'CAPACITY_OPTIMIZED' - DIVERSIFIED = 'DIVERSIFIED' - LOWEST_PRICE = 'LOWEST_PRICE' - PRIORITIZED = 'PRIORITIZED' - - -@dataclass -class FleetSpotOption: - allocation_strategy: Optional['FleetSpotOptionAllocationStrategy'] = None - instance_pools_to_use_count: Optional[int] = None - max_total_price: Optional[float] = None - - def as_dict(self) -> dict: - body = {} - if self.allocation_strategy is not None: body['allocation_strategy'] = self.allocation_strategy.value - if self.instance_pools_to_use_count is not None: - body['instance_pools_to_use_count'] = self.instance_pools_to_use_count - if self.max_total_price is not None: body['max_total_price'] = self.max_total_price - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'FleetSpotOption': - return cls(allocation_strategy=_enum(d, 'allocation_strategy', FleetSpotOptionAllocationStrategy), - instance_pools_to_use_count=d.get('instance_pools_to_use_count', None), - max_total_price=d.get('max_total_price', None)) - - -class FleetSpotOptionAllocationStrategy(Enum): - """lowest-price | diversified | capacity-optimized""" - - CAPACITY_OPTIMIZED = 'CAPACITY_OPTIMIZED' - DIVERSIFIED = 'DIVERSIFIED' - LOWEST_PRICE = 'LOWEST_PRICE' - PRIORITIZED = 'PRIORITIZED' - - @dataclass class GcpAttributes: availability: Optional['GcpAvailability'] = None @@ -1996,7 +1903,6 @@ class GetInstancePool: enable_elastic_disk: Optional[bool] = None gcp_attributes: Optional['InstancePoolGcpAttributes'] = None idle_instance_autotermination_minutes: Optional[int] = None - instance_pool_fleet_attributes: Optional['InstancePoolFleetAttributes'] = None instance_pool_name: Optional[str] = None max_capacity: Optional[int] = None min_idle_instances: Optional[int] = None @@ -2018,8 +1924,6 @@ def as_dict(self) -> dict: if self.gcp_attributes: body['gcp_attributes'] = self.gcp_attributes.as_dict() if self.idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = self.idle_instance_autotermination_minutes - if self.instance_pool_fleet_attributes: - body['instance_pool_fleet_attributes'] = self.instance_pool_fleet_attributes.as_dict() if self.instance_pool_id is not None: body['instance_pool_id'] = self.instance_pool_id if self.instance_pool_name is not None: body['instance_pool_name'] = self.instance_pool_name if self.max_capacity is not None: body['max_capacity'] = self.max_capacity @@ -2044,8 +1948,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'GetInstancePool': enable_elastic_disk=d.get('enable_elastic_disk', None), gcp_attributes=_from_dict(d, 'gcp_attributes', InstancePoolGcpAttributes), idle_instance_autotermination_minutes=d.get('idle_instance_autotermination_minutes', None), - instance_pool_fleet_attributes=_from_dict(d, 'instance_pool_fleet_attributes', - InstancePoolFleetAttributes), instance_pool_id=d.get('instance_pool_id', None), instance_pool_name=d.get('instance_pool_name', None), max_capacity=d.get('max_capacity', None), @@ -2305,7 +2207,6 @@ class InstancePoolAndStats: enable_elastic_disk: Optional[bool] = None gcp_attributes: Optional['InstancePoolGcpAttributes'] = None idle_instance_autotermination_minutes: Optional[int] = None - instance_pool_fleet_attributes: Optional['InstancePoolFleetAttributes'] = None instance_pool_id: Optional[str] = None instance_pool_name: Optional[str] = None max_capacity: Optional[int] = None @@ -2328,8 +2229,6 @@ def as_dict(self) -> dict: if self.gcp_attributes: body['gcp_attributes'] = self.gcp_attributes.as_dict() if self.idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = self.idle_instance_autotermination_minutes - if self.instance_pool_fleet_attributes: - body['instance_pool_fleet_attributes'] = self.instance_pool_fleet_attributes.as_dict() if self.instance_pool_id is not None: body['instance_pool_id'] = self.instance_pool_id if self.instance_pool_name is not None: body['instance_pool_name'] = self.instance_pool_name if self.max_capacity is not None: body['max_capacity'] = self.max_capacity @@ -2354,8 +2253,6 @@ def from_dict(cls, d: Dict[str, any]) -> 'InstancePoolAndStats': enable_elastic_disk=d.get('enable_elastic_disk', None), gcp_attributes=_from_dict(d, 'gcp_attributes', InstancePoolGcpAttributes), idle_instance_autotermination_minutes=d.get('idle_instance_autotermination_minutes', None), - instance_pool_fleet_attributes=_from_dict(d, 'instance_pool_fleet_attributes', - InstancePoolFleetAttributes), instance_pool_id=d.get('instance_pool_id', None), instance_pool_name=d.get('instance_pool_name', None), max_capacity=d.get('max_capacity', None), @@ -2426,28 +2323,6 @@ class InstancePoolAzureAttributesAvailability(Enum): SPOT_WITH_FALLBACK_AZURE = 'SPOT_WITH_FALLBACK_AZURE' -@dataclass -class InstancePoolFleetAttributes: - fleet_on_demand_option: Optional['FleetOnDemandOption'] = None - fleet_spot_option: Optional['FleetSpotOption'] = None - launch_template_overrides: Optional['List[FleetLaunchTemplateOverride]'] = None - - def as_dict(self) -> dict: - body = {} - if self.fleet_on_demand_option: body['fleet_on_demand_option'] = self.fleet_on_demand_option.as_dict() - if self.fleet_spot_option: body['fleet_spot_option'] = self.fleet_spot_option.as_dict() - if self.launch_template_overrides: - body['launch_template_overrides'] = [v.as_dict() for v in self.launch_template_overrides] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'InstancePoolFleetAttributes': - return cls(fleet_on_demand_option=_from_dict(d, 'fleet_on_demand_option', FleetOnDemandOption), - fleet_spot_option=_from_dict(d, 'fleet_spot_option', FleetSpotOption), - launch_template_overrides=_repeated(d, 'launch_template_overrides', - FleetLaunchTemplateOverride)) - - @dataclass class InstancePoolGcpAttributes: gcp_availability: Optional['GcpAvailability'] = None @@ -5244,7 +5119,6 @@ def create(self, enable_elastic_disk: Optional[bool] = None, gcp_attributes: Optional[InstancePoolGcpAttributes] = None, idle_instance_autotermination_minutes: Optional[int] = None, - instance_pool_fleet_attributes: Optional[InstancePoolFleetAttributes] = None, max_capacity: Optional[int] = None, min_idle_instances: Optional[int] = None, preloaded_docker_images: Optional[List[DockerImage]] = None, @@ -5287,8 +5161,6 @@ def create(self, will be automatically terminated after a default timeout. If specified, the threshold must be between 0 and 10000 minutes. Users can also set this value to 0 to instantly remove idle instances from the cache if min cache size could still hold. - :param instance_pool_fleet_attributes: :class:`InstancePoolFleetAttributes` (optional) - The fleet related setting to power the instance pool. :param max_capacity: int (optional) Maximum number of outstanding instances to keep in the pool, including both instances used by clusters and idle instances. Clusters that require further instance provisioning will fail during @@ -5313,8 +5185,6 @@ def create(self, if gcp_attributes is not None: body['gcp_attributes'] = gcp_attributes.as_dict() if idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = idle_instance_autotermination_minutes - if instance_pool_fleet_attributes is not None: - body['instance_pool_fleet_attributes'] = instance_pool_fleet_attributes.as_dict() if instance_pool_name is not None: body['instance_pool_name'] = instance_pool_name if max_capacity is not None: body['max_capacity'] = max_capacity if min_idle_instances is not None: body['min_idle_instances'] = min_idle_instances @@ -5354,7 +5224,6 @@ def edit(self, enable_elastic_disk: Optional[bool] = None, gcp_attributes: Optional[InstancePoolGcpAttributes] = None, idle_instance_autotermination_minutes: Optional[int] = None, - instance_pool_fleet_attributes: Optional[InstancePoolFleetAttributes] = None, max_capacity: Optional[int] = None, min_idle_instances: Optional[int] = None, preloaded_docker_images: Optional[List[DockerImage]] = None, @@ -5399,8 +5268,6 @@ def edit(self, will be automatically terminated after a default timeout. If specified, the threshold must be between 0 and 10000 minutes. Users can also set this value to 0 to instantly remove idle instances from the cache if min cache size could still hold. - :param instance_pool_fleet_attributes: :class:`InstancePoolFleetAttributes` (optional) - The fleet related setting to power the instance pool. :param max_capacity: int (optional) Maximum number of outstanding instances to keep in the pool, including both instances used by clusters and idle instances. Clusters that require further instance provisioning will fail during @@ -5425,8 +5292,6 @@ def edit(self, if gcp_attributes is not None: body['gcp_attributes'] = gcp_attributes.as_dict() if idle_instance_autotermination_minutes is not None: body['idle_instance_autotermination_minutes'] = idle_instance_autotermination_minutes - if instance_pool_fleet_attributes is not None: - body['instance_pool_fleet_attributes'] = instance_pool_fleet_attributes.as_dict() if instance_pool_id is not None: body['instance_pool_id'] = instance_pool_id if instance_pool_name is not None: body['instance_pool_name'] = instance_pool_name if max_capacity is not None: body['max_capacity'] = max_capacity @@ -5722,7 +5587,7 @@ def all_cluster_statuses(self) -> ListAllClusterLibraryStatusesResponse: res = self._api.do('GET', '/api/2.0/libraries/all-cluster-statuses', headers=headers) return ListAllClusterLibraryStatusesResponse.from_dict(res) - def cluster_status(self, cluster_id: str) -> ClusterLibraryStatuses: + def cluster_status(self, cluster_id: str) -> Iterator[LibraryFullStatus]: """Get status. Get the status of libraries on a cluster. A status will be available for all libraries installed on @@ -5741,14 +5606,14 @@ def cluster_status(self, cluster_id: str) -> ClusterLibraryStatuses: :param cluster_id: str Unique identifier of the cluster whose status should be retrieved. - :returns: :class:`ClusterLibraryStatuses` + :returns: Iterator over :class:`LibraryFullStatus` """ query = {} if cluster_id is not None: query['cluster_id'] = cluster_id headers = {'Accept': 'application/json', } - res = self._api.do('GET', '/api/2.0/libraries/cluster-status', query=query, headers=headers) - return ClusterLibraryStatuses.from_dict(res) + json = self._api.do('GET', '/api/2.0/libraries/cluster-status', query=query, headers=headers) + return [LibraryFullStatus.from_dict(v) for v in json.get('library_statuses', [])] def install(self, cluster_id: str, libraries: List[Library]): """Add a library. diff --git a/databricks/sdk/service/files.py b/databricks/sdk/service/files.py index 28a824b17..e96263950 100755 --- a/databricks/sdk/service/files.py +++ b/databricks/sdk/service/files.py @@ -2,7 +2,7 @@ import logging from dataclasses import dataclass -from typing import Dict, Iterator, List, Optional +from typing import BinaryIO, Dict, Iterator, List, Optional from ._internal import _repeated @@ -87,6 +87,11 @@ def from_dict(cls, d: Dict[str, any]) -> 'Delete': return cls(path=d.get('path', None), recursive=d.get('recursive', None)) +@dataclass +class DownloadResponse: + contents: Optional[BinaryIO] = None + + @dataclass class FileInfo: file_size: Optional[int] = None @@ -433,3 +438,75 @@ def read(self, path: str, *, length: Optional[int] = None, offset: Optional[int] headers = {'Accept': 'application/json', } res = self._api.do('GET', '/api/2.0/dbfs/read', query=query, headers=headers) return ReadResponse.from_dict(res) + + +class FilesAPI: + """The Files API allows you to read, write, and delete files and directories in Unity Catalog volumes.""" + + def __init__(self, api_client): + self._api = api_client + + def delete(self, file_path: str): + """Delete a file or directory. + + Deletes a file or directory. + + :param file_path: str + The absolute path of the file or directory in DBFS. + + + """ + + headers = {} + self._api.do('DELETE', f'/api/2.0/fs/files/{file_path}', headers=headers) + + def download(self, file_path: str) -> DownloadResponse: + """Download a file. + + Downloads a file of up to 2 GiB. + + :param file_path: str + The absolute path of the file or directory in DBFS. + + :returns: :class:`DownloadResponse` + """ + + headers = {'Accept': 'application/octet-stream', } + res = self._api.do('GET', f'/api/2.0/fs/files/{file_path}', headers=headers, raw=True) + return DownloadResponse(contents=res) + + def get_status(self, path: str) -> FileInfo: + """Get file or directory status. + + Returns the status of a file or directory. + + :param path: str + The absolute path of the file or directory in the Files API, omitting the initial slash. + + :returns: :class:`FileInfo` + """ + + query = {} + if path is not None: query['path'] = path + headers = {'Accept': 'application/json', } + res = self._api.do('GET', '/api/2.0/fs/get-status', query=query, headers=headers) + return FileInfo.from_dict(res) + + def upload(self, file_path: str, contents: BinaryIO, *, overwrite: Optional[bool] = None): + """Upload a file. + + Uploads a file of up to 2 GiB. + + :param file_path: str + The absolute path of the file or directory in DBFS. + :param contents: BinaryIO + :param overwrite: bool (optional) + If true, an existing file will be overwritten. + + + """ + + query = {} + if overwrite is not None: query['overwrite'] = overwrite + headers = {'Content-Type': 'application/octet-stream', } + self._api.do('PUT', f'/api/2.0/fs/files/{file_path}', query=query, headers=headers, data=contents) diff --git a/databricks/sdk/service/ml.py b/databricks/sdk/service/ml.py index 60c3e68e4..65641970a 100755 --- a/databricks/sdk/service/ml.py +++ b/databricks/sdk/service/ml.py @@ -746,31 +746,31 @@ def from_dict(cls, d: Dict[str, any]) -> 'FileInfo': @dataclass -class GetExperimentByNameResponse: - experiment: Optional['Experiment'] = None +class GetExperimentPermissionLevelsResponse: + permission_levels: Optional['List[ExperimentPermissionsDescription]'] = None def as_dict(self) -> dict: body = {} - if self.experiment: body['experiment'] = self.experiment.as_dict() + if self.permission_levels: body['permission_levels'] = [v.as_dict() for v in self.permission_levels] return body @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'GetExperimentByNameResponse': - return cls(experiment=_from_dict(d, 'experiment', Experiment)) + def from_dict(cls, d: Dict[str, any]) -> 'GetExperimentPermissionLevelsResponse': + return cls(permission_levels=_repeated(d, 'permission_levels', ExperimentPermissionsDescription)) @dataclass -class GetExperimentPermissionLevelsResponse: - permission_levels: Optional['List[ExperimentPermissionsDescription]'] = None +class GetExperimentResponse: + experiment: Optional['Experiment'] = None def as_dict(self) -> dict: body = {} - if self.permission_levels: body['permission_levels'] = [v.as_dict() for v in self.permission_levels] + if self.experiment: body['experiment'] = self.experiment.as_dict() return body @classmethod - def from_dict(cls, d: Dict[str, any]) -> 'GetExperimentPermissionLevelsResponse': - return cls(permission_levels=_repeated(d, 'permission_levels', ExperimentPermissionsDescription)) + def from_dict(cls, d: Dict[str, any]) -> 'GetExperimentResponse': + return cls(experiment=_from_dict(d, 'experiment', Experiment)) @dataclass @@ -2563,7 +2563,7 @@ def delete_tag(self, run_id: str, key: str): headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } self._api.do('POST', '/api/2.0/mlflow/runs/delete-tag', body=body, headers=headers) - def get_by_name(self, experiment_name: str) -> GetExperimentByNameResponse: + def get_by_name(self, experiment_name: str) -> GetExperimentResponse: """Get metadata. Gets metadata for an experiment. @@ -2577,16 +2577,16 @@ def get_by_name(self, experiment_name: str) -> GetExperimentByNameResponse: :param experiment_name: str Name of the associated experiment. - :returns: :class:`GetExperimentByNameResponse` + :returns: :class:`GetExperimentResponse` """ query = {} if experiment_name is not None: query['experiment_name'] = experiment_name headers = {'Accept': 'application/json', } res = self._api.do('GET', '/api/2.0/mlflow/experiments/get-by-name', query=query, headers=headers) - return GetExperimentByNameResponse.from_dict(res) + return GetExperimentResponse.from_dict(res) - def get_experiment(self, experiment_id: str) -> Experiment: + def get_experiment(self, experiment_id: str) -> GetExperimentResponse: """Get an experiment. Gets metadata for an experiment. This method works on deleted experiments. @@ -2594,14 +2594,14 @@ def get_experiment(self, experiment_id: str) -> Experiment: :param experiment_id: str ID of the associated experiment. - :returns: :class:`Experiment` + :returns: :class:`GetExperimentResponse` """ query = {} if experiment_id is not None: query['experiment_id'] = experiment_id headers = {'Accept': 'application/json', } res = self._api.do('GET', '/api/2.0/mlflow/experiments/get', query=query, headers=headers) - return Experiment.from_dict(res) + return GetExperimentResponse.from_dict(res) def get_experiment_permission_levels(self, experiment_id: str) -> GetExperimentPermissionLevelsResponse: """Get experiment permission levels. @@ -2641,7 +2641,7 @@ def get_history(self, max_results: Optional[int] = None, page_token: Optional[str] = None, run_id: Optional[str] = None, - run_uuid: Optional[str] = None) -> GetMetricHistoryResponse: + run_uuid: Optional[str] = None) -> Iterator[Metric]: """Get history of a given metric within a run. Gets a list of all values for the specified metric for a given run. @@ -2659,7 +2659,7 @@ def get_history(self, [Deprecated, use run_id instead] ID of the run from which to fetch metric values. This field will be removed in a future MLflow version. - :returns: :class:`GetMetricHistoryResponse` + :returns: Iterator over :class:`Metric` """ query = {} @@ -2669,8 +2669,16 @@ def get_history(self, if run_id is not None: query['run_id'] = run_id if run_uuid is not None: query['run_uuid'] = run_uuid headers = {'Accept': 'application/json', } - res = self._api.do('GET', '/api/2.0/mlflow/metrics/get-history', query=query, headers=headers) - return GetMetricHistoryResponse.from_dict(res) + + while True: + json = self._api.do('GET', '/api/2.0/mlflow/metrics/get-history', query=query, headers=headers) + if 'metrics' not in json or not json['metrics']: + return + for v in json['metrics']: + yield Metric.from_dict(v) + if 'next_page_token' not in json or not json['next_page_token']: + return + query['page_token'] = json['next_page_token'] def get_run(self, run_id: str, *, run_uuid: Optional[str] = None) -> GetRunResponse: """Get a run. diff --git a/databricks/sdk/service/provisioning.py b/databricks/sdk/service/provisioning.py index fff48339e..c1b99daa1 100755 --- a/databricks/sdk/service/provisioning.py +++ b/databricks/sdk/service/provisioning.py @@ -247,6 +247,7 @@ class CreateWorkspaceRequest: 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 @@ -265,6 +266,7 @@ def as_dict(self) -> dict: if self.cloud_resource_container: body['cloud_resource_container'] = self.cloud_resource_container.as_dict() if self.credentials_id is not None: body['credentials_id'] = self.credentials_id + if self.custom_tags: body['custom_tags'] = self.custom_tags if self.deployment_name is not None: body['deployment_name'] = self.deployment_name if self.gcp_managed_network_config: body['gcp_managed_network_config'] = self.gcp_managed_network_config.as_dict() @@ -289,6 +291,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'CreateWorkspaceRequest': cloud=d.get('cloud', None), cloud_resource_container=_from_dict(d, 'cloud_resource_container', CloudResourceContainer), credentials_id=d.get('credentials_id', None), + custom_tags=d.get('custom_tags', None), deployment_name=d.get('deployment_name', None), gcp_managed_network_config=_from_dict(d, 'gcp_managed_network_config', GcpManagedNetworkConfig), @@ -330,6 +333,9 @@ def from_dict(cls, d: Dict[str, any]) -> 'Credential': credentials_name=d.get('credentials_name', None)) +CustomTags = Dict[str, str] + + @dataclass class CustomerFacingGcpCloudResourceContainer: """The general workspace configurations that are specific to Google Cloud.""" @@ -779,6 +785,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'StsRole': class UpdateWorkspaceRequest: 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_id: Optional[str] = None storage_configuration_id: Optional[str] = None @@ -789,6 +796,7 @@ def as_dict(self) -> dict: body = {} if self.aws_region is not None: body['aws_region'] = self.aws_region if self.credentials_id is not None: body['credentials_id'] = self.credentials_id + if self.custom_tags: body['custom_tags'] = self.custom_tags if self.managed_services_customer_managed_key_id is not None: body['managed_services_customer_managed_key_id'] = self.managed_services_customer_managed_key_id if self.network_id is not None: body['network_id'] = self.network_id @@ -803,6 +811,7 @@ def as_dict(self) -> dict: def from_dict(cls, d: Dict[str, any]) -> 'UpdateWorkspaceRequest': return cls(aws_region=d.get('aws_region', None), credentials_id=d.get('credentials_id', None), + custom_tags=d.get('custom_tags', None), managed_services_customer_managed_key_id=d.get('managed_services_customer_managed_key_id', None), network_id=d.get('network_id', None), @@ -910,6 +919,7 @@ class Workspace: cloud_resource_container: Optional['CloudResourceContainer'] = None creation_time: Optional[int] = 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 @@ -934,6 +944,7 @@ def as_dict(self) -> dict: body['cloud_resource_container'] = self.cloud_resource_container.as_dict() if self.creation_time is not None: body['creation_time'] = self.creation_time if self.credentials_id is not None: body['credentials_id'] = self.credentials_id + if self.custom_tags: body['custom_tags'] = self.custom_tags if self.deployment_name is not None: body['deployment_name'] = self.deployment_name if self.gcp_managed_network_config: body['gcp_managed_network_config'] = self.gcp_managed_network_config.as_dict() @@ -964,6 +975,7 @@ def from_dict(cls, d: Dict[str, any]) -> 'Workspace': cloud_resource_container=_from_dict(d, 'cloud_resource_container', CloudResourceContainer), creation_time=d.get('creation_time', None), credentials_id=d.get('credentials_id', None), + custom_tags=d.get('custom_tags', None), deployment_name=d.get('deployment_name', None), gcp_managed_network_config=_from_dict(d, 'gcp_managed_network_config', GcpManagedNetworkConfig), @@ -1778,6 +1790,7 @@ def create(self, 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, @@ -1810,6 +1823,10 @@ def create(self, The general workspace configurations that are specific to cloud providers. :param credentials_id: str (optional) ID of the workspace's credential configuration object. + :param custom_tags: Dict[str,str] (optional) + The custom tags key-value pairing that is attached to this workspace. The key-value pair is a string + of utf-8 characters. The value can be an empty string, with maximum length of 255 characters. The + key can be of maximum length of 127 characters, and cannot be empty. :param deployment_name: str (optional) The deployment name defines part of the subdomain for the workspace. The workspace URL for web application and REST APIs is `.cloud.databricks.com`. For example, if the @@ -1894,6 +1911,7 @@ def create(self, if cloud_resource_container is not None: body['cloud_resource_container'] = cloud_resource_container.as_dict() if credentials_id is not None: body['credentials_id'] = credentials_id + if custom_tags is not None: body['custom_tags'] = custom_tags if deployment_name is not None: body['deployment_name'] = deployment_name if gcp_managed_network_config is not None: body['gcp_managed_network_config'] = gcp_managed_network_config.as_dict() @@ -1926,6 +1944,7 @@ def create_and_wait( 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, @@ -1941,6 +1960,7 @@ def create_and_wait( 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, @@ -2022,6 +2042,7 @@ def update(self, *, 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_id: Optional[str] = None, storage_configuration_id: Optional[str] = None, @@ -2047,7 +2068,8 @@ def update(self, for workspace storage. - Private access settings ID to add PrivateLink support. You can add or update the private access settings ID to upgrade a workspace to add support for front-end, back-end, or both types of connectivity. You cannot remove (downgrade) any existing front-end or back-end PrivateLink - support on a workspace. + support on a workspace. - Custom tags. Given you provide an empty custom tags, the update would not be + applied. After calling the `PATCH` operation to update the workspace configuration, make repeated `GET` requests with the workspace ID and check the workspace status. The workspace is successful if the @@ -2079,7 +2101,7 @@ def update(self, storage. - Private access settings ID to add PrivateLink support. You can add or update the private access settings ID to upgrade a workspace to add support for front-end, back-end, or both types of connectivity. You cannot remove (downgrade) any existing front-end or back-end PrivateLink support on - a workspace. + a workspace. - Custom tags. Given you provide an empty custom tags, the update would not be applied. **Important**: To update a running workspace, your workspace must have no running compute resources that run in your workspace's VPC in the Classic data plane. For example, stop all all-purpose @@ -2127,6 +2149,10 @@ def update(self, :param credentials_id: str (optional) ID of the workspace's credential configuration object. This parameter is available for updating both failed and running workspaces. + :param custom_tags: Dict[str,str] (optional) + The custom tags key-value pairing that is attached to this workspace. The key-value pair is a string + of utf-8 characters. The value can be an empty string, with maximum length of 255 characters. The + key can be of maximum length of 127 characters, and cannot be empty. :param managed_services_customer_managed_key_id: str (optional) The ID of the workspace's managed services encryption key configuration object. This parameter is available only for updating failed workspaces. @@ -2148,6 +2174,7 @@ def update(self, body = {} if aws_region is not None: body['aws_region'] = aws_region if credentials_id is not None: body['credentials_id'] = credentials_id + if custom_tags is not None: body['custom_tags'] = custom_tags if managed_services_customer_managed_key_id is not None: body['managed_services_customer_managed_key_id'] = managed_services_customer_managed_key_id if network_id is not None: body['network_id'] = network_id @@ -2167,6 +2194,7 @@ def update_and_wait( *, 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_id: Optional[str] = None, storage_configuration_id: Optional[str] = None, @@ -2174,6 +2202,7 @@ def update_and_wait( 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_id=network_id, storage_configuration_id=storage_configuration_id, diff --git a/databricks/sdk/service/sharing.py b/databricks/sdk/service/sharing.py index 0067658d0..d6dd39b74 100755 --- a/databricks/sdk/service/sharing.py +++ b/databricks/sdk/service/sharing.py @@ -549,6 +549,7 @@ class Privilege(Enum): CREATE_FUNCTION = 'CREATE_FUNCTION' CREATE_MANAGED_STORAGE = 'CREATE_MANAGED_STORAGE' CREATE_MATERIALIZED_VIEW = 'CREATE_MATERIALIZED_VIEW' + CREATE_MODEL = 'CREATE_MODEL' CREATE_PROVIDER = 'CREATE_PROVIDER' CREATE_RECIPIENT = 'CREATE_RECIPIENT' CREATE_SCHEMA = 'CREATE_SCHEMA' diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py index a41d5ac04..88db53895 100755 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/service/sql.py @@ -800,6 +800,8 @@ class ExecuteStatementRequest: disposition: Optional['Disposition'] = None format: Optional['Format'] = None on_wait_timeout: Optional['TimeoutAction'] = None + parameters: Optional['List[StatementParameterListItem]'] = None + row_limit: Optional[int] = None schema: Optional[str] = None statement: Optional[str] = None wait_timeout: Optional[str] = None @@ -812,6 +814,8 @@ def as_dict(self) -> dict: if self.disposition is not None: body['disposition'] = self.disposition.value if self.format is not None: body['format'] = self.format.value if self.on_wait_timeout is not None: body['on_wait_timeout'] = self.on_wait_timeout.value + if self.parameters: body['parameters'] = [v.as_dict() for v in self.parameters] + if self.row_limit is not None: body['row_limit'] = self.row_limit if self.schema is not None: body['schema'] = self.schema if self.statement is not None: body['statement'] = self.statement if self.wait_timeout is not None: body['wait_timeout'] = self.wait_timeout @@ -825,6 +829,8 @@ def from_dict(cls, d: Dict[str, any]) -> 'ExecuteStatementRequest': disposition=_enum(d, 'disposition', Disposition), format=_enum(d, 'format', Format), on_wait_timeout=_enum(d, 'on_wait_timeout', TimeoutAction), + parameters=_repeated(d, 'parameters', StatementParameterListItem), + row_limit=d.get('row_limit', None), schema=d.get('schema', None), statement=d.get('statement', None), wait_timeout=d.get('wait_timeout', None), @@ -1934,6 +1940,24 @@ class State(Enum): STOPPING = 'STOPPING' +@dataclass +class StatementParameterListItem: + name: str + type: Optional[str] = None + value: Optional[str] = None + + def as_dict(self) -> dict: + body = {} + if self.name is not None: body['name'] = self.name + if self.type is not None: body['type'] = self.type + if self.value is not None: body['value'] = self.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'StatementParameterListItem': + return cls(name=d.get('name', None), type=d.get('type', None), value=d.get('value', None)) + + class StatementState(Enum): """Statement execution state: - `PENDING`: waiting for warehouse - `RUNNING`: running - `SUCCEEDED`: execution was successful, result data available for fetch - `FAILED`: execution @@ -3206,6 +3230,8 @@ def execute_statement(self, disposition: Optional[Disposition] = None, format: Optional[Format] = None, on_wait_timeout: Optional[TimeoutAction] = None, + parameters: Optional[List[StatementParameterListItem]] = None, + row_limit: Optional[int] = None, schema: Optional[str] = None, statement: Optional[str] = None, wait_timeout: Optional[str] = None, @@ -3286,6 +3312,36 @@ def execute_statement(self, `CANCEL` → the statement execution is canceled and the call returns immediately with a `CANCELED` state. + :param parameters: List[:class:`StatementParameterListItem`] (optional) + A list of parameters to pass into a SQL statement containing parameter markers. A parameter consists + of a name, a value, and optionally a type. To represent a NULL value, the `value` field may be + omitted. If the `type` field is omitted, the value is interpreted as a string. + + If the type is given, parameters will be checked for type correctness according to the given type. A + value is correct if the provided string can be converted to the requested type using the `cast` + function. The exact semantics are described in the section [`cast` function] of the SQL language + reference. + + For example, the following statement contains two parameters, `my_id` and `my_date`: + + SELECT * FROM my_table WHERE name = :my_name AND date = :my_date + + The parameters can be passed in the request body as follows: + + { ..., "statement": "SELECT * FROM my_table WHERE name = :my_name AND date = :my_date", + "parameters": [ { "name": "my_name", "value": "the name" }, { "name": "my_date", "value": + "2020-01-01", "type": "DATE" } ] } + + Currently, positional parameters denoted by a `?` marker are not supported by the SQL Statement + Execution API. + + Also see the section [Parameter markers] of the SQL language reference. + + [Parameter markers]: https://docs.databricks.com/sql/language-manual/sql-ref-parameter-marker.html + [`cast` function]: https://docs.databricks.com/sql/language-manual/functions/cast.html + :param row_limit: int (optional) + Applies the given row limit to the statement's result set with identical semantics as the SQL + `LIMIT` clause. :param schema: str (optional) Sets default schema for statement execution, similar to [`USE SCHEMA`] in SQL. @@ -3308,6 +3364,8 @@ def execute_statement(self, if disposition is not None: body['disposition'] = disposition.value if format is not None: body['format'] = format.value if on_wait_timeout is not None: body['on_wait_timeout'] = on_wait_timeout.value + if parameters is not None: body['parameters'] = [v.as_dict() for v in parameters] + if row_limit is not None: body['row_limit'] = row_limit if schema is not None: body['schema'] = schema if statement is not None: body['statement'] = statement if wait_timeout is not None: body['wait_timeout'] = wait_timeout diff --git a/docs/account/metastores.rst b/docs/account/metastores.rst index 4625ef887..f2ba17011 100644 --- a/docs/account/metastores.rst +++ b/docs/account/metastores.rst @@ -96,7 +96,7 @@ Account Metastores Gets all Unity Catalog metastores associated with an account specified by ID. - :returns: :class:`ListMetastoresResponse` + :returns: Iterator over :class:`MetastoreInfo` .. py:method:: update(metastore_id [, metastore_info]) diff --git a/docs/account/storage_credentials.rst b/docs/account/storage_credentials.rst index 32d6b5c10..b114f763b 100644 --- a/docs/account/storage_credentials.rst +++ b/docs/account/storage_credentials.rst @@ -114,7 +114,7 @@ Account Storage Credentials :param metastore_id: str Unity Catalog metastore ID - :returns: :class:`ListStorageCredentialsResponse` + :returns: Iterator over :class:`StorageCredentialInfo` .. py:method:: update(metastore_id, name [, credential_info]) diff --git a/docs/account/workspaces.rst b/docs/account/workspaces.rst index 7225f776c..a5d741a22 100644 --- a/docs/account/workspaces.rst +++ b/docs/account/workspaces.rst @@ -9,7 +9,7 @@ Workspaces 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. - .. py:method:: create(workspace_name [, aws_region, cloud, cloud_resource_container, credentials_id, deployment_name, gcp_managed_network_config, gke_config, location, managed_services_customer_managed_key_id, network_id, pricing_tier, private_access_settings_id, storage_configuration_id, storage_customer_managed_key_id]) + .. py:method:: create(workspace_name [, aws_region, cloud, cloud_resource_container, credentials_id, custom_tags, deployment_name, gcp_managed_network_config, gke_config, location, managed_services_customer_managed_key_id, network_id, pricing_tier, private_access_settings_id, storage_configuration_id, storage_customer_managed_key_id]) Usage: @@ -64,6 +64,10 @@ Workspaces The general workspace configurations that are specific to cloud providers. :param credentials_id: str (optional) ID of the workspace's credential configuration object. + :param custom_tags: Dict[str,str] (optional) + The custom tags key-value pairing that is attached to this workspace. The key-value pair is a string + of utf-8 characters. The value can be an empty string, with maximum length of 255 characters. The + key can be of maximum length of 127 characters, and cannot be empty. :param deployment_name: str (optional) The deployment name defines part of the subdomain for the workspace. The workspace URL for web application and REST APIs is `.cloud.databricks.com`. For example, if the @@ -238,7 +242,7 @@ Workspaces :returns: Iterator over :class:`Workspace` - .. py:method:: update(workspace_id [, aws_region, credentials_id, managed_services_customer_managed_key_id, network_id, storage_configuration_id, storage_customer_managed_key_id]) + .. py:method:: update(workspace_id [, aws_region, credentials_id, custom_tags, managed_services_customer_managed_key_id, network_id, storage_configuration_id, storage_customer_managed_key_id]) Usage: @@ -300,7 +304,8 @@ Workspaces for workspace storage. - Private access settings ID to add PrivateLink support. You can add or update the private access settings ID to upgrade a workspace to add support for front-end, back-end, or both types of connectivity. You cannot remove (downgrade) any existing front-end or back-end PrivateLink - support on a workspace. + support on a workspace. - Custom tags. Given you provide an empty custom tags, the update would not be + applied. After calling the `PATCH` operation to update the workspace configuration, make repeated `GET` requests with the workspace ID and check the workspace status. The workspace is successful if the @@ -332,7 +337,7 @@ Workspaces storage. - Private access settings ID to add PrivateLink support. You can add or update the private access settings ID to upgrade a workspace to add support for front-end, back-end, or both types of connectivity. You cannot remove (downgrade) any existing front-end or back-end PrivateLink support on - a workspace. + a workspace. - Custom tags. Given you provide an empty custom tags, the update would not be applied. **Important**: To update a running workspace, your workspace must have no running compute resources that run in your workspace's VPC in the Classic data plane. For example, stop all all-purpose @@ -380,6 +385,10 @@ Workspaces :param credentials_id: str (optional) ID of the workspace's credential configuration object. This parameter is available for updating both failed and running workspaces. + :param custom_tags: Dict[str,str] (optional) + The custom tags key-value pairing that is attached to this workspace. The key-value pair is a string + of utf-8 characters. The value can be an empty string, with maximum length of 255 characters. The + key can be of maximum length of 127 characters, and cannot be empty. :param managed_services_customer_managed_key_id: str (optional) The ID of the workspace's managed services encryption key configuration object. This parameter is available only for updating failed workspaces. diff --git a/docs/workspace/catalogs.rst b/docs/workspace/catalogs.rst index a9818cc49..dae181a51 100644 --- a/docs/workspace/catalogs.rst +++ b/docs/workspace/catalogs.rst @@ -9,7 +9,7 @@ Catalogs 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. - .. py:method:: create(name [, comment, connection_name, properties, provider_name, share_name, storage_root]) + .. py:method:: create(name [, comment, connection_name, options, properties, provider_name, share_name, storage_root]) Usage: @@ -37,6 +37,8 @@ Catalogs User-provided free-form text description. :param connection_name: str (optional) The name of the connection to an external data source. + :param options: Dict[str,str] (optional) + A map of key-value properties attached to the securable. :param properties: Dict[str,str] (optional) A map of key-value properties attached to the securable. :param provider_name: str (optional) @@ -118,7 +120,7 @@ Catalogs :returns: Iterator over :class:`CatalogInfo` - .. py:method:: update(name [, comment, isolation_mode, owner, properties]) + .. py:method:: update(name [, comment, isolation_mode, options, owner, properties]) Usage: @@ -148,6 +150,8 @@ Catalogs User-provided free-form text description. :param isolation_mode: :class:`IsolationMode` (optional) Whether the current securable is accessible from all workspaces or a specific set of workspaces. + :param options: Dict[str,str] (optional) + A map of key-value properties attached to the securable. :param owner: str (optional) Username of current owner of catalog. :param properties: Dict[str,str] (optional) diff --git a/docs/workspace/experiments.rst b/docs/workspace/experiments.rst index 5ae3e6776..7ea3b6234 100644 --- a/docs/workspace/experiments.rst +++ b/docs/workspace/experiments.rst @@ -163,7 +163,7 @@ Experiments :param experiment_name: str Name of the associated experiment. - :returns: :class:`GetExperimentByNameResponse` + :returns: :class:`GetExperimentResponse` .. py:method:: get_experiment(experiment_id) @@ -192,7 +192,7 @@ Experiments :param experiment_id: str ID of the associated experiment. - :returns: :class:`Experiment` + :returns: :class:`GetExperimentResponse` .. py:method:: get_experiment_permission_levels(experiment_id) @@ -238,7 +238,7 @@ Experiments [Deprecated, use run_id instead] ID of the run from which to fetch metric values. This field will be removed in a future MLflow version. - :returns: :class:`GetMetricHistoryResponse` + :returns: Iterator over :class:`Metric` .. py:method:: get_run(run_id [, run_uuid]) diff --git a/docs/workspace/instance_pools.rst b/docs/workspace/instance_pools.rst index 5c8d2c95a..fcc649664 100644 --- a/docs/workspace/instance_pools.rst +++ b/docs/workspace/instance_pools.rst @@ -17,7 +17,7 @@ Instance Pools Databricks does not charge DBUs while instances are idle in the pool. Instance provider billing does apply. See pricing. - .. py:method:: create(instance_pool_name, node_type_id [, aws_attributes, azure_attributes, custom_tags, disk_spec, enable_elastic_disk, gcp_attributes, idle_instance_autotermination_minutes, instance_pool_fleet_attributes, max_capacity, min_idle_instances, preloaded_docker_images, preloaded_spark_versions]) + .. py:method:: create(instance_pool_name, node_type_id [, aws_attributes, azure_attributes, custom_tags, disk_spec, enable_elastic_disk, gcp_attributes, idle_instance_autotermination_minutes, max_capacity, min_idle_instances, preloaded_docker_images, preloaded_spark_versions]) Usage: @@ -74,8 +74,6 @@ Instance Pools will be automatically terminated after a default timeout. If specified, the threshold must be between 0 and 10000 minutes. Users can also set this value to 0 to instantly remove idle instances from the cache if min cache size could still hold. - :param instance_pool_fleet_attributes: :class:`InstancePoolFleetAttributes` (optional) - The fleet related setting to power the instance pool. :param max_capacity: int (optional) Maximum number of outstanding instances to keep in the pool, including both instances used by clusters and idle instances. Clusters that require further instance provisioning will fail during @@ -104,7 +102,7 @@ Instance Pools - .. py:method:: edit(instance_pool_id, instance_pool_name, node_type_id [, aws_attributes, azure_attributes, custom_tags, disk_spec, enable_elastic_disk, gcp_attributes, idle_instance_autotermination_minutes, instance_pool_fleet_attributes, max_capacity, min_idle_instances, preloaded_docker_images, preloaded_spark_versions]) + .. py:method:: edit(instance_pool_id, instance_pool_name, node_type_id [, aws_attributes, azure_attributes, custom_tags, disk_spec, enable_elastic_disk, gcp_attributes, idle_instance_autotermination_minutes, max_capacity, min_idle_instances, preloaded_docker_images, preloaded_spark_versions]) Usage: @@ -167,8 +165,6 @@ Instance Pools will be automatically terminated after a default timeout. If specified, the threshold must be between 0 and 10000 minutes. Users can also set this value to 0 to instantly remove idle instances from the cache if min cache size could still hold. - :param instance_pool_fleet_attributes: :class:`InstancePoolFleetAttributes` (optional) - The fleet related setting to power the instance pool. :param max_capacity: int (optional) Maximum number of outstanding instances to keep in the pool, including both instances used by clusters and idle instances. Clusters that require further instance provisioning will fail during diff --git a/docs/workspace/libraries.rst b/docs/workspace/libraries.rst index ea6d9e456..500ef263c 100644 --- a/docs/workspace/libraries.rst +++ b/docs/workspace/libraries.rst @@ -50,7 +50,7 @@ ManagedLibraries :param cluster_id: str Unique identifier of the cluster whose status should be retrieved. - :returns: :class:`ClusterLibraryStatuses` + :returns: Iterator over :class:`LibraryFullStatus` .. py:method:: install(cluster_id, libraries) diff --git a/docs/workspace/model_versions.rst b/docs/workspace/model_versions.rst new file mode 100644 index 000000000..4c51eb84a --- /dev/null +++ b/docs/workspace/model_versions.rst @@ -0,0 +1,112 @@ +Model Versions +============== +.. py:class:: 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). + + .. py:method:: delete(full_name, version) + + Delete a Model Version. + + Deletes a model version from the specified registered model. Any aliases assigned to the model version + will also be deleted. + + The caller must be a metastore admin or an owner of the parent registered model. For the latter case, + the caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + + + + + .. py:method:: get(full_name, version) + + Get a Model Version. + + Get a model version. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the parent + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + + :returns: :class:`RegisteredModelInfo` + + + .. py:method:: get_by_alias(full_name, alias) + + Get Model Version By Alias. + + Get a model version by alias. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param alias: str + The name of the alias + + :returns: :class:`ModelVersionInfo` + + + .. py:method:: list(full_name [, max_results, page_token]) + + List Model Versions. + + List model versions. You can list model versions under a particular schema, or list all model versions + in the current metastore. + + The returned models are filtered based on the privileges of the calling user. For example, the + metastore admin is able to list all the model versions. A regular user needs to be the owner or have + the **EXECUTE** privilege on the parent registered model to recieve the model versions in the + response. For the latter case, the caller must also be the owner or have the **USE_CATALOG** privilege + on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + There is no guarantee of a specific ordering of the elements in the response. + + :param full_name: str + The full three-level name of the registered model under which to list model versions + :param max_results: int (optional) + Max number of model versions to return + :param page_token: str (optional) + Opaque token to send for the next page of results (pagination). + + :returns: Iterator over :class:`ModelVersionInfo` + + + .. py:method:: update(full_name, version [, comment]) + + Update a Model Version. + + Updates the specified model version. + + The caller must be a metastore admin or an owner of the parent registered model. For the latter case, + the caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + Currently only the comment of the model version can be updated. + + :param full_name: str + The three-level (fully qualified) name of the model version + :param version: int + The integer version number of the model version + :param comment: str (optional) + The comment attached to the model version + + :returns: :class:`ModelVersionInfo` + \ No newline at end of file diff --git a/docs/workspace/registered_models.rst b/docs/workspace/registered_models.rst new file mode 100644 index 000000000..c546734fd --- /dev/null +++ b/docs/workspace/registered_models.rst @@ -0,0 +1,181 @@ +Registered Models +================= +.. py:class:: 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. + + .. py:method:: create(catalog_name, schema_name, name [, comment, storage_location]) + + Create a Registered Model. + + Creates a new registered model in Unity Catalog. + + File storage for model versions in the registered model will be located in the default location which + is specified by the parent schema, or the parent catalog, or the Metastore. + + For registered model creation to succeed, the user must satisfy the following conditions: - The caller + must be a metastore admin, or be the owner of the parent catalog and schema, or have the + **USE_CATALOG** privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + - The caller must have the **CREATE MODEL** or **CREATE FUNCTION** privilege on the parent schema. + + :param catalog_name: str + The name of the catalog where the schema and the registered model reside + :param schema_name: str + The name of the schema where the registered model resides + :param name: str + The name of the registered model + :param comment: str (optional) + The comment attached to the registered model + :param storage_location: str (optional) + The storage location on the cloud under which model version data files are stored + + :returns: :class:`RegisteredModelInfo` + + + .. py:method:: delete(full_name) + + Delete a Registered Model. + + Deletes a registered model and all its model versions from the specified parent catalog and schema. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + + + + + .. py:method:: delete_alias(full_name, alias) + + Delete a Registered Model Alias. + + Deletes a registered model alias. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param alias: str + The name of the alias + + + + + .. py:method:: get(full_name) + + Get a Registered Model. + + Get a registered model. + + The caller must be a metastore admin or an owner of (or have the **EXECUTE** privilege on) the + registered model. For the latter case, the caller must also be the owner or have the **USE_CATALOG** + privilege on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + The three-level (fully qualified) name of the registered model + + :returns: :class:`RegisteredModelInfo` + + + .. py:method:: list( [, catalog_name, max_results, page_token, schema_name]) + + List Registered Models. + + List registered models. You can list registered models under a particular schema, or list all + registered models in the current metastore. + + The returned models are filtered based on the privileges of the calling user. For example, the + metastore admin is able to list all the registered models. A regular user needs to be the owner or + have the **EXECUTE** privilege on the registered model to recieve the registered models in the + response. For the latter case, the caller must also be the owner or have the **USE_CATALOG** privilege + on the parent catalog and the **USE_SCHEMA** privilege on the parent schema. + + There is no guarantee of a specific ordering of the elements in the response. + + :param catalog_name: str (optional) + The identifier of the catalog under which to list registered models. If specified, schema_name must + be specified. + :param max_results: int (optional) + Max number of registered models to return. If catalog and schema are unspecified, max_results must + be specified. If max_results is unspecified, we return all results, starting from the page specified + by page_token. + :param page_token: str (optional) + Opaque token to send for the next page of results (pagination). + :param schema_name: str (optional) + The identifier of the schema under which to list registered models. If specified, catalog_name must + be specified. + + :returns: Iterator over :class:`RegisteredModelInfo` + + + .. py:method:: set_alias(full_name, alias, version_num) + + Set a Registered Model Alias. + + Set an alias on the specified registered model. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + :param full_name: str + Full name of the registered model + :param alias: str + The name of the alias + :param version_num: int + The version number of the model version to which the alias points + + :returns: :class:`RegisteredModelAlias` + + + .. py:method:: update(full_name [, comment, name, owner]) + + Update a Registered Model. + + Updates the specified registered model. + + The caller must be a metastore admin or an owner of the registered model. For the latter case, the + caller must also be the owner or have the **USE_CATALOG** privilege on the parent catalog and the + **USE_SCHEMA** privilege on the parent schema. + + Currently only the name, the owner or the comment of the registered model can be updated. + + :param full_name: str + The three-level (fully qualified) name of the registered model + :param comment: str (optional) + The comment attached to the registered model + :param name: str (optional) + The name of the registered model + :param owner: str (optional) + The identifier of the user who owns the registered model + + :returns: :class:`RegisteredModelInfo` + \ No newline at end of file diff --git a/docs/workspace/statement_execution.rst b/docs/workspace/statement_execution.rst index 67ce0dd5d..fcad43273 100644 --- a/docs/workspace/statement_execution.rst +++ b/docs/workspace/statement_execution.rst @@ -170,7 +170,7 @@ Statement Execution - .. py:method:: execute_statement( [, byte_limit, catalog, disposition, format, on_wait_timeout, schema, statement, wait_timeout, warehouse_id]) + .. py:method:: execute_statement( [, byte_limit, catalog, disposition, format, on_wait_timeout, parameters, row_limit, schema, statement, wait_timeout, warehouse_id]) Execute a SQL statement. @@ -248,6 +248,36 @@ Statement Execution `CANCEL` → the statement execution is canceled and the call returns immediately with a `CANCELED` state. + :param parameters: List[:class:`StatementParameterListItem`] (optional) + A list of parameters to pass into a SQL statement containing parameter markers. A parameter consists + of a name, a value, and optionally a type. To represent a NULL value, the `value` field may be + omitted. If the `type` field is omitted, the value is interpreted as a string. + + If the type is given, parameters will be checked for type correctness according to the given type. A + value is correct if the provided string can be converted to the requested type using the `cast` + function. The exact semantics are described in the section [`cast` function] of the SQL language + reference. + + For example, the following statement contains two parameters, `my_id` and `my_date`: + + SELECT * FROM my_table WHERE name = :my_name AND date = :my_date + + The parameters can be passed in the request body as follows: + + { ..., "statement": "SELECT * FROM my_table WHERE name = :my_name AND date = :my_date", + "parameters": [ { "name": "my_name", "value": "the name" }, { "name": "my_date", "value": + "2020-01-01", "type": "DATE" } ] } + + Currently, positional parameters denoted by a `?` marker are not supported by the SQL Statement + Execution API. + + Also see the section [Parameter markers] of the SQL language reference. + + [Parameter markers]: https://docs.databricks.com/sql/language-manual/sql-ref-parameter-marker.html + [`cast` function]: https://docs.databricks.com/sql/language-manual/functions/cast.html + :param row_limit: int (optional) + Applies the given row limit to the statement's result set with identical semantics as the SQL + `LIMIT` clause. :param schema: str (optional) Sets default schema for statement execution, similar to [`USE SCHEMA`] in SQL. diff --git a/docs/workspace/workspace-catalog.rst b/docs/workspace/workspace-catalog.rst index ccdb80ce9..9a0342e31 100644 --- a/docs/workspace/workspace-catalog.rst +++ b/docs/workspace/workspace-catalog.rst @@ -14,10 +14,10 @@ Configure data governance with Unity Catalog for metastores, catalogs, schemas, functions grants metastores + model_versions + registered_models schemas - securable_tags storage_credentials - subentity_tags system_schemas table_constraints tables diff --git a/tests/integration/test_files.py b/tests/integration/test_files.py index 632d987ba..537e55581 100644 --- a/tests/integration/test_files.py +++ b/tests/integration/test_files.py @@ -1,10 +1,11 @@ import io import pathlib -from typing import List +from typing import Callable, List import pytest from databricks.sdk.core import DatabricksError +from databricks.sdk.service.catalog import VolumeType def test_local_io(random): @@ -183,13 +184,46 @@ def test_dbfs_upload_download(w, random, junk, tmp_path): assert f.read() == b"some text data" -def test_files_api_upload_download(w, random): - pytest.skip() - f = io.BytesIO(b"some text data") - target_file = f'/Volumes/bogdanghita/default/v3_shared/sdk-testing/{random(10)}.txt' - w.files.upload(target_file, f) +class ResourceWithCleanup: + cleanup: Callable[[], None] - with w.files.download(target_file) as f: - assert f.read() == b"some text data" + def __init__(self, cleanup): + self.cleanup = cleanup + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.cleanup() + + @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)) + + @staticmethod + def create_volume(w, catalog, schema, volume): + res = w.volumes.create(catalog_name=catalog, + schema_name=schema, + name=volume, + volume_type=VolumeType.MANAGED) + return ResourceWithCleanup(lambda: w.volumes.delete(res.full_name)) + + +def test_files_api_upload_download(ucws, random): + w = ucws + schema = 'filesit-' + random() + volume = 'filesit-' + random() + with ResourceWithCleanup.create_schema(w, 'main', schema): + with ResourceWithCleanup.create_volume(w, 'main', schema, volume): + f = io.BytesIO(b"some text data") + target_file = f'/Volumes/main/{schema}/{volume}/filesit-{random()}.txt' + w.files.upload(target_file, f) + + res = w.files.get_status(target_file) + assert not res.is_dir + + with w.files.download(target_file).contents as f: + assert f.read() == b"some text data" - w.files.delete(target_file) + w.files.delete(target_file)