From 11d4f357c86bd2b647c6801674dc3a1077a14ff6 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi <88379306+tanmay-db@users.noreply.github.com> Date: Wed, 10 Jan 2024 21:31:48 +0530 Subject: [PATCH] Update SDK to Latest OpenAPI Specification (#501) ## Changes Generating SDK for latest OpenAPI specification ## Tests Unit tests (will run integration test on Release PR) - [ ] `make test` run locally - [ ] `make fmt` applied - [ ] relevant integration tests applied --- .codegen/_openapi_sha | 2 +- .gitattributes | 2 + databricks/sdk/__init__.py | 21 + databricks/sdk/service/catalog.py | 247 +++- databricks/sdk/service/dashboards.py | 75 + databricks/sdk/service/iam.py | 1 + databricks/sdk/service/jobs.py | 11 +- databricks/sdk/service/pipelines.py | 30 +- databricks/sdk/service/settings.py | 43 +- databricks/sdk/service/sharing.py | 10 +- databricks/sdk/service/sql.py | 161 ++- databricks/sdk/service/vectorsearch.py | 1206 +++++++++++++++++ .../list_external_locations_on_aws.py | 3 +- .../list_storage_credentials_on_aws.py | 3 +- 14 files changed, 1708 insertions(+), 107 deletions(-) create mode 100755 databricks/sdk/service/dashboards.py create mode 100755 databricks/sdk/service/vectorsearch.py diff --git a/.codegen/_openapi_sha b/.codegen/_openapi_sha index 2c3fb6e1..c9ce6cc2 100644 --- a/.codegen/_openapi_sha +++ b/.codegen/_openapi_sha @@ -1 +1 @@ -d3853c8dee5806d04da2ae8910f273ffb35719a5 \ No newline at end of file +0e0d4cbe87193e36c73b8b2be3b0dd0f1b013e00 \ No newline at end of file diff --git a/.gitattributes b/.gitattributes index c28d21d3..dd1547c9 100755 --- a/.gitattributes +++ b/.gitattributes @@ -3,6 +3,7 @@ databricks/sdk/errors/mapping.py linguist-generated=true databricks/sdk/service/billing.py linguist-generated=true databricks/sdk/service/catalog.py linguist-generated=true databricks/sdk/service/compute.py linguist-generated=true +databricks/sdk/service/dashboards.py linguist-generated=true databricks/sdk/service/files.py linguist-generated=true databricks/sdk/service/iam.py linguist-generated=true databricks/sdk/service/jobs.py linguist-generated=true @@ -14,6 +15,7 @@ databricks/sdk/service/serving.py linguist-generated=true databricks/sdk/service/settings.py linguist-generated=true databricks/sdk/service/sharing.py linguist-generated=true databricks/sdk/service/sql.py linguist-generated=true +databricks/sdk/service/vectorsearch.py linguist-generated=true databricks/sdk/service/workspace.py linguist-generated=true examples/alerts/create_alerts.py linguist-generated=true examples/alerts/get_alerts.py linguist-generated=true diff --git a/databricks/sdk/__init__.py b/databricks/sdk/__init__.py index 1f380ba0..cecf6d62 100755 --- a/databricks/sdk/__init__.py +++ b/databricks/sdk/__init__.py @@ -24,6 +24,7 @@ InstancePoolsAPI, InstanceProfilesAPI, LibrariesAPI, PolicyFamiliesAPI) +from databricks.sdk.service.dashboards import LakeviewAPI from databricks.sdk.service.files import DbfsAPI, FilesAPI from databricks.sdk.service.iam import (AccountAccessControlAPI, AccountAccessControlProxyAPI, @@ -62,6 +63,8 @@ QueryHistoryAPI, QueryVisualizationsAPI, StatementExecutionAPI, WarehousesAPI) +from databricks.sdk.service.vectorsearch import (VectorSearchEndpointsAPI, + VectorSearchIndexesAPI) from databricks.sdk.service.workspace import (GitCredentialsAPI, ReposAPI, SecretsAPI, WorkspaceAPI) @@ -169,6 +172,7 @@ def __init__(self, self._instance_profiles = InstanceProfilesAPI(self._api_client) self._ip_access_lists = IpAccessListsAPI(self._api_client) self._jobs = JobsAPI(self._api_client) + self._lakeview = LakeviewAPI(self._api_client) self._libraries = LibrariesAPI(self._api_client) self._metastores = MetastoresAPI(self._api_client) self._model_registry = ModelRegistryAPI(self._api_client) @@ -198,6 +202,8 @@ def __init__(self, self._token_management = TokenManagementAPI(self._api_client) self._tokens = TokensAPI(self._api_client) self._users = UsersAPI(self._api_client) + self._vector_search_endpoints = VectorSearchEndpointsAPI(self._api_client) + self._vector_search_indexes = VectorSearchIndexesAPI(self._api_client) self._volumes = VolumesAPI(self._api_client) self._warehouses = WarehousesAPI(self._api_client) self._workspace = WorkspaceExt(self._api_client) @@ -361,6 +367,11 @@ def jobs(self) -> JobsAPI: """The Jobs API allows you to create, edit, and delete jobs.""" return self._jobs + @property + def lakeview(self) -> LakeviewAPI: + """These APIs provide specific management operations for Lakeview dashboards.""" + return self._lakeview + @property def libraries(self) -> LibrariesAPI: """The Libraries API allows you to install and uninstall libraries and get the status of libraries on a cluster.""" @@ -506,6 +517,16 @@ def users(self) -> UsersAPI: """User identities recognized by Databricks and represented by email addresses.""" return self._users + @property + def vector_search_endpoints(self) -> VectorSearchEndpointsAPI: + """**Endpoint**: Represents the compute resources to host vector search indexes.""" + return self._vector_search_endpoints + + @property + def vector_search_indexes(self) -> VectorSearchIndexesAPI: + """**Index**: An efficient representation of your embedding vectors that supports real-time and efficient approximate nearest neighbor (ANN) search queries.""" + return self._vector_search_indexes + @property def volumes(self) -> VolumesAPI: """Volumes are a Unity Catalog (UC) capability for accessing, storing, governing, organizing and processing files.""" diff --git a/databricks/sdk/service/catalog.py b/databricks/sdk/service/catalog.py index 6e51eb64..243723ec 100755 --- a/databricks/sdk/service/catalog.py +++ b/databricks/sdk/service/catalog.py @@ -2242,17 +2242,23 @@ class ListExternalLocationsResponse: external_locations: Optional[List[ExternalLocationInfo]] = None """An array of external locations.""" + next_page_token: Optional[str] = None + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" + def as_dict(self) -> dict: """Serializes the ListExternalLocationsResponse into a dictionary suitable for use as a JSON request body.""" body = {} if self.external_locations: body['external_locations'] = [v.as_dict() for v in self.external_locations] + 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]) -> ListExternalLocationsResponse: """Deserializes the ListExternalLocationsResponse from a dictionary.""" - return cls(external_locations=_repeated_dict(d, 'external_locations', ExternalLocationInfo)) + return cls(external_locations=_repeated_dict(d, 'external_locations', ExternalLocationInfo), + next_page_token=d.get('next_page_token', None)) @dataclass @@ -2260,16 +2266,22 @@ class ListFunctionsResponse: functions: Optional[List[FunctionInfo]] = None """An array of function information objects.""" + next_page_token: Optional[str] = None + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" + def as_dict(self) -> dict: """Serializes the ListFunctionsResponse into a dictionary suitable for use as a JSON request body.""" body = {} if self.functions: body['functions'] = [v.as_dict() for v in self.functions] + 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]) -> ListFunctionsResponse: """Deserializes the ListFunctionsResponse from a dictionary.""" - return cls(functions=_repeated_dict(d, 'functions', FunctionInfo)) + return cls(functions=_repeated_dict(d, 'functions', FunctionInfo), + next_page_token=d.get('next_page_token', None)) @dataclass @@ -2294,7 +2306,8 @@ class ListModelVersionsResponse: model_versions: Optional[List[ModelVersionInfo]] = None next_page_token: Optional[str] = None - """Token to retrieve the next page of results""" + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" def as_dict(self) -> dict: """Serializes the ListModelVersionsResponse into a dictionary suitable for use as a JSON request body.""" @@ -2334,28 +2347,39 @@ def from_dict(cls, d: Dict[str, any]) -> ListRegisteredModelsResponse: @dataclass class ListSchemasResponse: + next_page_token: Optional[str] = None + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" + schemas: Optional[List[SchemaInfo]] = None """An array of schema information objects.""" def as_dict(self) -> dict: """Serializes the ListSchemasResponse into a dictionary suitable for use as a JSON request body.""" body = {} + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token if self.schemas: body['schemas'] = [v.as_dict() for v in self.schemas] return body @classmethod def from_dict(cls, d: Dict[str, any]) -> ListSchemasResponse: """Deserializes the ListSchemasResponse from a dictionary.""" - return cls(schemas=_repeated_dict(d, 'schemas', SchemaInfo)) + return cls(next_page_token=d.get('next_page_token', None), + schemas=_repeated_dict(d, 'schemas', SchemaInfo)) @dataclass class ListStorageCredentialsResponse: + next_page_token: Optional[str] = None + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" + storage_credentials: Optional[List[StorageCredentialInfo]] = None def as_dict(self) -> dict: """Serializes the ListStorageCredentialsResponse into a dictionary suitable for use as a JSON request body.""" body = {} + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token if self.storage_credentials: body['storage_credentials'] = [v.as_dict() for v in self.storage_credentials] return body @@ -2363,7 +2387,8 @@ def as_dict(self) -> dict: @classmethod def from_dict(cls, d: Dict[str, any]) -> ListStorageCredentialsResponse: """Deserializes the ListStorageCredentialsResponse from a dictionary.""" - return cls(storage_credentials=_repeated_dict(d, 'storage_credentials', StorageCredentialInfo)) + return cls(next_page_token=d.get('next_page_token', None), + storage_credentials=_repeated_dict(d, 'storage_credentials', StorageCredentialInfo)) @dataclass @@ -2386,7 +2411,8 @@ def from_dict(cls, d: Dict[str, any]) -> ListSystemSchemasResponse: @dataclass class ListTableSummariesResponse: next_page_token: Optional[str] = None - """Opaque token for pagination. Omitted if there are no more results.""" + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" tables: Optional[List[TableSummary]] = None """List of table summaries.""" @@ -2408,8 +2434,8 @@ def from_dict(cls, d: Dict[str, any]) -> ListTableSummariesResponse: @dataclass class ListTablesResponse: next_page_token: Optional[str] = None - """Opaque token for pagination. Omitted if there are no more results. page_token should be set to - this value for fetching the next page.""" + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" tables: Optional[List[TableInfo]] = None """An array of table information objects.""" @@ -3340,23 +3366,6 @@ def from_dict(cls, d: Dict[str, any]) -> TableConstraint: primary_key_constraint=_from_dict(d, 'primary_key_constraint', PrimaryKeyConstraint)) -@dataclass -class TableConstraintList: - table_constraints: Optional[List[TableConstraint]] = None - """List of table constraints. Note: this field is not set in the output of the __listTables__ API.""" - - def as_dict(self) -> dict: - """Serializes the TableConstraintList into a dictionary suitable for use as a JSON request body.""" - body = {} - if self.table_constraints: body['table_constraints'] = [v.as_dict() for v in self.table_constraints] - return body - - @classmethod - def from_dict(cls, d: Dict[str, any]) -> TableConstraintList: - """Deserializes the TableConstraintList from a dictionary.""" - return cls(table_constraints=_repeated_dict(d, 'table_constraints', TableConstraint)) - - @dataclass class TableDependency: """A table that is dependent on a SQL object.""" @@ -3451,7 +3460,8 @@ class TableInfo: storage_location: Optional[str] = None """Storage root URL for table (for **MANAGED**, **EXTERNAL** tables)""" - table_constraints: Optional[TableConstraintList] = None + table_constraints: Optional[List[TableConstraint]] = None + """List of table constraints. Note: this field is not set in the output of the __listTables__ API.""" table_id: Optional[str] = None """Name of table, relative to parent schema.""" @@ -3508,7 +3518,7 @@ def as_dict(self) -> dict: if self.storage_credential_name is not None: body['storage_credential_name'] = self.storage_credential_name if self.storage_location is not None: body['storage_location'] = self.storage_location - if self.table_constraints: body['table_constraints'] = self.table_constraints.as_dict() + if self.table_constraints: body['table_constraints'] = [v.as_dict() for v in self.table_constraints] if self.table_id is not None: body['table_id'] = self.table_id if self.table_type is not None: body['table_type'] = self.table_type.value if self.updated_at is not None: body['updated_at'] = self.updated_at @@ -3547,7 +3557,7 @@ def from_dict(cls, d: Dict[str, any]) -> TableInfo: sql_path=d.get('sql_path', None), storage_credential_name=d.get('storage_credential_name', None), storage_location=d.get('storage_location', None), - table_constraints=_from_dict(d, 'table_constraints', TableConstraintList), + table_constraints=_repeated_dict(d, 'table_constraints', TableConstraint), table_id=d.get('table_id', None), table_type=_enum(d, 'table_type', TableType), updated_at=d.get('updated_at', None), @@ -5239,20 +5249,45 @@ def get(self, name: str) -> ExternalLocationInfo: res = self._api.do('GET', f'/api/2.1/unity-catalog/external-locations/{name}', headers=headers) return ExternalLocationInfo.from_dict(res) - def list(self) -> Iterator[ExternalLocationInfo]: + def list(self, + *, + max_results: Optional[int] = None, + page_token: Optional[str] = None) -> Iterator[ExternalLocationInfo]: """List external locations. Gets an array of external locations (__ExternalLocationInfo__ objects) from the metastore. The caller must be a metastore admin, the owner of the external location, or a user that has some privilege on - the external location. There is no guarantee of a specific ordering of the elements in the array. + the external location. For unpaginated request, there is no guarantee of a specific ordering of the + elements in the array. For paginated request, elements are ordered by their name. + + :param max_results: int (optional) + Maximum number of external locations to return. If not set, all the external locations are returned + (not recommended). - when set to a value greater than 0, the page length is the minimum of this + value and a server configured value; - when set to 0, the page length is set to a server configured + value (recommended); - when set to a value less than 0, an invalid parameter error is returned; + :param page_token: str (optional) + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`ExternalLocationInfo` """ + 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', } - json = self._api.do('GET', '/api/2.1/unity-catalog/external-locations', headers=headers) - parsed = ListExternalLocationsResponse.from_dict(json).external_locations - return parsed if parsed is not None else [] + + while True: + json = self._api.do('GET', + '/api/2.1/unity-catalog/external-locations', + query=query, + headers=headers) + if 'external_locations' not in json or not json['external_locations']: + return + for v in json['external_locations']: + yield ExternalLocationInfo.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, name: str, @@ -5391,30 +5426,52 @@ def get(self, name: str) -> FunctionInfo: res = self._api.do('GET', f'/api/2.1/unity-catalog/functions/{name}', headers=headers) return FunctionInfo.from_dict(res) - def list(self, catalog_name: str, schema_name: str) -> Iterator[FunctionInfo]: + def list(self, + catalog_name: str, + schema_name: str, + *, + max_results: Optional[int] = None, + page_token: Optional[str] = None) -> Iterator[FunctionInfo]: """List functions. List functions within the specified parent catalog and schema. If the user is a metastore admin, all functions are returned in the output list. Otherwise, the user must have the **USE_CATALOG** privilege on the catalog and the **USE_SCHEMA** privilege on the schema, and the output list contains only - functions for which either the user has the **EXECUTE** privilege or the user is the owner. There is - no guarantee of a specific ordering of the elements in the array. + functions for which either the user has the **EXECUTE** privilege or the user is the owner. For + unpaginated request, there is no guarantee of a specific ordering of the elements in the array. For + paginated request, elements are ordered by their name. :param catalog_name: str Name of parent catalog for functions of interest. :param schema_name: str Parent schema of functions. + :param max_results: int (optional) + Maximum number of functions to return. If not set, all the functions are returned (not recommended). + - when set to a value greater than 0, the page length is the minimum of this value and a server + configured value; - when set to 0, the page length is set to a server configured value + (recommended); - when set to a value less than 0, an invalid parameter error is returned; + :param page_token: str (optional) + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`FunctionInfo` """ 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', } - json = self._api.do('GET', '/api/2.1/unity-catalog/functions', query=query, headers=headers) - parsed = ListFunctionsResponse.from_dict(json).functions - return parsed if parsed is not None else [] + + while True: + json = self._api.do('GET', '/api/2.1/unity-catalog/functions', query=query, headers=headers) + if 'functions' not in json or not json['functions']: + return + for v in json['functions']: + yield FunctionInfo.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, name: str, *, owner: Optional[str] = None) -> FunctionInfo: """Update a function. @@ -5892,9 +5949,13 @@ def list(self, :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 + Maximum number of model versions to return. If not set, the page length is set to a server + configured value (100, as of 1/3/2024). - when set to a value greater than 0, the page length is the + minimum of this value and a server configured value(1000, as of 1/3/2024); - when set to 0, the page + length is set to a server configured value (100, as of 1/3/2024) (recommended); - when set to a + value less than 0, an invalid parameter error is returned; :param page_token: str (optional) - Opaque token to send for the next page of results (pagination). + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`ModelVersionInfo` """ @@ -6269,26 +6330,47 @@ def get(self, full_name: str) -> SchemaInfo: res = self._api.do('GET', f'/api/2.1/unity-catalog/schemas/{full_name}', headers=headers) return SchemaInfo.from_dict(res) - def list(self, catalog_name: str) -> Iterator[SchemaInfo]: + def list(self, + catalog_name: str, + *, + max_results: Optional[int] = None, + page_token: Optional[str] = None) -> Iterator[SchemaInfo]: """List schemas. Gets an array of schemas for a catalog in the metastore. If the caller is the metastore admin or the owner of the parent catalog, all schemas for the catalog will be retrieved. Otherwise, only schemas - owned by the caller (or for which the caller has the **USE_SCHEMA** privilege) will be retrieved. - There is no guarantee of a specific ordering of the elements in the array. + owned by the caller (or for which the caller has the **USE_SCHEMA** privilege) will be retrieved. For + unpaginated request, there is no guarantee of a specific ordering of the elements in the array. For + paginated request, elements are ordered by their name. :param catalog_name: str Parent catalog for schemas of interest. + :param max_results: int (optional) + Maximum number of schemas to return. If not set, all the schemas are returned (not recommended). - + when set to a value greater than 0, the page length is the minimum of this value and a server + configured value; - when set to 0, the page length is set to a server configured value + (recommended); - when set to a value less than 0, an invalid parameter error is returned; + :param page_token: str (optional) + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`SchemaInfo` """ 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 headers = {'Accept': 'application/json', } - json = self._api.do('GET', '/api/2.1/unity-catalog/schemas', query=query, headers=headers) - parsed = ListSchemasResponse.from_dict(json).schemas - return parsed if parsed is not None else [] + + while True: + json = self._api.do('GET', '/api/2.1/unity-catalog/schemas', query=query, headers=headers) + if 'schemas' not in json or not json['schemas']: + return + for v in json['schemas']: + yield SchemaInfo.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, @@ -6442,21 +6524,47 @@ def get(self, name: str) -> StorageCredentialInfo: res = self._api.do('GET', f'/api/2.1/unity-catalog/storage-credentials/{name}', headers=headers) return StorageCredentialInfo.from_dict(res) - def list(self) -> Iterator[StorageCredentialInfo]: + def list(self, + *, + max_results: Optional[int] = None, + page_token: Optional[str] = None) -> Iterator[StorageCredentialInfo]: """List credentials. Gets an array of storage credentials (as __StorageCredentialInfo__ objects). The array is limited to only those storage credentials the caller has permission to access. If the caller is a metastore - admin, all storage credentials will be retrieved. There is no guarantee of a specific ordering of the - elements in the array. + admin, retrieval of credentials is unrestricted. For unpaginated request, there is no guarantee of a + specific ordering of the elements in the array. For paginated request, elements are ordered by their + name. + + :param max_results: int (optional) + Maximum number of storage credentials to return. If not set, all the storage credentials are + returned (not recommended). - when set to a value greater than 0, the page length is the minimum of + this value and a server configured value; - when set to 0, the page length is set to a server + configured value (recommended); - when set to a value less than 0, an invalid parameter error is + returned; + :param page_token: str (optional) + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`StorageCredentialInfo` """ + 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', } - json = self._api.do('GET', '/api/2.1/unity-catalog/storage-credentials', headers=headers) - parsed = ListStorageCredentialsResponse.from_dict(json).storage_credentials - return parsed if parsed is not None else [] + + while True: + json = self._api.do('GET', + '/api/2.1/unity-catalog/storage-credentials', + query=query, + headers=headers) + if 'storage_credentials' not in json or not json['storage_credentials']: + return + for v in json['storage_credentials']: + yield StorageCredentialInfo.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, name: str, @@ -6788,6 +6896,8 @@ def list(self, *, include_delta_metadata: Optional[bool] = None, max_results: Optional[int] = None, + omit_columns: Optional[bool] = None, + omit_properties: Optional[bool] = None, page_token: Optional[str] = None) -> Iterator[TableInfo]: """List tables. @@ -6804,11 +6914,14 @@ def list(self, :param include_delta_metadata: bool (optional) Whether delta metadata should be included in the response. :param max_results: int (optional) - Maximum number of tables to return (page length). If not set, all accessible tables in the schema - are returned. If set to: - - * greater than 0, page length is the minimum of this value and a server configured value. * equal to - 0, page length is set to a server configured value. * lesser than 0, invalid parameter error. + Maximum number of tables to return. If not set, all the tables are returned (not recommended). - + when set to a value greater than 0, the page length is the minimum of this value and a server + configured value; - when set to 0, the page length is set to a server configured value + (recommended); - when set to a value less than 0, an invalid parameter error is returned; + :param omit_columns: bool (optional) + Whether to omit the columns of the table from the response or not. + :param omit_properties: bool (optional) + Whether to omit the properties of the table from the response or not. :param page_token: str (optional) Opaque token to send for the next page of results (pagination). @@ -6819,6 +6932,8 @@ def list(self, if catalog_name is not None: query['catalog_name'] = catalog_name if include_delta_metadata is not None: query['include_delta_metadata'] = include_delta_metadata if max_results is not None: query['max_results'] = max_results + if omit_columns is not None: query['omit_columns'] = omit_columns + if omit_properties is not None: query['omit_properties'] = omit_properties 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', } @@ -6845,10 +6960,10 @@ def list_summaries(self, Gets an array of summaries for tables for a schema and catalog within the metastore. The table summaries returned are either: - * summaries for all tables (within the current metastore and parent catalog and schema), when the user - is a metastore admin, or: * summaries for all tables and schemas (within the current metastore and - parent catalog) for which the user has ownership or the **SELECT** privilege on the table and - ownership or **USE_SCHEMA** privilege on the schema, provided that the user also has ownership or the + * summaries for tables (within the current metastore and parent catalog and schema), when the user is + a metastore admin, or: * summaries for tables and schemas (within the current metastore and parent + catalog) for which the user has ownership or the **SELECT** privilege on the table and ownership or + **USE_SCHEMA** privilege on the schema, provided that the user also has ownership or the **USE_CATALOG** privilege on the parent catalog. There is no guarantee of a specific ordering of the elements in the array. @@ -6856,9 +6971,13 @@ def list_summaries(self, :param catalog_name: str Name of parent catalog for tables of interest. :param max_results: int (optional) - Maximum number of tables to return (page length). Defaults to 10000. + Maximum number of summaries for tables to return. If not set, the page length is set to a server + configured value (10000, as of 1/5/2024). - when set to a value greater than 0, the page length is + the minimum of this value and a server configured value (10000, as of 1/5/2024); - when set to 0, + the page length is set to a server configured value (10000, as of 1/5/2024) (recommended); - when + set to a value less than 0, an invalid parameter error is returned; :param page_token: str (optional) - Opaque token to send for the next page of results (pagination). + Opaque pagination token to go to next page based on previous query. :param schema_name_pattern: str (optional) A sql LIKE pattern (% and _) for schema names. All schemas will be returned if not set or empty. :param table_name_pattern: str (optional) diff --git a/databricks/sdk/service/dashboards.py b/databricks/sdk/service/dashboards.py new file mode 100755 index 00000000..51ae22e1 --- /dev/null +++ b/databricks/sdk/service/dashboards.py @@ -0,0 +1,75 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Dict, Optional + +_LOG = logging.getLogger('databricks.sdk') + +# all definitions in this file are in alphabetical order + + +@dataclass +class PublishRequest: + dashboard_id: Optional[str] = None + """UUID identifying the dashboard to be published.""" + + embed_credentials: Optional[bool] = None + """Flag to indicate if the publisher's credentials should be embedded in the published dashboard. + These embedded credentials will be used to execute the published dashboard's queries.""" + + warehouse_id: Optional[str] = None + """The ID of the warehouse that can be used to override the warehouse which was set in the draft.""" + + def as_dict(self) -> dict: + """Serializes the PublishRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dashboard_id is not None: body['dashboard_id'] = self.dashboard_id + if self.embed_credentials is not None: body['embed_credentials'] = self.embed_credentials + if self.warehouse_id is not None: body['warehouse_id'] = self.warehouse_id + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> PublishRequest: + """Deserializes the PublishRequest from a dictionary.""" + return cls(dashboard_id=d.get('dashboard_id', None), + embed_credentials=d.get('embed_credentials', None), + warehouse_id=d.get('warehouse_id', None)) + + +class LakeviewAPI: + """These APIs provide specific management operations for Lakeview dashboards. Generic resource management can + be done with Workspace API (import, export, get-status, list, delete).""" + + def __init__(self, api_client): + self._api = api_client + + def publish(self, + dashboard_id: str, + *, + embed_credentials: Optional[bool] = None, + warehouse_id: Optional[str] = None): + """Publish dashboard. + + Publish the current draft dashboard. + + :param dashboard_id: str + UUID identifying the dashboard to be published. + :param embed_credentials: bool (optional) + Flag to indicate if the publisher's credentials should be embedded in the published dashboard. These + embedded credentials will be used to execute the published dashboard's queries. + :param warehouse_id: str (optional) + The ID of the warehouse that can be used to override the warehouse which was set in the draft. + + + """ + body = {} + if embed_credentials is not None: body['embed_credentials'] = embed_credentials + if warehouse_id is not None: body['warehouse_id'] = warehouse_id + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + self._api.do('POST', + f'/api/2.0/lakeview/dashboards/{dashboard_id}/published', + body=body, + headers=headers) diff --git a/databricks/sdk/service/iam.py b/databricks/sdk/service/iam.py index b1f4dcbf..e162771f 100755 --- a/databricks/sdk/service/iam.py +++ b/databricks/sdk/service/iam.py @@ -1125,6 +1125,7 @@ def from_dict(cls, d: Dict[str, any]) -> User: class UserSchema(Enum): URN_IETF_PARAMS_SCIM_SCHEMAS_CORE_2_0_USER = 'urn:ietf:params:scim:schemas:core:2.0:User' + URN_IETF_PARAMS_SCIM_SCHEMAS_EXTENSION_WORKSPACE_2_0_USER = 'urn:ietf:params:scim:schemas:extension:workspace:2.0:User' class WorkspacePermission(Enum): diff --git a/databricks/sdk/service/jobs.py b/databricks/sdk/service/jobs.py index 25cc245f..7aae7b5b 100755 --- a/databricks/sdk/service/jobs.py +++ b/databricks/sdk/service/jobs.py @@ -5185,9 +5185,10 @@ def repair_run_and_wait( sql_params=sql_params).result(timeout=timeout) def reset(self, job_id: int, new_settings: JobSettings): - """Overwrite all settings for a job. + """Update all job settings (reset). - Overwrite all settings for the given job. Use the Update endpoint to update job settings partially. + Overwrite all settings for the given job. Use the [_Update_ endpoint](:method:jobs/update) to update + job settings partially. :param job_id: int The canonical identifier of the job to reset. This field is required. @@ -5497,10 +5498,10 @@ def update(self, *, fields_to_remove: Optional[List[str]] = None, new_settings: Optional[JobSettings] = None): - """Partially update a job. + """Update job settings partially. - Add, update, or remove specific settings of an existing job. Use the ResetJob to overwrite all job - settings. + Add, update, or remove specific settings of an existing job. Use the [_Reset_ + endpoint](:method:jobs/reset) to overwrite all job settings. :param job_id: int The canonical identifier of the job to update. This field is required. diff --git a/databricks/sdk/service/pipelines.py b/databricks/sdk/service/pipelines.py index 560e7699..c7ae2329 100755 --- a/databricks/sdk/service/pipelines.py +++ b/databricks/sdk/service/pipelines.py @@ -832,6 +832,11 @@ class PipelineCluster: """Attributes related to clusters running on Google Cloud Platform. If not specified at cluster creation, a set of default values will be used.""" + init_scripts: Optional[List[compute.InitScriptInfo]] = None + """The configuration for storing init scripts. Any number of destinations can be specified. The + scripts are executed sequentially in the order provided. If `cluster_log_conf` is specified, + init script logs are sent to `//init_scripts`.""" + instance_pool_id: Optional[str] = None """The optional ID of the instance pool to which the cluster belongs.""" @@ -895,6 +900,7 @@ def as_dict(self) -> dict: body['driver_instance_pool_id'] = self.driver_instance_pool_id if self.driver_node_type_id is not None: body['driver_node_type_id'] = self.driver_node_type_id if self.gcp_attributes: body['gcp_attributes'] = self.gcp_attributes.as_dict() + if self.init_scripts: body['init_scripts'] = [v.as_dict() for v in self.init_scripts] if self.instance_pool_id is not None: body['instance_pool_id'] = self.instance_pool_id if self.label is not None: body['label'] = self.label if self.node_type_id is not None: body['node_type_id'] = self.node_type_id @@ -917,6 +923,7 @@ def from_dict(cls, d: Dict[str, any]) -> PipelineCluster: driver_instance_pool_id=d.get('driver_instance_pool_id', None), driver_node_type_id=d.get('driver_node_type_id', None), gcp_attributes=_from_dict(d, 'gcp_attributes', compute.GcpAttributes), + init_scripts=_repeated_dict(d, 'init_scripts', compute.InitScriptInfo), instance_pool_id=d.get('instance_pool_id', None), label=d.get('label', None), node_type_id=d.get('node_type_id', None), @@ -1402,6 +1409,10 @@ class StartUpdate: full_refresh_selection are empty, this is a full graph update. Full Refresh on a table means that the states of the table will be reset before the refresh.""" + validate_only: Optional[bool] = None + """If true, this update only validates the correctness of pipeline source code but does not + materialize or publish any datasets.""" + def as_dict(self) -> dict: """Serializes the StartUpdate into a dictionary suitable for use as a JSON request body.""" body = {} @@ -1411,6 +1422,7 @@ def as_dict(self) -> dict: body['full_refresh_selection'] = [v for v in self.full_refresh_selection] if self.pipeline_id is not None: body['pipeline_id'] = self.pipeline_id if self.refresh_selection: body['refresh_selection'] = [v for v in self.refresh_selection] + if self.validate_only is not None: body['validate_only'] = self.validate_only return body @classmethod @@ -1420,7 +1432,8 @@ def from_dict(cls, d: Dict[str, any]) -> StartUpdate: full_refresh=d.get('full_refresh', None), full_refresh_selection=d.get('full_refresh_selection', None), pipeline_id=d.get('pipeline_id', None), - refresh_selection=d.get('refresh_selection', None)) + refresh_selection=d.get('refresh_selection', None), + validate_only=d.get('validate_only', None)) class StartUpdateCause(Enum): @@ -1486,6 +1499,10 @@ class UpdateInfo: update_id: Optional[str] = None """The ID of this update.""" + validate_only: Optional[bool] = None + """If true, this update only validates the correctness of pipeline source code but does not + materialize or publish any datasets.""" + def as_dict(self) -> dict: """Serializes the UpdateInfo into a dictionary suitable for use as a JSON request body.""" body = {} @@ -1500,6 +1517,7 @@ def as_dict(self) -> dict: if self.refresh_selection: body['refresh_selection'] = [v for v in self.refresh_selection] if self.state is not None: body['state'] = self.state.value if self.update_id is not None: body['update_id'] = self.update_id + if self.validate_only is not None: body['validate_only'] = self.validate_only return body @classmethod @@ -1514,7 +1532,8 @@ def from_dict(cls, d: Dict[str, any]) -> UpdateInfo: pipeline_id=d.get('pipeline_id', None), refresh_selection=d.get('refresh_selection', None), state=_enum(d, 'state', UpdateInfoState), - update_id=d.get('update_id', None)) + update_id=d.get('update_id', None), + validate_only=d.get('validate_only', None)) class UpdateInfoCause(Enum): @@ -2007,7 +2026,8 @@ def start_update(self, cause: Optional[StartUpdateCause] = None, full_refresh: Optional[bool] = None, full_refresh_selection: Optional[List[str]] = None, - refresh_selection: Optional[List[str]] = None) -> StartUpdateResponse: + refresh_selection: Optional[List[str]] = None, + validate_only: Optional[bool] = None) -> StartUpdateResponse: """Start a pipeline. Starts a new update for the pipeline. If there is already an active update for the pipeline, the @@ -2025,6 +2045,9 @@ def start_update(self, A list of tables to update without fullRefresh. If both refresh_selection and full_refresh_selection are empty, this is a full graph update. Full Refresh on a table means that the states of the table will be reset before the refresh. + :param validate_only: bool (optional) + If true, this update only validates the correctness of pipeline source code but does not materialize + or publish any datasets. :returns: :class:`StartUpdateResponse` """ @@ -2034,6 +2057,7 @@ def start_update(self, if full_refresh_selection is not None: body['full_refresh_selection'] = [v for v in full_refresh_selection] if refresh_selection is not None: body['refresh_selection'] = [v for v in refresh_selection] + if validate_only is not None: body['validate_only'] = validate_only headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } res = self._api.do('POST', f'/api/2.0/pipelines/{pipeline_id}/updates', body=body, headers=headers) return StartUpdateResponse.from_dict(res) diff --git a/databricks/sdk/service/settings.py b/databricks/sdk/service/settings.py index 15bca640..352263cc 100755 --- a/databricks/sdk/service/settings.py +++ b/databricks/sdk/service/settings.py @@ -182,7 +182,7 @@ class CreateTokenRequest: lifetime_seconds: Optional[int] = None """The lifetime of the token, in seconds. - If the ifetime is not specified, this token remains valid indefinitely.""" + If the lifetime is not specified, this token remains valid indefinitely.""" def as_dict(self) -> dict: """Serializes the CreateTokenRequest into a dictionary suitable for use as a JSON request body.""" @@ -582,6 +582,23 @@ def from_dict(cls, d: Dict[str, any]) -> ListNetworkConnectivityConfigurationsRe next_page_token=d.get('next_page_token', None)) +@dataclass +class ListPublicTokensResponse: + token_infos: Optional[List[PublicTokenInfo]] = None + """The information for each token.""" + + def as_dict(self) -> dict: + """Serializes the ListPublicTokensResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.token_infos: body['token_infos'] = [v.as_dict() for v in self.token_infos] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> ListPublicTokensResponse: + """Deserializes the ListPublicTokensResponse from a dictionary.""" + return cls(token_infos=_repeated_dict(d, 'token_infos', PublicTokenInfo)) + + @dataclass class ListTokensResponse: token_infos: Optional[List[TokenInfo]] = None @@ -998,7 +1015,7 @@ class ReplaceIpAccessList: """Specifies whether this IP access list is enabled.""" ip_access_list_id: Optional[str] = None - """The ID for the corresponding IP access list to modify""" + """The ID for the corresponding IP access list""" ip_addresses: Optional[List[str]] = None @@ -1283,7 +1300,7 @@ class UpdateIpAccessList: """Specifies whether this IP access list is enabled.""" ip_access_list_id: Optional[str] = None - """The ID for the corresponding IP access list to modify""" + """The ID for the corresponding IP access list""" ip_addresses: Optional[List[str]] = None @@ -1391,6 +1408,7 @@ def delete(self, ip_access_list_id: str): Deletes an IP access list, specified by its list ID. :param ip_access_list_id: str + The ID for the corresponding IP access list """ @@ -1406,6 +1424,7 @@ def get(self, ip_access_list_id: str) -> GetIpAccessListResponse: Gets an IP access list, specified by its list ID. :param ip_access_list_id: str + The ID for the corresponding IP access list :returns: :class:`GetIpAccessListResponse` """ @@ -1451,6 +1470,7 @@ def replace(self, effect. :param ip_access_list_id: str + The ID for the corresponding IP access list :param label: str Label for the IP access list. This **cannot** be empty. :param list_type: :class:`ListType` @@ -1499,6 +1519,7 @@ def update(self, It can take a few minutes for the changes to take effect. :param ip_access_list_id: str + The ID for the corresponding IP access list :param enabled: bool (optional) Specifies whether this IP access list is enabled. :param ip_addresses: List[str] (optional) @@ -1715,7 +1736,7 @@ def delete(self, ip_access_list_id: str): Deletes an IP access list, specified by its list ID. :param ip_access_list_id: str - The ID for the corresponding IP access list to modify + The ID for the corresponding IP access list """ @@ -1729,7 +1750,7 @@ def get(self, ip_access_list_id: str) -> FetchIpAccessListResponse: Gets an IP access list, specified by its list ID. :param ip_access_list_id: str - The ID for the corresponding IP access list to modify + The ID for the corresponding IP access list :returns: :class:`FetchIpAccessListResponse` """ @@ -1772,7 +1793,7 @@ def replace(self, :method:workspaceconf/setStatus. :param ip_access_list_id: str - The ID for the corresponding IP access list to modify + The ID for the corresponding IP access list :param label: str Label for the IP access list. This **cannot** be empty. :param list_type: :class:`ListType` @@ -1819,7 +1840,7 @@ def update(self, no effect until you enable the feature. See :method:workspaceconf/setStatus. :param ip_access_list_id: str - The ID for the corresponding IP access list to modify + The ID for the corresponding IP access list :param enabled: bool (optional) Specifies whether this IP access list is enabled. :param ip_addresses: List[str] (optional) @@ -2363,7 +2384,7 @@ def create(self, :param lifetime_seconds: int (optional) The lifetime of the token, in seconds. - If the ifetime is not specified, this token remains valid indefinitely. + If the lifetime is not specified, this token remains valid indefinitely. :returns: :class:`CreateTokenResponse` """ @@ -2391,17 +2412,17 @@ def delete(self, token_id: str): headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } self._api.do('POST', '/api/2.0/token/delete', body=body, headers=headers) - def list(self) -> Iterator[TokenInfo]: + def list(self) -> Iterator[PublicTokenInfo]: """List tokens. Lists all the valid tokens for a user-workspace pair. - :returns: Iterator over :class:`TokenInfo` + :returns: Iterator over :class:`PublicTokenInfo` """ headers = {'Accept': 'application/json', } json = self._api.do('GET', '/api/2.0/token/list', headers=headers) - parsed = ListTokensResponse.from_dict(json).token_infos + parsed = ListPublicTokensResponse.from_dict(json).token_infos return parsed if parsed is not None else [] diff --git a/databricks/sdk/service/sharing.py b/databricks/sdk/service/sharing.py index bcc21cf5..37753bdc 100755 --- a/databricks/sdk/service/sharing.py +++ b/databricks/sdk/service/sharing.py @@ -584,7 +584,8 @@ class ListCleanRoomsResponse: """An array of clean rooms. Remote details (central) are not included.""" next_page_token: Optional[str] = None - """Token to retrieve the next page of results. Absent if there are no more pages.""" + """Opaque token to retrieve the next page of results. Absent if there are no more pages. + __page_token__ should be set to this value for the next request (for the next page of results).""" def as_dict(self) -> dict: """Serializes the ListCleanRoomsResponse into a dictionary suitable for use as a JSON request body.""" @@ -1626,9 +1627,12 @@ def list(self, array. :param max_results: int (optional) - Maximum number of clean rooms to return. + Maximum number of clean rooms to return. If not set, all the clean rooms are returned (not + recommended). - when set to a value greater than 0, the page length is the minimum of this value and + a server configured value; - when set to 0, the page length is set to a server configured value + (recommended); - when set to a value less than 0, an invalid parameter error is returned; :param page_token: str (optional) - Pagination token to go to next page based on previous query. + Opaque pagination token to go to next page based on previous query. :returns: Iterator over :class:`CleanRoomInfo` """ diff --git a/databricks/sdk/service/sql.py b/databricks/sdk/service/sql.py index 3e3a6042..e0d0b0ac 100755 --- a/databricks/sdk/service/sql.py +++ b/databricks/sdk/service/sql.py @@ -23,8 +23,8 @@ class AccessControl: group_name: Optional[str] = None permission_level: Optional[PermissionLevel] = None - """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_MANAGE`: Can manage the - query""" + """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_EDIT`: Can edit the query + * `CAN_MANAGE`: Can manage the query""" user_name: Optional[str] = None @@ -681,8 +681,8 @@ class Dashboard: """The identifier of the workspace folder containing the object.""" permission_tier: Optional[PermissionLevel] = None - """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_MANAGE`: Can manage the - query""" + """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_EDIT`: Can edit the query + * `CAN_MANAGE`: Can manage the query""" slug: Optional[str] = None """URL slug. Usually mirrors the query name with dashes (`-`) instead of spaces. Appears in the URL @@ -745,6 +745,33 @@ def from_dict(cls, d: Dict[str, any]) -> Dashboard: widgets=_repeated_dict(d, 'widgets', Widget)) +@dataclass +class DashboardEditContent: + dashboard_id: Optional[str] = None + + name: Optional[str] = None + """The title of this dashboard that appears in list views and at the top of the dashboard page.""" + + run_as_role: Optional[RunAsRole] = None + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" + + def as_dict(self) -> dict: + """Serializes the DashboardEditContent into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dashboard_id is not None: body['dashboard_id'] = self.dashboard_id + if self.name is not None: body['name'] = self.name + if self.run_as_role is not None: body['run_as_role'] = self.run_as_role.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DashboardEditContent: + """Deserializes the DashboardEditContent from a dictionary.""" + return cls(dashboard_id=d.get('dashboard_id', None), + name=d.get('name', None), + run_as_role=_enum(d, 'run_as_role', RunAsRole)) + + @dataclass class DashboardOptions: moved_to_trash_at: Optional[str] = None @@ -763,6 +790,49 @@ def from_dict(cls, d: Dict[str, any]) -> DashboardOptions: return cls(moved_to_trash_at=d.get('moved_to_trash_at', None)) +@dataclass +class DashboardPostContent: + name: str + """The title of this dashboard that appears in list views and at the top of the dashboard page.""" + + dashboard_filters_enabled: Optional[bool] = None + """Indicates whether the dashboard filters are enabled""" + + is_favorite: Optional[bool] = None + """Indicates whether this dashboard object should appear in the current user's favorites list.""" + + parent: Optional[str] = None + """The identifier of the workspace folder containing the object.""" + + run_as_role: Optional[RunAsRole] = None + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" + + tags: Optional[List[str]] = None + + def as_dict(self) -> dict: + """Serializes the DashboardPostContent into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.dashboard_filters_enabled is not None: + body['dashboard_filters_enabled'] = self.dashboard_filters_enabled + if self.is_favorite is not None: body['is_favorite'] = self.is_favorite + if self.name is not None: body['name'] = self.name + if self.parent is not None: body['parent'] = self.parent + if self.run_as_role is not None: body['run_as_role'] = self.run_as_role.value + if self.tags: body['tags'] = [v for v in self.tags] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DashboardPostContent: + """Deserializes the DashboardPostContent from a dictionary.""" + return cls(dashboard_filters_enabled=d.get('dashboard_filters_enabled', None), + is_favorite=d.get('is_favorite', None), + name=d.get('name', None), + parent=d.get('parent', None), + run_as_role=_enum(d, 'run_as_role', RunAsRole), + tags=d.get('tags', None)) + + @dataclass class DataSource: """A JSON object representing a DBSQL data source / SQL warehouse.""" @@ -1488,7 +1558,13 @@ class ExternalLink: external_link: Optional[str] = None """A presigned URL pointing to a chunk of result data, hosted by an external service, with a short expiration time (<= 15 minutes). As this URL contains a temporary credential, it should be - considered sensitive and the client should expose this URL in a log.""" + considered sensitive and the client should not expose this URL in a log.""" + + http_headers: Optional[Dict[str, str]] = None + """HTTP headers that must be included with a GET request to the `external_link`. Each header is + provided as a key-value pair. Headers are typically used to pass a decryption key to the + external service. The values of these headers should be considered sensitive and the client + should not expose these values in a log.""" next_chunk_index: Optional[int] = None """When fetching, provides the `chunk_index` for the _next_ chunk. If absent, indicates there are @@ -1513,6 +1589,7 @@ def as_dict(self) -> dict: if self.chunk_index is not None: body['chunk_index'] = self.chunk_index if self.expiration is not None: body['expiration'] = self.expiration if self.external_link is not None: body['external_link'] = self.external_link + if self.http_headers: body['http_headers'] = self.http_headers if self.next_chunk_index is not None: body['next_chunk_index'] = self.next_chunk_index if self.next_chunk_internal_link is not None: body['next_chunk_internal_link'] = self.next_chunk_internal_link @@ -1527,6 +1604,7 @@ def from_dict(cls, d: Dict[str, any]) -> ExternalLink: chunk_index=d.get('chunk_index', None), expiration=d.get('expiration', None), external_link=d.get('external_link', None), + http_headers=d.get('http_headers', None), next_chunk_index=d.get('next_chunk_index', None), next_chunk_internal_link=d.get('next_chunk_internal_link', None), row_count=d.get('row_count', None), @@ -2025,9 +2103,10 @@ class ParameterType(Enum): class PermissionLevel(Enum): - """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_MANAGE`: Can manage the - query""" + """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_EDIT`: Can edit the query + * `CAN_MANAGE`: Can manage the query""" + CAN_EDIT = 'CAN_EDIT' CAN_MANAGE = 'CAN_MANAGE' CAN_RUN = 'CAN_RUN' CAN_VIEW = 'CAN_VIEW' @@ -2100,8 +2179,8 @@ class Query: """The identifier of the workspace folder containing the object.""" permission_tier: Optional[PermissionLevel] = None - """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_MANAGE`: Can manage the - query""" + """* `CAN_VIEW`: Can view the query * `CAN_RUN`: Can run the query * `CAN_EDIT`: Can edit the query + * `CAN_MANAGE`: Can manage the query""" query: Optional[str] = None """The text of the query to be run.""" @@ -2110,7 +2189,8 @@ class Query: """A SHA-256 hash of the query text along with the authenticated user ID.""" run_as_role: Optional[RunAsRole] = None - """Run as role""" + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" tags: Optional[List[str]] = None @@ -2206,6 +2286,10 @@ class QueryEditContent: query_id: Optional[str] = None + run_as_role: Optional[RunAsRole] = None + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" + def as_dict(self) -> dict: """Serializes the QueryEditContent into a dictionary suitable for use as a JSON request body.""" body = {} @@ -2215,6 +2299,7 @@ def as_dict(self) -> dict: if self.options: body['options'] = self.options if self.query is not None: body['query'] = self.query if self.query_id is not None: body['query_id'] = self.query_id + if self.run_as_role is not None: body['run_as_role'] = self.run_as_role.value return body @classmethod @@ -2225,7 +2310,8 @@ def from_dict(cls, d: Dict[str, any]) -> QueryEditContent: name=d.get('name', None), options=d.get('options', None), query=d.get('query', None), - query_id=d.get('query_id', None)) + query_id=d.get('query_id', None), + run_as_role=_enum(d, 'run_as_role', RunAsRole)) @dataclass @@ -2626,7 +2712,8 @@ class QueryPostContent: """The text of the query to be run.""" run_as_role: Optional[RunAsRole] = None - """Run as role""" + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" def as_dict(self) -> dict: """Serializes the QueryPostContent into a dictionary suitable for use as a JSON request body.""" @@ -2848,7 +2935,8 @@ def from_dict(cls, d: Dict[str, any]) -> ResultSchema: class RunAsRole(Enum): - """Run as role""" + """Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior)""" OWNER = 'owner' VIEWER = 'viewer' @@ -3903,12 +3991,12 @@ def create(self, :param dashboard_filters_enabled: bool (optional) Indicates whether the dashboard filters are enabled :param is_favorite: bool (optional) - Indicates whether this query object should appear in the current user's favorites list. The - application uses this flag to determine whether or not the "favorite star " should selected. + Indicates whether this dashboard object should appear in the current user's favorites list. :param parent: str (optional) The identifier of the workspace folder containing the object. :param run_as_role: :class:`RunAsRole` (optional) - Run as role + Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior) :param tags: List[str] (optional) :returns: :class:`Dashboard` @@ -4010,6 +4098,37 @@ def restore(self, dashboard_id: str): headers = {'Accept': 'application/json', } self._api.do('POST', f'/api/2.0/preview/sql/dashboards/trash/{dashboard_id}', headers=headers) + def update(self, + dashboard_id: str, + *, + name: Optional[str] = None, + run_as_role: Optional[RunAsRole] = None) -> Dashboard: + """Change a dashboard definition. + + Modify this dashboard definition. This operation only affects attributes of the dashboard object. It + does not add, modify, or remove widgets. + + **Note**: You cannot undo this operation. + + :param dashboard_id: str + :param name: str (optional) + The title of this dashboard that appears in list views and at the top of the dashboard page. + :param run_as_role: :class:`RunAsRole` (optional) + Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior) + + :returns: :class:`Dashboard` + """ + body = {} + if name is not None: body['name'] = name + if run_as_role is not None: body['run_as_role'] = run_as_role.value + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', + f'/api/2.0/preview/sql/dashboards/{dashboard_id}', + body=body, + headers=headers) + return Dashboard.from_dict(res) + class DataSourcesAPI: """This API is provided to assist you in making new query objects. When creating a query object, you may @@ -4175,7 +4294,8 @@ def create(self, :param query: str (optional) The text of the query to be run. :param run_as_role: :class:`RunAsRole` (optional) - Run as role + Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior) :returns: :class:`Query` """ @@ -4297,7 +4417,8 @@ def update(self, description: Optional[str] = None, name: Optional[str] = None, options: Optional[Any] = None, - query: Optional[str] = None) -> Query: + query: Optional[str] = None, + run_as_role: Optional[RunAsRole] = None) -> Query: """Change a query definition. Modify this query definition. @@ -4320,6 +4441,9 @@ def update(self, overridden at runtime. :param query: str (optional) The text of the query to be run. + :param run_as_role: :class:`RunAsRole` (optional) + Sets the **Run as** role for the object. Must be set to one of `"viewer"` (signifying "run as + viewer" behavior) or `"owner"` (signifying "run as owner" behavior) :returns: :class:`Query` """ @@ -4329,6 +4453,7 @@ def update(self, if name is not None: body['name'] = name if options is not None: body['options'] = options if query is not None: body['query'] = query + if run_as_role is not None: body['run_as_role'] = run_as_role.value headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } res = self._api.do('POST', f'/api/2.0/preview/sql/queries/{query_id}', body=body, headers=headers) return Query.from_dict(res) diff --git a/databricks/sdk/service/vectorsearch.py b/databricks/sdk/service/vectorsearch.py new file mode 100755 index 00000000..95ad717c --- /dev/null +++ b/databricks/sdk/service/vectorsearch.py @@ -0,0 +1,1206 @@ +# Code generated from OpenAPI specs by Databricks SDK Generator. DO NOT EDIT. + +from __future__ import annotations + +import logging +import random +import time +from dataclasses import dataclass +from datetime import timedelta +from enum import Enum +from typing import Callable, Dict, Iterator, List, Optional + +from ..errors import OperationFailed +from ._internal import Wait, _enum, _from_dict, _repeated_dict + +_LOG = logging.getLogger('databricks.sdk') + +# all definitions in this file are in alphabetical order + + +@dataclass +class ColumnInfo: + name: Optional[str] = None + """Name of the column.""" + + def as_dict(self) -> dict: + """Serializes the ColumnInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.name is not None: body['name'] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> ColumnInfo: + """Deserializes the ColumnInfo from a dictionary.""" + return cls(name=d.get('name', None)) + + +@dataclass +class CreateEndpoint: + name: str + """Name of endpoint""" + + endpoint_type: EndpointType + """Type of endpoint.""" + + def as_dict(self) -> dict: + """Serializes the CreateEndpoint into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.endpoint_type is not None: body['endpoint_type'] = self.endpoint_type.value + if self.name is not None: body['name'] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> CreateEndpoint: + """Deserializes the CreateEndpoint from a dictionary.""" + return cls(endpoint_type=_enum(d, 'endpoint_type', EndpointType), name=d.get('name', None)) + + +@dataclass +class CreateVectorIndexRequest: + name: str + """Name of the index""" + + primary_key: str + """Primary key of the index""" + + index_type: VectorIndexType + """There are 2 types of Vector Search indexes: + + - `DELTA_SYNC`: An index that automatically syncs with a source Delta Table, automatically and + incrementally updating the index as the underlying data in the Delta Table changes. - + `DIRECT_ACCESS`: An index that supports direct read and write of vectors and metadata through + our REST and SDK APIs. With this model, the user manages index updates.""" + + delta_sync_vector_index_spec: Optional[DeltaSyncVectorIndexSpecRequest] = None + """Specification for Delta Sync Index. Required if `index_type` is `DELTA_SYNC`.""" + + direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None + """Specification for Direct Vector Access Index. Required if `index_type` is `DIRECT_ACCESS`.""" + + endpoint_name: Optional[str] = None + """Name of the endpoint to be used for serving the index""" + + def as_dict(self) -> dict: + """Serializes the CreateVectorIndexRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.delta_sync_vector_index_spec: + body['delta_sync_vector_index_spec'] = self.delta_sync_vector_index_spec.as_dict() + if self.direct_access_index_spec: + body['direct_access_index_spec'] = self.direct_access_index_spec.as_dict() + if self.endpoint_name is not None: body['endpoint_name'] = self.endpoint_name + if self.index_type is not None: body['index_type'] = self.index_type.value + if self.name is not None: body['name'] = self.name + if self.primary_key is not None: body['primary_key'] = self.primary_key + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> CreateVectorIndexRequest: + """Deserializes the CreateVectorIndexRequest from a dictionary.""" + return cls(delta_sync_vector_index_spec=_from_dict(d, 'delta_sync_vector_index_spec', + DeltaSyncVectorIndexSpecRequest), + direct_access_index_spec=_from_dict(d, 'direct_access_index_spec', + DirectAccessVectorIndexSpec), + endpoint_name=d.get('endpoint_name', None), + index_type=_enum(d, 'index_type', VectorIndexType), + name=d.get('name', None), + primary_key=d.get('primary_key', None)) + + +@dataclass +class CreateVectorIndexResponse: + vector_index: Optional[VectorIndex] = None + + def as_dict(self) -> dict: + """Serializes the CreateVectorIndexResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.vector_index: body['vector_index'] = self.vector_index.as_dict() + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> CreateVectorIndexResponse: + """Deserializes the CreateVectorIndexResponse from a dictionary.""" + return cls(vector_index=_from_dict(d, 'vector_index', VectorIndex)) + + +@dataclass +class DeleteDataResult: + """Result of the upsert or delete operation.""" + + failed_primary_keys: Optional[List[str]] = None + """List of primary keys for rows that failed to process.""" + + success_row_count: Optional[int] = None + """Count of successfully processed rows.""" + + def as_dict(self) -> dict: + """Serializes the DeleteDataResult into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.failed_primary_keys: body['failed_primary_keys'] = [v for v in self.failed_primary_keys] + if self.success_row_count is not None: body['success_row_count'] = self.success_row_count + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DeleteDataResult: + """Deserializes the DeleteDataResult from a dictionary.""" + return cls(failed_primary_keys=d.get('failed_primary_keys', None), + success_row_count=d.get('success_row_count', None)) + + +class DeleteDataStatus(Enum): + """Status of the delete operation.""" + + FAILURE = 'FAILURE' + PARTIAL_SUCCESS = 'PARTIAL_SUCCESS' + SUCCESS = 'SUCCESS' + + +@dataclass +class DeleteDataVectorIndexRequest: + """Request payload for deleting data from a vector index.""" + + primary_keys: List[str] + """List of primary keys for the data to be deleted.""" + + name: Optional[str] = None + """Name of the vector index where data is to be deleted. Must be a Direct Vector Access Index.""" + + def as_dict(self) -> dict: + """Serializes the DeleteDataVectorIndexRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.name is not None: body['name'] = self.name + if self.primary_keys: body['primary_keys'] = [v for v in self.primary_keys] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DeleteDataVectorIndexRequest: + """Deserializes the DeleteDataVectorIndexRequest from a dictionary.""" + return cls(name=d.get('name', None), primary_keys=d.get('primary_keys', None)) + + +@dataclass +class DeleteDataVectorIndexResponse: + """Response to a delete data vector index request.""" + + result: Optional[DeleteDataResult] = None + """Result of the upsert or delete operation.""" + + status: Optional[DeleteDataStatus] = None + """Status of the delete operation.""" + + def as_dict(self) -> dict: + """Serializes the DeleteDataVectorIndexResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.result: body['result'] = self.result.as_dict() + if self.status is not None: body['status'] = self.status.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DeleteDataVectorIndexResponse: + """Deserializes the DeleteDataVectorIndexResponse from a dictionary.""" + return cls(result=_from_dict(d, 'result', DeleteDataResult), + status=_enum(d, 'status', DeleteDataStatus)) + + +@dataclass +class DeltaSyncVectorIndexSpecRequest: + embedding_source_columns: Optional[List[EmbeddingSourceColumn]] = None + """The columns that contain the embedding source.""" + + embedding_vector_columns: Optional[List[EmbeddingVectorColumn]] = None + """The columns that contain the embedding vectors.""" + + pipeline_type: Optional[PipelineType] = None + """Pipeline execution mode. + + - `TRIGGERED`: If the pipeline uses the triggered execution mode, the system stops processing + after successfully refreshing the source table in the pipeline once, ensuring the table is + updated based on the data available when the update started. - `CONTINUOUS`: If the pipeline + uses continuous execution, the pipeline processes new data as it arrives in the source table to + keep vector index fresh.""" + + source_table: Optional[str] = None + """The name of the source table.""" + + def as_dict(self) -> dict: + """Serializes the DeltaSyncVectorIndexSpecRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_source_columns: + body['embedding_source_columns'] = [v.as_dict() for v in self.embedding_source_columns] + if self.embedding_vector_columns: + body['embedding_vector_columns'] = [v.as_dict() for v in self.embedding_vector_columns] + if self.pipeline_type is not None: body['pipeline_type'] = self.pipeline_type.value + if self.source_table is not None: body['source_table'] = self.source_table + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DeltaSyncVectorIndexSpecRequest: + """Deserializes the DeltaSyncVectorIndexSpecRequest from a dictionary.""" + return cls(embedding_source_columns=_repeated_dict(d, 'embedding_source_columns', + EmbeddingSourceColumn), + embedding_vector_columns=_repeated_dict(d, 'embedding_vector_columns', + EmbeddingVectorColumn), + pipeline_type=_enum(d, 'pipeline_type', PipelineType), + source_table=d.get('source_table', None)) + + +@dataclass +class DeltaSyncVectorIndexSpecResponse: + embedding_source_columns: Optional[List[EmbeddingSourceColumn]] = None + """The columns that contain the embedding source.""" + + embedding_vector_columns: Optional[List[EmbeddingVectorColumn]] = None + """The columns that contain the embedding vectors.""" + + pipeline_id: Optional[str] = None + """The ID of the pipeline that is used to sync the index.""" + + pipeline_type: Optional[PipelineType] = None + """Pipeline execution mode. + + - `TRIGGERED`: If the pipeline uses the triggered execution mode, the system stops processing + after successfully refreshing the source table in the pipeline once, ensuring the table is + updated based on the data available when the update started. - `CONTINUOUS`: If the pipeline + uses continuous execution, the pipeline processes new data as it arrives in the source table to + keep vector index fresh.""" + + source_table: Optional[str] = None + """The name of the source table.""" + + def as_dict(self) -> dict: + """Serializes the DeltaSyncVectorIndexSpecResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_source_columns: + body['embedding_source_columns'] = [v.as_dict() for v in self.embedding_source_columns] + if self.embedding_vector_columns: + body['embedding_vector_columns'] = [v.as_dict() for v in self.embedding_vector_columns] + if self.pipeline_id is not None: body['pipeline_id'] = self.pipeline_id + if self.pipeline_type is not None: body['pipeline_type'] = self.pipeline_type.value + if self.source_table is not None: body['source_table'] = self.source_table + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DeltaSyncVectorIndexSpecResponse: + """Deserializes the DeltaSyncVectorIndexSpecResponse from a dictionary.""" + return cls(embedding_source_columns=_repeated_dict(d, 'embedding_source_columns', + EmbeddingSourceColumn), + embedding_vector_columns=_repeated_dict(d, 'embedding_vector_columns', + EmbeddingVectorColumn), + pipeline_id=d.get('pipeline_id', None), + pipeline_type=_enum(d, 'pipeline_type', PipelineType), + source_table=d.get('source_table', None)) + + +@dataclass +class DirectAccessVectorIndexSpec: + embedding_vector_columns: Optional[List[EmbeddingVectorColumn]] = None + + schema_json: Optional[str] = None + """The schema of the index in JSON format. + + Supported types are `integer`, `long`, `float`, `double`, `boolean`, `string`, `date`, + `timestamp`. + + Supported types for vector column: `array`, `array`,`.""" + + def as_dict(self) -> dict: + """Serializes the DirectAccessVectorIndexSpec into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_vector_columns: + body['embedding_vector_columns'] = [v.as_dict() for v in self.embedding_vector_columns] + if self.schema_json is not None: body['schema_json'] = self.schema_json + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> DirectAccessVectorIndexSpec: + """Deserializes the DirectAccessVectorIndexSpec from a dictionary.""" + return cls(embedding_vector_columns=_repeated_dict(d, 'embedding_vector_columns', + EmbeddingVectorColumn), + schema_json=d.get('schema_json', None)) + + +@dataclass +class EmbeddingConfig: + embedding_model_endpoint_name: Optional[str] = None + """Name of the embedding model endpoint""" + + def as_dict(self) -> dict: + """Serializes the EmbeddingConfig into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_model_endpoint_name is not None: + body['embedding_model_endpoint_name'] = self.embedding_model_endpoint_name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> EmbeddingConfig: + """Deserializes the EmbeddingConfig from a dictionary.""" + return cls(embedding_model_endpoint_name=d.get('embedding_model_endpoint_name', None)) + + +@dataclass +class EmbeddingSourceColumn: + embedding_config: Optional[EmbeddingConfig] = None + + name: Optional[str] = None + """Name of the column""" + + def as_dict(self) -> dict: + """Serializes the EmbeddingSourceColumn into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_config: body['embedding_config'] = self.embedding_config.as_dict() + if self.name is not None: body['name'] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> EmbeddingSourceColumn: + """Deserializes the EmbeddingSourceColumn from a dictionary.""" + return cls(embedding_config=_from_dict(d, 'embedding_config', EmbeddingConfig), + name=d.get('name', None)) + + +@dataclass +class EmbeddingVectorColumn: + embedding_dimension: Optional[int] = None + """Dimension of the embedding vector""" + + name: Optional[str] = None + """Name of the column""" + + def as_dict(self) -> dict: + """Serializes the EmbeddingVectorColumn into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.embedding_dimension is not None: body['embedding_dimension'] = self.embedding_dimension + if self.name is not None: body['name'] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> EmbeddingVectorColumn: + """Deserializes the EmbeddingVectorColumn from a dictionary.""" + return cls(embedding_dimension=d.get('embedding_dimension', None), name=d.get('name', None)) + + +@dataclass +class EndpointInfo: + creation_timestamp: Optional[int] = None + """Timestamp of endpoint creation""" + + creator: Optional[str] = None + """Creator of the endpoint""" + + endpoint_status: Optional[EndpointStatus] = None + """Current status of the endpoint""" + + endpoint_type: Optional[EndpointType] = None + """Type of endpoint.""" + + id: Optional[str] = None + """Unique identifier of the endpoint""" + + last_updated_timestamp: Optional[int] = None + """Timestamp of last update to the endpoint""" + + last_updated_user: Optional[str] = None + """User who last updated the endpoint""" + + name: Optional[str] = None + """Name of endpoint""" + + num_indexes: Optional[int] = None + """Number of indexes on the endpoint""" + + def as_dict(self) -> dict: + """Serializes the EndpointInfo into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.creation_timestamp is not None: body['creation_timestamp'] = self.creation_timestamp + if self.creator is not None: body['creator'] = self.creator + if self.endpoint_status: body['endpoint_status'] = self.endpoint_status.as_dict() + if self.endpoint_type is not None: body['endpoint_type'] = self.endpoint_type.value + if self.id is not None: body['id'] = self.id + if self.last_updated_timestamp is not None: + body['last_updated_timestamp'] = self.last_updated_timestamp + if self.last_updated_user is not None: body['last_updated_user'] = self.last_updated_user + if self.name is not None: body['name'] = self.name + if self.num_indexes is not None: body['num_indexes'] = self.num_indexes + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> EndpointInfo: + """Deserializes the EndpointInfo from a dictionary.""" + return cls(creation_timestamp=d.get('creation_timestamp', None), + creator=d.get('creator', None), + endpoint_status=_from_dict(d, 'endpoint_status', EndpointStatus), + endpoint_type=_enum(d, 'endpoint_type', EndpointType), + id=d.get('id', None), + last_updated_timestamp=d.get('last_updated_timestamp', None), + last_updated_user=d.get('last_updated_user', None), + name=d.get('name', None), + num_indexes=d.get('num_indexes', None)) + + +@dataclass +class EndpointStatus: + """Status information of an endpoint""" + + message: Optional[str] = None + """Additional status message""" + + state: Optional[EndpointStatusState] = None + """Current state of the endpoint""" + + def as_dict(self) -> dict: + """Serializes the EndpointStatus into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.message is not None: body['message'] = self.message + if self.state is not None: body['state'] = self.state.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> EndpointStatus: + """Deserializes the EndpointStatus from a dictionary.""" + return cls(message=d.get('message', None), state=_enum(d, 'state', EndpointStatusState)) + + +class EndpointStatusState(Enum): + """Current state of the endpoint""" + + OFFLINE = 'OFFLINE' + ONLINE = 'ONLINE' + PROVISIONING = 'PROVISIONING' + + +class EndpointType(Enum): + """Type of endpoint.""" + + STANDARD = 'STANDARD' + + +@dataclass +class ListEndpointResponse: + endpoints: Optional[List[EndpointInfo]] = None + """An array of Endpoint objects""" + + next_page_token: Optional[str] = None + """A token that can be used to get the next page of results. If not present, there are no more + results to show.""" + + def as_dict(self) -> dict: + """Serializes the ListEndpointResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.endpoints: body['endpoints'] = [v.as_dict() for v in self.endpoints] + 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]) -> ListEndpointResponse: + """Deserializes the ListEndpointResponse from a dictionary.""" + return cls(endpoints=_repeated_dict(d, 'endpoints', EndpointInfo), + next_page_token=d.get('next_page_token', None)) + + +@dataclass +class ListVectorIndexesResponse: + next_page_token: Optional[str] = None + """A token that can be used to get the next page of results. If not present, there are no more + results to show.""" + + vector_indexes: Optional[List[MiniVectorIndex]] = None + + def as_dict(self) -> dict: + """Serializes the ListVectorIndexesResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.next_page_token is not None: body['next_page_token'] = self.next_page_token + if self.vector_indexes: body['vector_indexes'] = [v.as_dict() for v in self.vector_indexes] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> ListVectorIndexesResponse: + """Deserializes the ListVectorIndexesResponse from a dictionary.""" + return cls(next_page_token=d.get('next_page_token', None), + vector_indexes=_repeated_dict(d, 'vector_indexes', MiniVectorIndex)) + + +@dataclass +class MiniVectorIndex: + creator: Optional[str] = None + """The user who created the index.""" + + endpoint_name: Optional[str] = None + """Name of the endpoint associated with the index""" + + index_type: Optional[VectorIndexType] = None + """There are 2 types of Vector Search indexes: + + - `DELTA_SYNC`: An index that automatically syncs with a source Delta Table, automatically and + incrementally updating the index as the underlying data in the Delta Table changes. - + `DIRECT_ACCESS`: An index that supports direct read and write of vectors and metadata through + our REST and SDK APIs. With this model, the user manages index updates.""" + + name: Optional[str] = None + """Name of the index""" + + primary_key: Optional[str] = None + """Primary key of the index""" + + def as_dict(self) -> dict: + """Serializes the MiniVectorIndex into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.creator is not None: body['creator'] = self.creator + if self.endpoint_name is not None: body['endpoint_name'] = self.endpoint_name + if self.index_type is not None: body['index_type'] = self.index_type.value + if self.name is not None: body['name'] = self.name + if self.primary_key is not None: body['primary_key'] = self.primary_key + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> MiniVectorIndex: + """Deserializes the MiniVectorIndex from a dictionary.""" + return cls(creator=d.get('creator', None), + endpoint_name=d.get('endpoint_name', None), + index_type=_enum(d, 'index_type', VectorIndexType), + name=d.get('name', None), + primary_key=d.get('primary_key', None)) + + +class PipelineType(Enum): + """Pipeline execution mode. + + - `TRIGGERED`: If the pipeline uses the triggered execution mode, the system stops processing + after successfully refreshing the source table in the pipeline once, ensuring the table is + updated based on the data available when the update started. - `CONTINUOUS`: If the pipeline + uses continuous execution, the pipeline processes new data as it arrives in the source table to + keep vector index fresh.""" + + CONTINUOUS = 'CONTINUOUS' + TRIGGERED = 'TRIGGERED' + + +@dataclass +class QueryVectorIndexRequest: + columns: List[str] + """List of column names to include in the response.""" + + filters_json: Optional[str] = None + """JSON string representing query filters. + + Example filters: - `{"id <": 5}`: Filter for id less than 5. - `{"id >": 5}`: Filter for id + greater than 5. - `{"id <=": 5}`: Filter for id less than equal to 5. - `{"id >=": 5}`: Filter + for id greater than equal to 5. - `{"id": 5}`: Filter for id equal to 5.""" + + index_name: Optional[str] = None + """Name of the vector index to query.""" + + num_results: Optional[int] = None + """Number of results to return. Defaults to 10.""" + + query_text: Optional[str] = None + """Query text. Required for Delta Sync Index using model endpoint.""" + + query_vector: Optional[List[float]] = None + """Query vector. Required for Direct Vector Access Index and Delta Sync Index using self-managed + vectors.""" + + def as_dict(self) -> dict: + """Serializes the QueryVectorIndexRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.columns: body['columns'] = [v for v in self.columns] + if self.filters_json is not None: body['filters_json'] = self.filters_json + if self.index_name is not None: body['index_name'] = self.index_name + if self.num_results is not None: body['num_results'] = self.num_results + if self.query_text is not None: body['query_text'] = self.query_text + if self.query_vector: body['query_vector'] = [v for v in self.query_vector] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> QueryVectorIndexRequest: + """Deserializes the QueryVectorIndexRequest from a dictionary.""" + return cls(columns=d.get('columns', None), + filters_json=d.get('filters_json', None), + index_name=d.get('index_name', None), + num_results=d.get('num_results', None), + query_text=d.get('query_text', None), + query_vector=d.get('query_vector', None)) + + +@dataclass +class QueryVectorIndexResponse: + manifest: Optional[ResultManifest] = None + """Metadata about the result set.""" + + result: Optional[ResultData] = None + """Data returned in the query result.""" + + def as_dict(self) -> dict: + """Serializes the QueryVectorIndexResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.manifest: body['manifest'] = self.manifest.as_dict() + if self.result: body['result'] = self.result.as_dict() + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> QueryVectorIndexResponse: + """Deserializes the QueryVectorIndexResponse from a dictionary.""" + return cls(manifest=_from_dict(d, 'manifest', ResultManifest), + result=_from_dict(d, 'result', ResultData)) + + +@dataclass +class ResultData: + """Data returned in the query result.""" + + data_array: Optional[List[List[str]]] = None + """Data rows returned in the query.""" + + row_count: Optional[int] = None + """Number of rows in the result set.""" + + def as_dict(self) -> dict: + """Serializes the ResultData into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.data_array: body['data_array'] = [v for v in self.data_array] + if self.row_count is not None: body['row_count'] = self.row_count + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> ResultData: + """Deserializes the ResultData from a dictionary.""" + return cls(data_array=d.get('data_array', None), row_count=d.get('row_count', None)) + + +@dataclass +class ResultManifest: + """Metadata about the result set.""" + + column_count: Optional[int] = None + """Number of columns in the result set.""" + + columns: Optional[List[ColumnInfo]] = None + """Information about each column in the result set.""" + + def as_dict(self) -> dict: + """Serializes the ResultManifest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.column_count is not None: body['column_count'] = self.column_count + if self.columns: body['columns'] = [v.as_dict() for v in self.columns] + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> ResultManifest: + """Deserializes the ResultManifest from a dictionary.""" + return cls(column_count=d.get('column_count', None), columns=_repeated_dict(d, 'columns', ColumnInfo)) + + +@dataclass +class UpsertDataResult: + """Result of the upsert or delete operation.""" + + failed_primary_keys: Optional[List[str]] = None + """List of primary keys for rows that failed to process.""" + + success_row_count: Optional[int] = None + """Count of successfully processed rows.""" + + def as_dict(self) -> dict: + """Serializes the UpsertDataResult into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.failed_primary_keys: body['failed_primary_keys'] = [v for v in self.failed_primary_keys] + if self.success_row_count is not None: body['success_row_count'] = self.success_row_count + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> UpsertDataResult: + """Deserializes the UpsertDataResult from a dictionary.""" + return cls(failed_primary_keys=d.get('failed_primary_keys', None), + success_row_count=d.get('success_row_count', None)) + + +class UpsertDataStatus(Enum): + """Status of the upsert operation.""" + + FAILURE = 'FAILURE' + PARTIAL_SUCCESS = 'PARTIAL_SUCCESS' + SUCCESS = 'SUCCESS' + + +@dataclass +class UpsertDataVectorIndexRequest: + """Request payload for upserting data into a vector index.""" + + inputs_json: str + """JSON string representing the data to be upserted.""" + + name: Optional[str] = None + """Name of the vector index where data is to be upserted. Must be a Direct Vector Access Index.""" + + def as_dict(self) -> dict: + """Serializes the UpsertDataVectorIndexRequest into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.inputs_json is not None: body['inputs_json'] = self.inputs_json + if self.name is not None: body['name'] = self.name + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> UpsertDataVectorIndexRequest: + """Deserializes the UpsertDataVectorIndexRequest from a dictionary.""" + return cls(inputs_json=d.get('inputs_json', None), name=d.get('name', None)) + + +@dataclass +class UpsertDataVectorIndexResponse: + """Response to an upsert data vector index request.""" + + result: Optional[UpsertDataResult] = None + """Result of the upsert or delete operation.""" + + status: Optional[UpsertDataStatus] = None + """Status of the upsert operation.""" + + def as_dict(self) -> dict: + """Serializes the UpsertDataVectorIndexResponse into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.result: body['result'] = self.result.as_dict() + if self.status is not None: body['status'] = self.status.value + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> UpsertDataVectorIndexResponse: + """Deserializes the UpsertDataVectorIndexResponse from a dictionary.""" + return cls(result=_from_dict(d, 'result', UpsertDataResult), + status=_enum(d, 'status', UpsertDataStatus)) + + +@dataclass +class VectorIndex: + creator: Optional[str] = None + """The user who created the index.""" + + delta_sync_vector_index_spec: Optional[DeltaSyncVectorIndexSpecResponse] = None + + direct_access_vector_index_spec: Optional[DirectAccessVectorIndexSpec] = None + + endpoint_name: Optional[str] = None + """Name of the endpoint associated with the index""" + + index_type: Optional[VectorIndexType] = None + """There are 2 types of Vector Search indexes: + + - `DELTA_SYNC`: An index that automatically syncs with a source Delta Table, automatically and + incrementally updating the index as the underlying data in the Delta Table changes. - + `DIRECT_ACCESS`: An index that supports direct read and write of vectors and metadata through + our REST and SDK APIs. With this model, the user manages index updates.""" + + name: Optional[str] = None + """Name of the index""" + + primary_key: Optional[str] = None + """Primary key of the index""" + + status: Optional[VectorIndexStatus] = None + + def as_dict(self) -> dict: + """Serializes the VectorIndex into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.creator is not None: body['creator'] = self.creator + if self.delta_sync_vector_index_spec: + body['delta_sync_vector_index_spec'] = self.delta_sync_vector_index_spec.as_dict() + if self.direct_access_vector_index_spec: + body['direct_access_vector_index_spec'] = self.direct_access_vector_index_spec.as_dict() + if self.endpoint_name is not None: body['endpoint_name'] = self.endpoint_name + if self.index_type is not None: body['index_type'] = self.index_type.value + if self.name is not None: body['name'] = self.name + if self.primary_key is not None: body['primary_key'] = self.primary_key + if self.status: body['status'] = self.status.as_dict() + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> VectorIndex: + """Deserializes the VectorIndex from a dictionary.""" + return cls(creator=d.get('creator', None), + delta_sync_vector_index_spec=_from_dict(d, 'delta_sync_vector_index_spec', + DeltaSyncVectorIndexSpecResponse), + direct_access_vector_index_spec=_from_dict(d, 'direct_access_vector_index_spec', + DirectAccessVectorIndexSpec), + endpoint_name=d.get('endpoint_name', None), + index_type=_enum(d, 'index_type', VectorIndexType), + name=d.get('name', None), + primary_key=d.get('primary_key', None), + status=_from_dict(d, 'status', VectorIndexStatus)) + + +@dataclass +class VectorIndexStatus: + index_url: Optional[str] = None + """Index API Url to be used to perform operations on the index""" + + indexed_row_count: Optional[int] = None + """Number of rows indexed""" + + message: Optional[str] = None + """Message associated with the index status""" + + ready: Optional[bool] = None + """Whether the index is ready for search""" + + def as_dict(self) -> dict: + """Serializes the VectorIndexStatus into a dictionary suitable for use as a JSON request body.""" + body = {} + if self.index_url is not None: body['index_url'] = self.index_url + if self.indexed_row_count is not None: body['indexed_row_count'] = self.indexed_row_count + if self.message is not None: body['message'] = self.message + if self.ready is not None: body['ready'] = self.ready + return body + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> VectorIndexStatus: + """Deserializes the VectorIndexStatus from a dictionary.""" + return cls(index_url=d.get('index_url', None), + indexed_row_count=d.get('indexed_row_count', None), + message=d.get('message', None), + ready=d.get('ready', None)) + + +class VectorIndexType(Enum): + """There are 2 types of Vector Search indexes: + + - `DELTA_SYNC`: An index that automatically syncs with a source Delta Table, automatically and + incrementally updating the index as the underlying data in the Delta Table changes. - + `DIRECT_ACCESS`: An index that supports direct read and write of vectors and metadata through + our REST and SDK APIs. With this model, the user manages index updates.""" + + DELTA_SYNC = 'DELTA_SYNC' + DIRECT_ACCESS = 'DIRECT_ACCESS' + + +class VectorSearchEndpointsAPI: + """**Endpoint**: Represents the compute resources to host vector search indexes.""" + + def __init__(self, api_client): + self._api = api_client + + def wait_get_endpoint_vector_search_endpoint_online( + self, + endpoint_name: str, + timeout=timedelta(minutes=20), + callback: Optional[Callable[[EndpointInfo], None]] = None) -> EndpointInfo: + deadline = time.time() + timeout.total_seconds() + target_states = (EndpointStatusState.ONLINE, ) + failure_states = (EndpointStatusState.OFFLINE, ) + status_message = 'polling...' + attempt = 1 + while time.time() < deadline: + poll = self.get_endpoint(endpoint_name=endpoint_name) + status = poll.endpoint_status.state + status_message = f'current status: {status}' + if poll.endpoint_status: + status_message = poll.endpoint_status.message + if status in target_states: + return poll + if callback: + callback(poll) + if status in failure_states: + msg = f'failed to reach ONLINE, got {status}: {status_message}' + raise OperationFailed(msg) + prefix = f"endpoint_name={endpoint_name}" + sleep = attempt + if sleep > 10: + # sleep 10s max per attempt + sleep = 10 + _LOG.debug(f'{prefix}: ({status}) {status_message} (sleeping ~{sleep}s)') + time.sleep(sleep + random.random()) + attempt += 1 + raise TimeoutError(f'timed out after {timeout}: {status_message}') + + def create_endpoint(self, name: str, endpoint_type: EndpointType) -> Wait[EndpointInfo]: + """Create an endpoint. + + Create a new endpoint. + + :param name: str + Name of endpoint + :param endpoint_type: :class:`EndpointType` + Type of endpoint. + + :returns: + Long-running operation waiter for :class:`EndpointInfo`. + See :method:wait_get_endpoint_vector_search_endpoint_online for more details. + """ + body = {} + if endpoint_type is not None: body['endpoint_type'] = endpoint_type.value + if name is not None: body['name'] = name + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + op_response = self._api.do('POST', '/api/2.0/vector-search/endpoints', body=body, headers=headers) + return Wait(self.wait_get_endpoint_vector_search_endpoint_online, + response=EndpointInfo.from_dict(op_response), + endpoint_name=op_response['name']) + + def create_endpoint_and_wait(self, name: str, endpoint_type: EndpointType, + timeout=timedelta(minutes=20)) -> EndpointInfo: + return self.create_endpoint(endpoint_type=endpoint_type, name=name).result(timeout=timeout) + + def delete_endpoint(self, endpoint_name: str, name: str): + """Delete an endpoint. + + :param endpoint_name: str + Name of the endpoint + :param name: str + Name of the endpoint to delete + + + """ + + headers = {} + self._api.do('DELETE', f'/api/2.0/vector-search/endpoints/{endpoint_name}', headers=headers) + + def get_endpoint(self, endpoint_name: str) -> EndpointInfo: + """Get an endpoint. + + :param endpoint_name: str + Name of the endpoint + + :returns: :class:`EndpointInfo` + """ + + headers = {'Accept': 'application/json', } + res = self._api.do('GET', f'/api/2.0/vector-search/endpoints/{endpoint_name}', headers=headers) + return EndpointInfo.from_dict(res) + + def list_endpoints(self, *, page_token: Optional[str] = None) -> Iterator[EndpointInfo]: + """List all endpoints. + + :param page_token: str (optional) + Token for pagination + + :returns: Iterator over :class:`EndpointInfo` + """ + + query = {} + if page_token is not None: query['page_token'] = page_token + headers = {'Accept': 'application/json', } + + while True: + json = self._api.do('GET', '/api/2.0/vector-search/endpoints', query=query, headers=headers) + if 'endpoints' not in json or not json['endpoints']: + return + for v in json['endpoints']: + yield EndpointInfo.from_dict(v) + if 'next_page_token' not in json or not json['next_page_token']: + return + query['page_token'] = json['next_page_token'] + + +class VectorSearchIndexesAPI: + """**Index**: An efficient representation of your embedding vectors that supports real-time and efficient + approximate nearest neighbor (ANN) search queries. + + There are 2 types of Vector Search indexes: * **Delta Sync Index**: An index that automatically syncs with + a source Delta Table, automatically and incrementally updating the index as the underlying data in the + Delta Table changes. * **Direct Vector Access Index**: An index that supports direct read and write of + vectors and metadata through our REST and SDK APIs. With this model, the user manages index updates.""" + + def __init__(self, api_client): + self._api = api_client + + def create_index(self, + name: str, + primary_key: str, + index_type: VectorIndexType, + *, + delta_sync_vector_index_spec: Optional[DeltaSyncVectorIndexSpecRequest] = None, + direct_access_index_spec: Optional[DirectAccessVectorIndexSpec] = None, + endpoint_name: Optional[str] = None) -> CreateVectorIndexResponse: + """Create an index. + + Create a new index. + + :param name: str + Name of the index + :param primary_key: str + Primary key of the index + :param index_type: :class:`VectorIndexType` + There are 2 types of Vector Search indexes: + + - `DELTA_SYNC`: An index that automatically syncs with a source Delta Table, automatically and + incrementally updating the index as the underlying data in the Delta Table changes. - + `DIRECT_ACCESS`: An index that supports direct read and write of vectors and metadata through our + REST and SDK APIs. With this model, the user manages index updates. + :param delta_sync_vector_index_spec: :class:`DeltaSyncVectorIndexSpecRequest` (optional) + Specification for Delta Sync Index. Required if `index_type` is `DELTA_SYNC`. + :param direct_access_index_spec: :class:`DirectAccessVectorIndexSpec` (optional) + Specification for Direct Vector Access Index. Required if `index_type` is `DIRECT_ACCESS`. + :param endpoint_name: str (optional) + Name of the endpoint to be used for serving the index + + :returns: :class:`CreateVectorIndexResponse` + """ + body = {} + if delta_sync_vector_index_spec is not None: + body['delta_sync_vector_index_spec'] = delta_sync_vector_index_spec.as_dict() + if direct_access_index_spec is not None: + body['direct_access_index_spec'] = direct_access_index_spec.as_dict() + if endpoint_name is not None: body['endpoint_name'] = endpoint_name + if index_type is not None: body['index_type'] = index_type.value + if name is not None: body['name'] = name + if primary_key is not None: body['primary_key'] = primary_key + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', '/api/2.0/vector-search/indexes', body=body, headers=headers) + return CreateVectorIndexResponse.from_dict(res) + + def delete_data_vector_index(self, name: str, primary_keys: List[str]) -> DeleteDataVectorIndexResponse: + """Delete data from index. + + Handles the deletion of data from a specified vector index. + + :param name: str + Name of the vector index where data is to be deleted. Must be a Direct Vector Access Index. + :param primary_keys: List[str] + List of primary keys for the data to be deleted. + + :returns: :class:`DeleteDataVectorIndexResponse` + """ + body = {} + if primary_keys is not None: body['primary_keys'] = [v for v in primary_keys] + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', + f'/api/2.0/vector-search/indexes/{name}/delete-data', + body=body, + headers=headers) + return DeleteDataVectorIndexResponse.from_dict(res) + + def delete_index(self, index_name: str): + """Delete an index. + + Delete an index. + + :param index_name: str + Name of the index + + + """ + + headers = {} + self._api.do('DELETE', f'/api/2.0/vector-search/indexes/{index_name}', headers=headers) + + def get_index(self, index_name: str) -> VectorIndex: + """Get an index. + + Get an index. + + :param index_name: str + Name of the index + + :returns: :class:`VectorIndex` + """ + + headers = {'Accept': 'application/json', } + res = self._api.do('GET', f'/api/2.0/vector-search/indexes/{index_name}', headers=headers) + return VectorIndex.from_dict(res) + + def list_indexes(self, + endpoint_name: str, + *, + page_token: Optional[str] = None) -> Iterator[MiniVectorIndex]: + """List indexes. + + List all indexes in the given endpoint. + + :param endpoint_name: str + Name of the endpoint + :param page_token: str (optional) + Token for pagination + + :returns: Iterator over :class:`MiniVectorIndex` + """ + + query = {} + if endpoint_name is not None: query['endpoint_name'] = endpoint_name + if page_token is not None: query['page_token'] = page_token + headers = {'Accept': 'application/json', } + + while True: + json = self._api.do('GET', '/api/2.0/vector-search/indexes', query=query, headers=headers) + if 'vector_indexes' not in json or not json['vector_indexes']: + return + for v in json['vector_indexes']: + yield MiniVectorIndex.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 query_index(self, + index_name: str, + columns: List[str], + *, + filters_json: Optional[str] = None, + num_results: Optional[int] = None, + query_text: Optional[str] = None, + query_vector: Optional[List[float]] = None) -> QueryVectorIndexResponse: + """Query an index. + + Query the specified vector index. + + :param index_name: str + Name of the vector index to query. + :param columns: List[str] + List of column names to include in the response. + :param filters_json: str (optional) + JSON string representing query filters. + + Example filters: - `{"id <": 5}`: Filter for id less than 5. - `{"id >": 5}`: Filter for id greater + than 5. - `{"id <=": 5}`: Filter for id less than equal to 5. - `{"id >=": 5}`: Filter for id + greater than equal to 5. - `{"id": 5}`: Filter for id equal to 5. + :param num_results: int (optional) + Number of results to return. Defaults to 10. + :param query_text: str (optional) + Query text. Required for Delta Sync Index using model endpoint. + :param query_vector: List[float] (optional) + Query vector. Required for Direct Vector Access Index and Delta Sync Index using self-managed + vectors. + + :returns: :class:`QueryVectorIndexResponse` + """ + body = {} + if columns is not None: body['columns'] = [v for v in columns] + if filters_json is not None: body['filters_json'] = filters_json + if num_results is not None: body['num_results'] = num_results + if query_text is not None: body['query_text'] = query_text + if query_vector is not None: body['query_vector'] = [v for v in query_vector] + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', + f'/api/2.0/vector-search/indexes/{index_name}/query', + body=body, + headers=headers) + return QueryVectorIndexResponse.from_dict(res) + + def sync_index(self, index_name: str): + """Synchronize an index. + + Triggers a synchronization process for a specified vector index. + + :param index_name: str + Name of the vector index to synchronize. Must be a Delta Sync Index. + + + """ + + headers = {} + self._api.do('POST', f'/api/2.0/vector-search/indexes/{index_name}/sync', headers=headers) + + def upsert_data_vector_index(self, name: str, inputs_json: str) -> UpsertDataVectorIndexResponse: + """Upsert data into an index. + + Handles the upserting of data into a specified vector index. + + :param name: str + Name of the vector index where data is to be upserted. Must be a Direct Vector Access Index. + :param inputs_json: str + JSON string representing the data to be upserted. + + :returns: :class:`UpsertDataVectorIndexResponse` + """ + body = {} + if inputs_json is not None: body['inputs_json'] = inputs_json + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + res = self._api.do('POST', + f'/api/2.0/vector-search/indexes/{name}/upsert-data', + body=body, + headers=headers) + return UpsertDataVectorIndexResponse.from_dict(res) diff --git a/examples/external_locations/list_external_locations_on_aws.py b/examples/external_locations/list_external_locations_on_aws.py index d847088b..2fb3b4b0 100755 --- a/examples/external_locations/list_external_locations_on_aws.py +++ b/examples/external_locations/list_external_locations_on_aws.py @@ -1,5 +1,6 @@ from databricks.sdk import WorkspaceClient +from databricks.sdk.service import catalog w = WorkspaceClient() -all = w.external_locations.list() +all = w.external_locations.list(catalog.ListExternalLocationsRequest()) diff --git a/examples/storage_credentials/list_storage_credentials_on_aws.py b/examples/storage_credentials/list_storage_credentials_on_aws.py index fad05f4a..d12a0aa4 100755 --- a/examples/storage_credentials/list_storage_credentials_on_aws.py +++ b/examples/storage_credentials/list_storage_credentials_on_aws.py @@ -1,5 +1,6 @@ from databricks.sdk import WorkspaceClient +from databricks.sdk.service import catalog w = WorkspaceClient() -all = w.storage_credentials.list() +all = w.storage_credentials.list(catalog.ListStorageCredentialsRequest())