databricks-sdk 0.18.0__py3-none-any.whl → 0.19.0__py3-none-any.whl

databricks/sdk/service/compute.py

@@ -65,15 +65,33 @@ class AddInstanceProfile:
 
 
 @dataclass
-class AutoScale:
-    min_workers: int
-    """The minimum number of workers to which the cluster can scale down when underutilized. It is also
-    the initial number of workers the cluster will have after creation."""
+class Adlsgen2Info:
+    destination: str
+    """abfss destination, e.g.
+    `abfss://<container-name>@<storage-account-name>.dfs.core.windows.net/<directory-name>`."""
+
+    def as_dict(self) -> dict:
+        """Serializes the Adlsgen2Info into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.destination is not None: body['destination'] = self.destination
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> Adlsgen2Info:
+        """Deserializes the Adlsgen2Info from a dictionary."""
+        return cls(destination=d.get('destination', None))
+
 
-    max_workers: int
+@dataclass
+class AutoScale:
+    max_workers: Optional[int] = None
     """The maximum number of workers to which the cluster can scale up when overloaded. Note that
     `max_workers` must be strictly greater than `min_workers`."""
 
+    min_workers: Optional[int] = None
+    """The minimum number of workers to which the cluster can scale down when underutilized. It is also
+    the initial number of workers the cluster will have after creation."""
+
     def as_dict(self) -> dict:
         """Serializes the AutoScale into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -2213,7 +2231,7 @@ class DataSecurityMode(Enum):
 
 @dataclass
 class DbfsStorageInfo:
-    destination: Optional[str] = None
+    destination: str
     """dbfs destination, e.g. `dbfs:/my/path`"""
 
     def as_dict(self) -> dict:
@@ -2957,6 +2975,18 @@ class GcpAttributes:
     
     [GCP documentation]: https://cloud.google.com/compute/docs/disks/local-ssd#choose_number_local_ssds"""
 
+    use_preemptible_executors: Optional[bool] = None
+    """This field determines whether the spark executors will be scheduled to run on preemptible VMs
+    (when set to true) versus standard compute engine VMs (when set to false; default). Note: Soon
+    to be deprecated, use the availability field instead."""
+
+    zone_id: Optional[str] = None
+    """Identifier for the availability zone in which the cluster resides. This can be one of the
+    following: - "HA" => High availability, spread nodes across availability zones for a Databricks
+    deployment region [default] - "AUTO" => Databricks picks an availability zone to schedule the
+    cluster on. - A GCP availability zone => Pick One of the available zones for (machine type +
+    region) from https://cloud.google.com/compute/docs/regions-zones."""
+
     def as_dict(self) -> dict:
         """Serializes the GcpAttributes into a dictionary suitable for use as a JSON request body."""
         body = {}
@@ -2965,6 +2995,9 @@ class GcpAttributes:
         if self.google_service_account is not None:
             body['google_service_account'] = self.google_service_account
         if self.local_ssd_count is not None: body['local_ssd_count'] = self.local_ssd_count
+        if self.use_preemptible_executors is not None:
+            body['use_preemptible_executors'] = self.use_preemptible_executors
+        if self.zone_id is not None: body['zone_id'] = self.zone_id
         return body
 
     @classmethod
@@ -2973,7 +3006,9 @@ class GcpAttributes:
         return cls(availability=_enum(d, 'availability', GcpAvailability),
                    boot_disk_size=d.get('boot_disk_size', None),
                    google_service_account=d.get('google_service_account', None),
-                   local_ssd_count=d.get('local_ssd_count', None))
+                   local_ssd_count=d.get('local_ssd_count', None),
+                   use_preemptible_executors=d.get('use_preemptible_executors', None),
+                   zone_id=d.get('zone_id', None))
 
 
 class GcpAvailability(Enum):
@@ -2985,6 +3020,23 @@ class GcpAvailability(Enum):
     PREEMPTIBLE_WITH_FALLBACK_GCP = 'PREEMPTIBLE_WITH_FALLBACK_GCP'
 
 
