update

Author: parthban-db

tests/integration/conftest.py

@@ -115,36 +115,58 @@ def inner(var: str) -> str:
 
 @pytest.fixture(scope="session")
 def schema(ucws, random):
-    schema = ucws.schemas.create("dbfs-" + random(), "main")
+    from databricks.sdk.catalog.v2.client import SchemasClient
+
+    sc = SchemasClient(config=ucws)
+    schema = sc.create("dbfs-" + random(), "main")
     yield schema
-    ucws.schemas.delete(schema.full_name)
+    sc.delete(schema.full_name)
 
 
 @pytest.fixture(scope="session")
 def volume(ucws, schema):
-    volume = ucws.volumes.create("main", schema.name, "dbfs-test", VolumeType.MANAGED)
+    from databricks.sdk.catalog.v2.client import VolumesClient
+
+    vc = VolumesClient(config=ucws)
+    volume = vc.create("main", schema.name, "dbfs-test", VolumeType.MANAGED)
     yield "/Volumes/" + volume.full_name.replace(".", "/")
-    ucws.volumes.delete(volume.full_name)
+    vc.delete(volume.full_name)
 
 
 @pytest.fixture(scope="session", params=[False, True])
 def files_api(request, ucws) -> FilesAPI:
+
     if request.param:
         # ensure new Files API client is used for files of any size
-        ucws.config.multipart_upload_min_stream_size = 0
+        ucws.multipart_upload_min_stream_size = 0
         # enable new Files API client
-        return FilesExt(ucws.api_client, ucws.config)
+        from databricks.sdk.databricks.core import ApiClient
+
+        client = ApiClient(cfg=ucws)
+
+        return FilesExt(client, ucws)
     else:
         # use the default client
-        return ucws.files
+        from databricks.sdk.databricks.core import ApiClient
+        from databricks.sdk.files.v2.files import FilesAPI
+
+        client = ApiClient(cfg=ucws)
+
+        api = FilesAPI(api_client=client)
+        return api
 
 
 @pytest.fixture()
 def workspace_dir(w, random):
-    directory = f"/Users/{w.current_user.me().user_name}/dir-{random(12)}"
-    w.workspace.mkdirs(directory)
+    from databricks.sdk.iam.v2.client import CurrentUserClient
+    from databricks.sdk.workspace.v2.client import WorkspaceClient
+
+    cuc = CurrentUserClient(config=w)
+    wc = WorkspaceClient(config=w)
+    directory = f"/Users/{cuc.me().user_name}/dir-{random(12)}"
+    wc.mkdirs(directory)
     yield directory
-    w.workspace.delete(directory, recursive=True)
+    wc.delete(directory, recursive=True)
 
 
 def _load_debug_env_if_runs_from_ide(key) -> bool:

tests/integration/test_client.py

@@ -6,9 +6,10 @@ def test_get_workspace_client(ucacct, env_or_skip):
     assert w.current_user.me().active
 
 
-def test_get_workspace_id(ucws, env_or_skip):
-    ws_id = int(env_or_skip("THIS_WORKSPACE_ID"))
-    assert ucws.get_workspace_id() == ws_id
+# TODO: need to decide if we want get_workspace_id on every client
+# def test_get_workspace_id(ucws, env_or_skip):
+#     ws_id = int(env_or_skip("THIS_WORKSPACE_ID"))
+#     assert ucws.get_workspace_id() == ws_id
 
 
 def test_creating_ws_client_from_ac_client_does_not_override_config(ucacct, env_or_skip):

tests/integration/test_clusters.py

@@ -1,4 +1,3 @@
-
 import pytest
 
 from databricks.sdk.compute.v2 import compute
@@ -28,6 +27,7 @@ def test_cluster_events(w, env_or_skip):
         count += 1
     assert count > 0
 
+
 # TODO: Re-enable this test after adding waiters to the SDK
 # def test_ensure_cluster_is_running(w, env_or_skip):
 #     cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID")

tests/integration/test_data_plane.py

@@ -1,4 +1,3 @@
-
 # TODO: Re-enable this after adding data plane services to the SDK
 # def test_data_plane_token_source(ucws, env_or_skip):
 #     endpoint = env_or_skip("SERVING_ENDPOINT_NAME")

tests/integration/test_dbutils.py

@@ -5,8 +5,8 @@
 import pytest
 
 from databricks.sdk.databricks.core import DatabricksError
-from databricks.sdk.databricks.errors import NotFound
 from databricks.sdk.databricks.dbutils import RemoteDbUtils
+from databricks.sdk.databricks.errors import NotFound
 
 
 def test_rest_dbfs_ls(w, env_or_skip):
@@ -16,9 +16,10 @@ def test_rest_dbfs_ls(w, env_or_skip):
 
     assert len(x) > 1
 
+
 # TODO: Re-enable this test after adding waiters to the SDK
 # def test_proxy_dbfs_mounts(w, env_or_skip):
-    
+
 #     w.cluster_id = env_or_skip("TEST_DEFAULT_CLUSTER_ID")
 
 #     dbu = RemoteDbUtils(config=w)
