Update Databricks CI test (#135)

Author: potter-potter

.github/workflows/e2e.yml

@@ -172,9 +172,9 @@ jobs:
         ASTRA_DB_API_ENDPOINT: ${{secrets.ASTRA_DB_ENDPOINT}}
         CLARIFAI_API_KEY: ${{secrets.CLARIFAI_API_KEY}}
         DATABRICKS_HOST: ${{secrets.DATABRICKS_HOST}}
-        DATABRICKS_USERNAME: ${{secrets.DATABRICKS_USERNAME}}
-        DATABRICKS_PASSWORD: ${{secrets.DATABRICKS_PASSWORD}}
         DATABRICKS_CATALOG: ${{secrets.DATABRICKS_CATALOG}}
+        DATABRICKS_CLIENT_ID: ${{secrets.DATABRICKS_CLIENT_ID}}
+        DATABRICKS_CLIENT_SECRET: ${{secrets.DATABRICKS_CLIENT_SECRET}}
         SHAREPOINT_CLIENT_ID: ${{secrets.SHAREPOINT_CLIENT_ID}}
         SHAREPOINT_CRED: ${{secrets.SHAREPOINT_CRED}}
         KDBAI_BEARER_TOKEN: ${{ secrets.KDBAI_BEARER_TOKEN }}

CHANGELOG.md

@@ -1,8 +1,10 @@
-## 0.0.21-dev1
+## 0.0.21-dev2
 
 ### Fixes
 
-* **Fix forward compatibility issues with `unstructured-client==0.26.0`. Update syntax and create a new SDK util file for reuse in the Partitioner and Chunker
+* **Fix forward compatibility issues with `unstructured-client==0.26.0`.** Update syntax and create a new SDK util file for reuse in the Partitioner and Chunker
+
+* **Update Databricks CI Test** Update to use client_id and client_secret auth. Also return files.upload method to one from open source.
 
 * **fix astra src bug**
 

test_e2e/dest/databricks-volumes.sh

@@ -23,8 +23,8 @@ source "$SCRIPT_DIR"/cleanup.sh
 function cleanup() {
   python "$SCRIPT_DIR"/python/test-databricks-volumes.py cleanup \
     --host "$DATABRICKS_HOST" \
-    --username "$DATABRICKS_USERNAME" \
-    --password "$DATABRICKS_PASSWORD" \
+    --client-id "$DATABRICKS_CLIENT_ID" \
+    --client-secret "$DATABRICKS_CLIENT_SECRET" \
     --volume "$DATABRICKS_VOLUME" \
     --catalog "$DATABRICKS_CATALOG" \
     --volume-path "$DATABRICKS_VOLUME_PATH"
@@ -48,16 +48,16 @@ PYTHONPATH=. ./unstructured_ingest/main.py \
   --work-dir "$WORK_DIR" \
   databricks-volumes \
   --host "$DATABRICKS_HOST" \
-  --username "$DATABRICKS_USERNAME" \
-  --password "$DATABRICKS_PASSWORD" \
+  --client-id "$DATABRICKS_CLIENT_ID" \
+  --client-secret "$DATABRICKS_CLIENT_SECRET" \
   --volume "$DATABRICKS_VOLUME" \
   --catalog "$DATABRICKS_CATALOG" \
   --volume-path "$DATABRICKS_VOLUME_PATH"
 
 python "$SCRIPT_DIR"/python/test-databricks-volumes.py test \
   --host "$DATABRICKS_HOST" \
-  --username "$DATABRICKS_USERNAME" \
-  --password "$DATABRICKS_PASSWORD" \
+  --client-id "$DATABRICKS_CLIENT_ID" \
+  --client-secret "$DATABRICKS_CLIENT_SECRET" \
   --volume "$DATABRICKS_VOLUME" \
   --catalog "$DATABRICKS_CATALOG" \
   --volume-path "$DATABRICKS_VOLUME_PATH"

test_e2e/python/test-databricks-volumes.py

@@ -16,20 +16,20 @@ def _get_volume_path(catalog: str, volume: str, volume_path: str):
 
 @cli.command()
 @click.option("--host", type=str, required=True)
-@click.option("--username", type=str, required=True)
-@click.option("--password", type=str, required=True)
+@click.option("--client-id", type=str, required=True)
+@click.option("--client-secret", type=str, required=True)
 @click.option("--catalog", type=str, required=True)
 @click.option("--volume", type=str, required=True)
 @click.option("--volume-path", type=str, required=True)
 def test(
     host: str,
-    username: str,
-    password: str,
+    client_id: str,
+    client_secret: str,
     catalog: str,
     volume: str,
     volume_path: str,
 ):
-    client = WorkspaceClient(host=host, username=username, password=password)
+    client = WorkspaceClient(host=host, client_id=client_id, client_secret=client_secret)
     files = list(
         client.files.list_directory_contents(_get_volume_path(catalog, volume, volume_path))
     )
@@ -53,20 +53,20 @@ def test(
 
 @cli.command()
 @click.option("--host", type=str, required=True)
-@click.option("--username", type=str, required=True)
-@click.option("--password", type=str, required=True)
+@click.option("--client-id", type=str, required=True)
+@click.option("--client-secret", type=str, required=True)
 @click.option("--catalog", type=str, required=True)
 @click.option("--volume", type=str, required=True)
 @click.option("--volume-path", type=str, required=True)
 def cleanup(
     host: str,
-    username: str,
-    password: str,
+    client_id: str,
+    client_secret: str,
     catalog: str,
     volume: str,
     volume_path: str,
 ):
-    client = WorkspaceClient(host=host, username=username, password=password)
+    client = WorkspaceClient(host=host, client_id=client_id, client_secret=client_secret)
 
     for file in client.files.list_directory_contents(
         _get_volume_path(catalog, volume, volume_path)

test_e2e/test-dest.sh

@@ -22,6 +22,7 @@ all_tests=(
   'chroma.sh'
   'clarifai.sh'
   'couchbase.sh'
+  'databricks-volumes.sh'
   'delta-table.sh'
   'dropbox.sh'
   'elasticsearch.sh'
@@ -68,7 +69,6 @@ tests_to_ignore=(
   'notion.sh'
   'dropbox.sh'
   'sharepoint.sh'
-  'databricks-volumes.sh'
 )
 
 for test in "${all_tests[@]}"; do

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.0.21-dev1"  # pragma: no cover
+__version__ = "0.0.21-dev2"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/databricks_volumes.py

@@ -42,8 +42,10 @@ class DatabricksVolumesAccessConfig(AccessConfig):
         description="The Databricks password part of basic authentication. "
         "Only possible when Host is *.cloud.databricks.com (AWS).",
     )
-    client_id: Optional[str] = Field(default=None)
-    client_secret: Optional[str] = Field(default=None)
+    client_id: Optional[str] = Field(default=None, description="Client ID of the OAuth app.")
+    client_secret: Optional[str] = Field(
+        default=None, description="Client Secret of the OAuth app."
+    )
     token: Optional[str] = Field(
         default=None,
         description="The Databricks personal access token (PAT) (AWS, Azure, and GCP) or "
@@ -140,11 +142,12 @@ def precheck(self) -> None:
 
     def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
         output_path = os.path.join(self.upload_config.path, path.name)
-        self.get_client().files.upload(
-            file_path=output_path,
-            contents=path,
-            overwrite=self.upload_config.overwrite,
-        )
+        with open(path, "rb") as elements_file:
+            self.get_client().files.upload(
+                file_path=output_path,
+                contents=elements_file,
+                overwrite=self.upload_config.overwrite,
+            )
 
 
 databricks_volumes_destination_entry = DestinationRegistryEntry(