+@dataclass
+class GcsStorageInfo:
+    destination: str
+    """GCS destination/URI, e.g. `gs://my-bucket/some-prefix`"""
+
+    def as_dict(self) -> dict:
+        """Serializes the GcsStorageInfo into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.destination is not None: body['destination'] = self.destination
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> GcsStorageInfo:
+        """Deserializes the GcsStorageInfo from a dictionary."""
+        return cls(destination=d.get('destination', None))
+
+
 @dataclass
 class GetClusterPermissionLevelsResponse:
     permission_levels: Optional[List[ClusterPermissionsDescription]] = None
@@ -3538,6 +3590,10 @@ class InitScriptExecutionDetailsStatus(Enum):
 
 @dataclass
 class InitScriptInfo:
+    abfss: Optional[Adlsgen2Info] = None
+    """destination needs to be provided. e.g. `{ "abfss" : { "destination" :
+    "abfss://<container-name>@<storage-account-name>.dfs.core.windows.net/<directory-name>" } }"""
+
     dbfs: Optional[DbfsStorageInfo] = None
     """destination needs to be provided. e.g. `{ "dbfs" : { "destination" : "dbfs:/home/cluster_log" }
     }`"""
@@ -3546,6 +3602,9 @@ class InitScriptInfo:
     """destination needs to be provided. e.g. `{ "file" : { "destination" : "file:/my/local/file.sh" }
     }`"""
 
+    gcs: Optional[GcsStorageInfo] = None
+    """destination needs to be provided. e.g. `{ "gcs": { "destination": "gs://my-bucket/file.sh" } }`"""
+
     s3: Optional[S3StorageInfo] = None
     """destination and either the region or endpoint need to be provided. e.g. `{ "s3": { "destination"
     : "s3://cluster_log_bucket/prefix", "region" : "us-west-2" } }` Cluster iam role is used to
@@ -3563,8 +3622,10 @@ class InitScriptInfo:
     def as_dict(self) -> dict:
         """Serializes the InitScriptInfo into a dictionary suitable for use as a JSON request body."""
         body = {}
+        if self.abfss: body['abfss'] = self.abfss.as_dict()
         if self.dbfs: body['dbfs'] = self.dbfs.as_dict()
         if self.file: body['file'] = self.file.as_dict()
+        if self.gcs: body['gcs'] = self.gcs.as_dict()
         if self.s3: body['s3'] = self.s3.as_dict()
         if self.volumes: body['volumes'] = self.volumes.as_dict()
         if self.workspace: body['workspace'] = self.workspace.as_dict()
@@ -3573,8 +3634,10 @@ class InitScriptInfo:
     @classmethod
     def from_dict(cls, d: Dict[str, any]) -> InitScriptInfo:
         """Deserializes the InitScriptInfo from a dictionary."""
-        return cls(dbfs=_from_dict(d, 'dbfs', DbfsStorageInfo),
+        return cls(abfss=_from_dict(d, 'abfss', Adlsgen2Info),
+                   dbfs=_from_dict(d, 'dbfs', DbfsStorageInfo),
                    file=_from_dict(d, 'file', LocalFileInfo),
+                   gcs=_from_dict(d, 'gcs', GcsStorageInfo),
                    s3=_from_dict(d, 's3', S3StorageInfo),
                    volumes=_from_dict(d, 'volumes', VolumesStorageInfo),
                    workspace=_from_dict(d, 'workspace', WorkspaceStorageInfo))
@@ -4428,7 +4491,7 @@ class ListSortOrder(Enum):
 
 @dataclass
 class LocalFileInfo:
-    destination: Optional[str] = None
+    destination: str
     """local file destination, e.g. `file:/my/local/file.sh`"""
 
     def as_dict(self) -> dict:
@@ -5027,6 +5090,11 @@ class RuntimeEngine(Enum):
 
 @dataclass
 class S3StorageInfo:
+    destination: str
+    """S3 destination, e.g. `s3://my-bucket/some-prefix` Note that logs will be delivered using cluster
+    iam role, please make sure you set cluster iam role and the role has write access to the
+    destination. Please also note that you cannot use AWS keys to deliver logs."""
+
     canned_acl: Optional[str] = None
     """(Optional) Set canned access control list for the logs, e.g. `bucket-owner-full-control`. If
     `canned_cal` is set, please make sure the cluster iam role has `s3:PutObjectAcl` permission on