@@ -29,11 +30,12 @@ def test_rest_dbfs_ls(w, env_or_skip):
 
 @pytest.fixture(params=["dbfs", "volumes"])
 def fs_and_base_path(request, ucws, volume):
+    dbu = RemoteDbUtils(config=ucws)
     if request.param == "dbfs":
-        fs = ucws.dbutils.fs
+        fs = dbu.fs
         base_path = "/tmp"
     else:
-        fs = ucws.dbutils.fs
+        fs = dbu.fs
         base_path = volume
     return fs, base_path
 
@@ -189,7 +191,7 @@ def test_secrets(w, random):
     logger.info(f"Before loading secret: {random_value}")
 
     from databricks.sdk.workspace.v2.client import SecretsClient
-    
+
     sc = SecretsClient(config=w)
     sc.create_scope(random_scope)
     sc.put_secret(random_scope, key_for_string, string_value=random_value)

tests/integration/test_files.py

@@ -7,9 +7,11 @@
 
 import pytest
 
+from databricks.sdk.catalog.v2.client import SchemasClient
 from databricks.sdk.databricks.core import DatabricksError
 from databricks.sdk.files.v2.client import DbfsClient
 
+
 def test_local_io(random):
     if platform.system() == "Windows":
         dummy_file = f"C:\\Windows\\Temp\\{random()}"
@@ -231,8 +233,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     @staticmethod
     def create_schema(w, catalog, schema):
-        res = w.schemas.create(catalog_name=catalog, name=schema)
-        return ResourceWithCleanup(lambda: w.schemas.delete(res.full_name))
+        sc = SchemasClient(config=w)
+        res = sc.create(catalog_name=catalog, name=schema)
+        return ResourceWithCleanup(lambda: sc.delete(res.full_name))
 
     @staticmethod
     def create_volume(w, catalog, schema, volume):

tests/integration/test_iam.py

@@ -1,10 +1,11 @@
 import pytest
 
 from databricks.sdk.databricks import errors
-from databricks.sdk.databricks.core import DatabricksError
-from databricks.sdk.iam.v2.client import AccountGroupsClient, AccountUsersClient, AccountServicePrincipalsClient
-from databricks.sdk.iam.v2.client import GroupsClient, UsersClient, ServicePrincipalsClient
-from databricks.sdk.databricks.core import ApiClient
+from databricks.sdk.databricks.core import ApiClient, DatabricksError
+from databricks.sdk.iam.v2.client import (AccountGroupsClient,
+                                          AccountServicePrincipalsClient,
+                                          AccountUsersClient, GroupsClient,
+                                          ServicePrincipalsClient, UsersClient)
 
 
 def test_filtering_groups(w, random):

tests/integration/test_workspace.py

@@ -1,8 +1,9 @@
 import io
 
-from databricks.sdk.workspace.v2.workspace import ImportFormat, Language
-from databricks.sdk.workspace.v2.client import WorkspaceClient
 from databricks.sdk.iam.v2.client import CurrentUserClient
+from databricks.sdk.workspace.v2.client import WorkspaceClient
+from databricks.sdk.workspace.v2.workspace import ImportFormat, Language
+
 
 def test_workspace_recursive_list(w, workspace_dir, random):
     wc = WorkspaceClient(config=w)
@@ -25,9 +26,9 @@ def test_workspace_recursive_list(w, workspace_dir, random):
 def test_workspace_upload_download_notebooks(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py"
-    
+
     wc.upload(notebook, io.BytesIO(b"print(1)"))
     with wc.download(notebook) as f:
         content = f.read()
@@ -39,7 +40,7 @@ def test_workspace_upload_download_notebooks(w, random):
 def test_workspace_unzip_notebooks(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py"
 
     # Big notebooks can be gzipped during transfer by the API (out of our control)
@@ -58,7 +59,7 @@ def test_workspace_unzip_notebooks(w, random):
 def test_workspace_download_connection_closed(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     notebook = f"/Users/{cuc.me().user_name}/notebook-{random(12)}.py"
 
     wc.upload(notebook, io.BytesIO(b"print(1)"))
@@ -74,7 +75,7 @@ def test_workspace_download_connection_closed(w, random):
 def test_workspace_upload_download_files(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     py_file = f"/Users/{cuc.me().user_name}/file-{random(12)}.py"
 
     wc.upload(py_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO)
@@ -88,7 +89,7 @@ def test_workspace_upload_download_files(w, random):
 def test_workspace_upload_download_txt_files(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     txt_file = f"/Users/{cuc.me().user_name}/txt-{random(12)}.txt"
 
     wc.upload(txt_file, io.BytesIO(b"print(1)"), format=ImportFormat.AUTO)
@@ -102,7 +103,7 @@ def test_workspace_upload_download_txt_files(w, random):
 def test_workspace_upload_download_notebooks_no_extension(w, random):
     wc = WorkspaceClient(config=w)
     cuc = CurrentUserClient(config=w)
-    
+
     nb = f"/Users/{cuc.me().user_name}/notebook-{random(12)}"
 
     wc.upload(

tests/test_client.py

@@ -1,5 +1,3 @@
-
-
 # from databricks.sdk import WorkspaceClient