@@ -5036,11 +5104,6 @@ class S3StorageInfo:
     for writing data, you may want to set `bucket-owner-full-control` to make bucket owner able to
     read the logs."""
 
-    destination: Optional[str] = None
-    """S3 destination, e.g. `s3://my-bucket/some-prefix` Note that logs will be delivered using cluster
-    iam role, please make sure you set cluster iam role and the role has write access to the
-    destination. Please also note that you cannot use AWS keys to deliver logs."""
-
     enable_encryption: Optional[bool] = None
     """(Optional) Flag to enable server side encryption, `false` by default."""
 
@@ -5371,7 +5434,7 @@ class UnpinCluster:
 
 @dataclass
 class VolumesStorageInfo:
-    destination: Optional[str] = None
+    destination: str
     """Unity Catalog Volumes file destination, e.g. `/Volumes/my-init.sh`"""
 
     def as_dict(self) -> dict:
@@ -5388,7 +5451,7 @@ class VolumesStorageInfo:
 
 @dataclass
 class WorkloadType:
-    clients: Optional[ClientsTypes] = None
+    clients: ClientsTypes
     """defined what type of clients can use the cluster. E.g. Notebooks, Jobs"""
 
     def as_dict(self) -> dict:
@@ -5405,7 +5468,7 @@ class WorkloadType:
 
 @dataclass
 class WorkspaceStorageInfo:
-    destination: Optional[str] = None
+    destination: str
     """workspace files destination, e.g. `/Users/user1@databricks.com/my-init.sh`"""
 
     def as_dict(self) -> dict:
@@ -6405,10 +6468,9 @@ class ClustersAPI:
 
         while True:
             json = self._api.do('POST', '/api/2.0/clusters/events', body=body, headers=headers)
-            if 'events' not in json or not json['events']:
-                return
-            for v in json['events']:
-                yield ClusterEvent.from_dict(v)
+            if 'events' in json:
+                for v in json['events']:
+                    yield ClusterEvent.from_dict(v)
             if 'next_page' not in json or not json['next_page']:
                 return
             body = json['next_page']
@@ -7733,10 +7795,9 @@ class PolicyFamiliesAPI:
 
         while True:
             json = self._api.do('GET', '/api/2.0/policy-families', query=query, headers=headers)
-            if 'policy_families' not in json or not json['policy_families']:
-                return
-            for v in json['policy_families']:
-                yield PolicyFamily.from_dict(v)
+            if 'policy_families' in json:
+                for v in json['policy_families']:
+                    yield PolicyFamily.from_dict(v)
             if 'next_page_token' not in json or not json['next_page_token']:
                 return
             query['page_token'] = json['next_page_token']

databricks/sdk/service/files.py

@@ -112,6 +112,43 @@ class Delete:
         return cls(path=d.get('path', None), recursive=d.get('recursive', None))
 
 
+@dataclass
+class DirectoryEntry:
+    file_size: Optional[int] = None
+    """The length of the file in bytes. This field is omitted for directories."""
+
+    is_directory: Optional[bool] = None
+    """True if the path is a directory."""
+
+    last_modified: Optional[int] = None
+    """Last modification time of given file in milliseconds since unix epoch."""
+
+    name: Optional[str] = None
+    """The name of the file or directory."""
+
+    path: Optional[str] = None
+    """The absolute path of the file or directory."""
+
+    def as_dict(self) -> dict:
+        """Serializes the DirectoryEntry into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.file_size is not None: body['file_size'] = self.file_size
+        if self.is_directory is not None: body['is_directory'] = self.is_directory
+        if self.last_modified is not None: body['last_modified'] = self.last_modified
+        if self.name is not None: body['name'] = self.name
+        if self.path is not None: body['path'] = self.path
+        return body
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, any]) -> DirectoryEntry:
+        """Deserializes the DirectoryEntry from a dictionary."""
+        return cls(file_size=d.get('file_size', None),
+                   is_directory=d.get('is_directory', None),
+                   last_modified=d.get('last_modified', None),
+                   name=d.get('name', None),
+                   path=d.get('path', None))
+
+
 @dataclass
 class DownloadResponse:
     contents: Optional[BinaryIO] = None
@@ -149,6 +186,28 @@ class FileInfo:
                    path=d.get('path', None))
 
 
+@dataclass
+class ListDirectoryResponse:
+    contents: Optional[List[DirectoryEntry]] = None
+    """Array of DirectoryEntry."""
+
+    next_page_token: Optional[str] = None
+    """A token, which can be sent as `page_token` to retrieve the next page."""
+
+    def as_dict(self) -> dict:
+        """Serializes the ListDirectoryResponse into a dictionary suitable for use as a JSON request body."""
+        body = {}
+        if self.contents: body['contents'] = [v.as_dict() for v in self.contents]
+        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]) -> ListDirectoryResponse:
+        """Deserializes the ListDirectoryResponse from a dictionary."""
+        return cls(contents=_repeated_dict(d, 'contents', DirectoryEntry),
+                   next_page_token=d.get('next_page_token', None))
+
+
 @dataclass
 class ListStatusResponse:
     files: Optional[List[FileInfo]] = None
@@ -505,59 +564,114 @@ class FilesAPI:
     def __init__(self, api_client):
         self._api = api_client
 
+    def create_directory(self, directory_path: str):
+        """Create a directory.
+        
+        Creates an empty directory. If called on an existing directory, the API returns a success response.
+        
+        :param directory_path: str
+          The absolute path of a directory.
+        
+        
+        """
+
+        headers = {}
+        self._api.do('PUT', f'/api/2.0/fs/directories{directory_path}', headers=headers)
+
     def delete(self, file_path: str):
-        """Delete a file or directory.
+        """Delete a file.
         
-        Deletes a file or directory.
+        Deletes a file.
         
         :param file_path: str
-          The absolute path of the file or directory in DBFS.
+          The absolute path of the file.
         
         
         """
 
         headers = {}
-        self._api.do('DELETE', f'/api/2.0/fs/files/{file_path}', headers=headers)
+        self._api.do('DELETE', f'/api/2.0/fs/files{file_path}', headers=headers)
+
+    def delete_directory(self, directory_path: str):
+        """Delete a directory.
+        
+        Deletes an empty directory. If the directory is not empty, the API returns a HTTP 400 error.
+        
+        :param directory_path: str
+          The absolute path of a directory.
+        
+        
+        """
+
+        headers = {}
+        self._api.do('DELETE', f'/api/2.0/fs/directories{directory_path}', headers=headers)
 
     def download(self, file_path: str) -> DownloadResponse:
         """Download a file.
         
-        Downloads a file of up to 2 GiB.
+        Downloads a file of up to 5 GiB.
         
         :param file_path: str
-          The absolute path of the file or directory in DBFS.
+          The absolute path of the file.
         
         :returns: :class:`DownloadResponse`
         """
 
         headers = {'Accept': 'application/octet-stream', }
-        res = self._api.do('GET', f'/api/2.0/fs/files/{file_path}', headers=headers, raw=True)
+        res = self._api.do('GET', f'/api/2.0/fs/files{file_path}', headers=headers, raw=True)
         return DownloadResponse(contents=res)
 
-    def get_status(self, path: str) -> FileInfo:
-        """Get file or directory status.
-        
-        Returns the status of a file or directory.
-        
-        :param path: str
-          The absolute path of the file or directory in the Files API, omitting the initial slash.
-        
-        :returns: :class:`FileInfo`
+    def list_directory_contents(self,
+                                directory_path: str,
+                                *,
+                                page_size: Optional[int] = None,
+                                page_token: Optional[str] = None) -> Iterator[DirectoryEntry]:
+        """List directory contents.
+        
+        Returns the contents of a directory. If there is no directory at the specified path, the API returns a
+        HTTP 404 error.
+        
+        :param directory_path: str
+          The absolute path of a directory.
+        :param page_size: int (optional)
+          The maximum number of directory entries to return. The API may return fewer than this value.
+          Receiving fewer results does not imply there are no more results. As long as the response contains a
+          next_page_token, there may be more results.
+          
+          If unspecified, at most 1000 directory entries will be returned. The maximum value is 1000. Values
+          above 1000 will be coerced to 1000.
+        :param page_token: str (optional)
+          A page token, received from a previous `list` call. Provide this to retrieve the subsequent page.
+          When paginating, all other parameters provided to `list` must match the call that provided the page
+          token.
+        
+        :returns: Iterator over :class:`DirectoryEntry`
         """
 
         query = {}
-        if path is not None: query['path'] = path
+        if page_size is not None: query['page_size'] = page_size
+        if page_token is not None: query['page_token'] = page_token
         headers = {'Accept': 'application/json', }
-        res = self._api.do('GET', '/api/2.0/fs/get-status', query=query, headers=headers)
-        return FileInfo.from_dict(res)
+
+        while True:
+            json = self._api.do('GET',
+                                f'/api/2.0/fs/directories{directory_path}',
+                                query=query,
+                                headers=headers)
+            if 'contents' in json:
+                for v in json['contents']:
+                    yield DirectoryEntry.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 upload(self, file_path: str, contents: BinaryIO, *, overwrite: Optional[bool] = None):
         """Upload a file.
         
-        Uploads a file of up to 2 GiB.
+        Uploads a file of up to 5 GiB.
         
         :param file_path: str
-          The absolute path of the file or directory in DBFS.
+          The absolute path of the file.
         :param contents: BinaryIO
         :param overwrite: bool (optional)
           If true, an existing file will be overwritten.
@@ -568,4 +682,4 @@ class FilesAPI:
         query = {}
         if overwrite is not None: query['overwrite'] = overwrite
         headers = {'Content-Type': 'application/octet-stream', }
-        self._api.do('PUT', f'/api/2.0/fs/files/{file_path}', query=query, headers=headers, data=contents)
+        self._api.do('PUT', f'/api/2.0/fs/files{file_path}', query=query, headers=headers, data=contents)

databricks/sdk/service/iam.py

@@ -1465,14 +1465,15 @@ class AccountGroupsAPI:
                                 f'/api/2.0/accounts/{self._api.account_id}/scim/v2/Groups',
                                 query=query,
                                 headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield Group.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield Group.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,
@@ -1705,14 +1706,15 @@ class AccountServicePrincipalsAPI:
                                 f'/api/2.0/accounts/{self._api.account_id}/scim/v2/ServicePrincipals',
                                 query=query,
                                 headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield ServicePrincipal.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield ServicePrincipal.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,
@@ -2005,14 +2007,15 @@ class AccountUsersAPI:
                                 f'/api/2.0/accounts/{self._api.account_id}/scim/v2/Users',
                                 query=query,
                                 headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield User.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield User.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,
@@ -2267,14 +2270,15 @@ class GroupsAPI:
         if "count" not in query: query['count'] = 100
         while True:
             json = self._api.do('GET', '/api/2.0/preview/scim/v2/Groups', query=query, headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield Group.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield Group.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,
@@ -2646,14 +2650,15 @@ class ServicePrincipalsAPI:
                                 '/api/2.0/preview/scim/v2/ServicePrincipals',
                                 query=query,
                                 headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield ServicePrincipal.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield ServicePrincipal.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,
@@ -2955,14 +2960,15 @@ class UsersAPI:
         if "count" not in query: query['count'] = 100
         while True:
             json = self._api.do('GET', '/api/2.0/preview/scim/v2/Users', query=query, headers=headers)
+            if 'Resources' in json:
+                for v in json['Resources']:
+                    i = v['id']
+                    if i in seen:
+                        continue
+                    seen.add(i)
+                    yield User.from_dict(v)
             if 'Resources' not in json or not json['Resources']:
                 return
-            for v in json['Resources']:
-                i = v['id']
-                if i in seen:
-                    continue
-                seen.add(i)
-                yield User.from_dict(v)
             query['startIndex'] += len(json['Resources'])
 
     def patch(self,