dasl-client 1.0.14__tar.gz → 1.0.16__tar.gz

{dasl_client-1.0.14 → dasl_client-1.0.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dasl_client
-Version: 1.0.14
+Version: 1.0.16
 Summary: The DASL client library used for interacting with the DASL workspace
 Home-page: https://github.com/antimatter/asl
 Author: Antimatter Team
@@ -11,8 +11,7 @@ License-File: LICENSE
 Requires-Dist: dasl_api==0.1.17
 Requires-Dist: databricks-sdk>=0.41.0
 Requires-Dist: pydantic>=2
-Requires-Dist: pyspark>=3.5.5
-Requires-Dist: typing_extensions==4.10.0
+Requires-Dist: typing_extensions>=4.10.0
 
 # DASL Client Library
 

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/client.py

@@ -1,4 +1,6 @@
 from copy import deepcopy
+from datetime import datetime, timedelta
+from time import sleep
 from typing import Any, Callable, Iterator, List, Optional, Tuple, TypeVar
 from pydantic import Field
 from pyspark.sql import DataFrame
@@ -608,6 +610,7 @@ class Client:
         self,
         warehouse: str,
         request: TransformRequest,
+        timeout: timedelta = timedelta(minutes=5),
     ) -> TransformResponse:
         """
         Run a sequence of ADHOC transforms against a SQL warehouse to
@@ -618,16 +621,29 @@ class Client:
         :return: a TransformResponse object containing the results
             after running the transforms.
         :raises: NotFoundError if the rule does not exist
+        :raises: Exception for a server-side error or timeout
         """
         with error_handler():
-            return TransformResponse.from_api_obj(
-                self._dbui_client().dbui_v1_transform(
-                    self._workspace(),
-                    warehouse,
-                    request.to_api_obj(),
-                )
+            status = self._dbui_client().dbui_v1_transform(
+                self._workspace(),
+                warehouse,
+                request.to_api_obj(),
             )
 
+            begin = datetime.now()
+            while datetime.now() - begin < timeout:
+                sleep(5)
+                status = self._dbui_client().dbui_v1_transform_status(
+                    self._workspace(), status.id
+                )
+
+                if status.status == "failure":
+                    raise Exception(f"adhoc transform failed with {status.error}")
+                elif status.status == "success":
+                    return TransformResponse.from_api_obj(status.result)
+
+            raise Exception("timed out waiting for adhoc transform result")
+
     def get_observable_events(
         self,
         warehouse: str,

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/preset_development/errors.py

@@ -149,7 +149,10 @@ class AssertionFailedError(PresetError):
     def __init__(self, expr: str, assertion_message: str, df: DataFrame):
         # Get the Databricks built-in functions out the namespace.
         ipython = get_ipython()
-        display = ipython.user_ns["display"]
+        if ipython is not None:
+            display = ipython.user_ns["display"]
+        else:
+            display = lambda x: x.show()
 
         self.expr = expr
         self.assertion_message = assertion_message

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/preset_development/preview_engine.py

@@ -54,10 +54,13 @@ class PreviewEngine:
                 self._preset.get("silver", None), self._pretransform_name
             )
 
+        self._bronze = None
         self._pre = None
         self._silver = []
         self._gold = []
-        self._result_df_map = {}
+        self._result_df_map: Tuple[
+            DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]
+        ] = (None, {}, {})
 
     def _validate_pretransform_name(
         self, silver: Dict[str, str], pretransform_name: str
@@ -181,8 +184,12 @@ class PreviewEngine:
 
         # Get the Databricks built-in functions out the namespace.
         ipython = get_ipython()
-        displayHTML = ipython.user_ns["displayHTML"]
-        display = ipython.user_ns["display"]
+        if ipython is not None:
+            displayHTML = ipython.user_ns["displayHTML"]
+            display = ipython.user_ns["display"]
+        else:
+            displayHTML = lambda x: print(x)
+            display = lambda x: x.show()
 
         def d(txt, lvl) -> None:
             displayHTML(
@@ -245,7 +252,7 @@ class PreviewEngine:
                 raise UnknownGoldTableError(name, gold_table_schema)
 
             # Performs the type check.
-            delta_df = self._spark.table(f"{gold_table_schema}.{name}").limit(0)
+            delta_df = self._spark.table(f"`{gold_table_schema}`.`{name}`").limit(0)
             unioned_df = delta_df.unionByName(df, allowMissingColumns=True)
 
             # Now we check no new columns.
@@ -286,7 +293,7 @@ class PreviewEngine:
             d("Resultant gold table preview", 3)
             display(unioned_df)
 
-    def evaluate(self, gold_table_schema: str) -> None:
+    def evaluate(self, gold_table_schema: str, display: bool = True) -> None:
         """
         Evaluates the loaded preset YAML using the input datasource configuration to load
         records. Finally, checks that the output from the Gold stages is compatible with
@@ -303,7 +310,9 @@ class PreviewEngine:
         ):
             if not any(
                 row.databaseName == schema_name
-                for row in self._spark.sql(f"SHOW SCHEMAS IN {catalog_name}").collect()
+                for row in self._spark.sql(
+                    f"SHOW SCHEMAS IN `{catalog_name}`"
+                ).collect()
             ):
                 raise InvalidGoldTableSchemaError(
                     gold_table_schema,
@@ -340,5 +349,12 @@ class PreviewEngine:
         self._compile_stages()
 
         with self._ds_params as df:
+            self._bronze = df
             self._result_df_map = self._run(df)
-            self._render_output(df, self._result_df_map, gold_table_schema)
+            if display:
+                self._render_output(df, self._result_df_map, gold_table_schema)
+
+    def results(
+        self,
+    ) -> Tuple[DataFrame, DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]]:
+        return self._bronze, *self._result_df_map

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/preset_development/preview_parameters.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.types import *
 from pyspark.sql.dataframe import DataFrame
@@ -115,6 +117,7 @@ class PreviewParameters:
         self._table = None
 
         self._pretransform_name = None
+        self._bronze_pre_transform: Optional[List[str]] = None
 
         self._df = None
 
@@ -166,10 +169,13 @@ class PreviewParameters:
                 self._record_limit
             )
 
+            if self._bronze_pre_transform is not None:
+                stream_df = stream_df.selectExpr(*self._bronze_pre_transform)
+
             query = (
                 stream_df.writeStream.format("memory")
                 .queryName("batch_data")
-                .trigger(once=True)
+                .trigger(availableNow=True)
                 .start()
             )
 
@@ -193,12 +199,17 @@ class PreviewParameters:
 
         # Get the Databricks built-in functions out the namespace.
         ipython = get_ipython()
-        dbutils = ipython.user_ns["dbutils"]
+        if ipython is not None:
+            dbutils = ipython.user_ns["dbutils"]
 
-        dbutils.fs.rm(
-            f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
-            recurse=True,
-        )
+            dbutils.fs.rm(
+                f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
+                recurse=True,
+            )
+        else:
+            print(
+                f"FYI, we are leaking temp data {self._autoloader_temp_schema_location}/{self._schema_uuid_str}"
+            )
 
     def from_input(self):
         """
@@ -286,6 +297,10 @@ class PreviewParameters:
         Returns:
             PreviewParameters: The current instance with updated configuration.
         """
+        if file_format.lower() == "jsonl":
+            self._autoloader_format = "json"
+            self.set_autoloader_multiline(False)
+            return self
         self._autoloader_format = file_format
         return self
 
@@ -350,6 +365,16 @@ class PreviewParameters:
         self._pretransform_name = pretransform_name
         return self
 
+    def set_bronze_pre_transform(self, expr: List[str]):
+        """
+        Sets a pre-transform expression that will run before data is written to bronze
+
+        Returns:
+            PreviewParameters: The current instance with updated configuration.
+        """
+        self._bronze_pre_transform = expr
+        return self
+
     def set_date_range(self, column: str, start_time: str, end_time: str):
         """
         Set the TIMESTAMP column and date range to use as the input data filter to

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/preset_development/stage.py

@@ -150,17 +150,19 @@ class Stage:
         if target_col not in existing_columns:
             raise ReferencedColumnMissingError("jsonExtract", target_col)
         schema = self._spark.sql(
-            f"SELECT schema_of_json_agg({target_col}) AS sc FROM {{df}}", df=df
+            f"SELECT schema_of_json_agg({self.auto_backtick(target_col)}) AS sc FROM {{df}}",
+            df=df,
         ).collect()[0][0]
         extract_df = self._spark.createDataFrame(data=[], schema=schema)
         columns = extract_df.columns
         columns = [
-            f"extract.{col} AS {col}"
+            self.auto_backtick(f"extract.{col}") + f" AS {self.auto_backtick(col)}"
             for col in columns
             if col not in omit_fields and col not in existing_columns
         ]
         columns += [
-            f"extract.{col} AS {duplicate_prefix}{col}"
+            self.auto_backtick(f"extract.{col}")
+            + f" AS {self.auto_backtick(duplicate_prefix + col)}"
             for col in columns
             if col not in omit_fields and col in existing_columns
         ]
@@ -176,7 +178,10 @@ class Stage:
             A DataFrame with the resultant operation's records.
         """
         return (
-            df.selectExpr("*", f"from_json({target_col}, '{schema}') AS extract")
+            df.selectExpr(
+                "*",
+                f"from_json({self.auto_backtick(target_col)}, '{schema}') AS extract",
+            )
             .selectExpr("*", *columns)
             .drop("extract")
         )
@@ -198,7 +203,10 @@ class Stage:
         """
         extract_df = self._spark.createDataFrame(data=[], schema=schema)
         schema = extract_df.drop(omit_fields).schema.simpleString()
-        return df.selectExpr("*", f"from_json({target_col}, '{schema}') AS {name}")
+        return df.selectExpr(
+            "*",
+            f"from_json({self.auto_backtick(target_col)}, '{schema}') AS {self.auto_backtick(name)}",
+        )
 
     def preserved_columns(
         self, df: DataFrame
@@ -224,7 +232,7 @@ class Stage:
         duplicate_prefix = self._duplicate_prefix()
         column_names = self._column_names()
         duplicate_renames = [
-            f"{col} AS {duplicate_prefix}{col}"
+            f"{self.auto_backtick(col)} AS {self.auto_backtick(duplicate_prefix + col)}"
             for col in preserved_columns
             if col in column_names
         ]
@@ -296,14 +304,46 @@ class Stage:
         """
         if field.get("from", None):
             # check that the from column exists in the df?
-            return f"{field['from']} AS {name}"
+            return f"{self.auto_backtick(field['from'])} AS {self.auto_backtick(name)}"
         elif field.get("literal", None):
-            return f"'{field['literal']}' AS {name}"
-        elif field.get("expr", None):
-            return f"{field['expr']} AS {name}"
+            return f"'{field['literal']}' AS {self.auto_backtick(name)}"
+        elif field.get("expr", None) is not None:
+            return f"{field['expr']} AS {self.auto_backtick(name)}"
         else:
             return ""
 
+    def is_backtick_escaped(self, name: str) -> bool:
+        """
+        check if a given (column) name is backtick escaped or not
+        :param name: column name
+        :return: bool
+        """
+        return name.startswith("`") and name.endswith("`")
+
+    def auto_backtick(self, name: str) -> str:
+        """
+        auto-backtick given name in case it isn't already backtick escaped.
+        if the name contains dots it will get split and each component backticked individually.
+        Returns the name wrapped in backticks or the passed name if it already had backticks.
+        :param name: column name
+        :return: str
+        """
+        if self.is_backtick_escaped(name):
+            return name
+        parts = name.split(".")
+        return ".".join(list(map(lambda s: f"`{s}`", parts)))
+
+    def force_apply_backticks(self, name: str) -> str:
+        """
+        forces application of backticks to the given (column) name as a single unit
+        if it already has backticks this is a noop
+        :param name: column name
+        :return: str
+        """
+        if self.is_backtick_escaped(name):
+            return name
+        return f"`{name}`"
+
     def process_node(self, name: str, node: Node) -> str:
         """
         Processes a single node in a tree of nodes.
@@ -319,7 +359,7 @@ class Stage:
                 child_expr = self.process_node(child_name, child_node)
                 fields_list.append(f"{child_expr}")
             joined_fields = ",\n".join(fields_list)
-            return f"struct(\n{joined_fields}\n) AS {name}"
+            return f"struct(\n{joined_fields}\n) AS {self.auto_backtick(name)}"
         else:
             return ""
 
@@ -341,12 +381,22 @@ class Stage:
         """
         Renders a list of field specifications containing both simple and
         STRUCT references into valid, STRUCT cognicient, SELECT expressions.
+        if a nested field is wrapped in backticks it will be treated as a simple field
+        for example field of name `col.with.dots` will NOT be treated as nested field.
 
         Returns:
             The SQL expression.
         """
-        simple_fields = [f for f in fields if "." not in f["name"]]
-        nested_fields = [f for f in fields if "." in f["name"]]
+        simple_fields = [
+            f
+            for f in fields
+            if "." not in f["name"] or self.is_backtick_escaped(f["name"])
+        ]
+        nested_fields = [
+            f
+            for f in fields
+            if "." in f["name"] and not self.is_backtick_escaped(f["name"])
+        ]
 
         result_parts = []
         for field in simple_fields:
@@ -358,7 +408,7 @@ class Stage:
             nested_str = self.parse_to_string(tree)
             result_parts.append(nested_str)
 
-        return [p for p in result_parts if p]
+        return [p for p in result_parts if p is not None and len(p) > 0]
 
     def select_expr(self, df: DataFrame) -> str:
         """
@@ -379,8 +429,12 @@ class Stage:
             if should_preserve:
                 if embed_col := preserve.get("embedColumn", None):
                     preserved_columns = self.preserved_columns_embed_column(df)
+                    # preserved_columns is obtained from df.columns
+                    # applying backticks to all of them is OK here
+                    # since they will never use "obj.key" to reference nested fields of structs
+                    # so we just go ahead and apply backticks to all across the board
                     select_fields += [
-                        f"struct({', '.join(preserved_columns)}) AS {embed_col}"
+                        f"struct({', '.join(list(map(lambda x: self.force_apply_backticks(x), preserved_columns)))}) AS {self.auto_backtick(embed_col)}"
                     ]
                 else:
                     (
@@ -388,8 +442,13 @@ class Stage:
                         duplicate_renames,
                         column_names,
                     ) = self.preserved_columns(df)
-                    select_fields += preserved_columns
-                    select_fields += duplicate_renames
+                    # see note above: same here - apply backticks to all columns across the board
+                    select_fields += list(
+                        map(lambda x: self.force_apply_backticks(x), preserved_columns)
+                    )
+                    select_fields += list(
+                        map(lambda x: self.force_apply_backticks(x), duplicate_renames)
+                    )
 
         return ["*"] + select_fields if self._stage == "temp_fields" else select_fields
 
@@ -475,7 +534,9 @@ class Stage:
                 df = (
                     df.alias("tmp")
                     .join(df_joined, on=[df[lhs] == df_joined[rhs]], how="left")
-                    .selectExpr("tmp.*", f"{select} AS {field.get('name')}")
+                    .selectExpr(
+                        "tmp.*", f"{select} AS {self.auto_backtick(field.get('name'))}"
+                    )
                 )
             elif csv := join.get("withCSV", None):
                 if path := csv.get("path", None):
@@ -485,7 +546,10 @@ class Stage:
                     df = (
                         df.alias("tmp")
                         .join(df_joined, on=[df[lhs] == df_joined[rhs]], how="left")
-                        .selectExpr("tmp.*", f"{select} AS {field.get('name')}")
+                        .selectExpr(
+                            "tmp.*",
+                            f"{select} AS {self.auto_backtick(field.get('name'))}",
+                        )
                     )
             else:
                 raise MissingJoinFieldError("withTable or withCSV (please supply 1)")
@@ -500,7 +564,10 @@ class Stage:
         """
         for field in self._fields:
             if field.get("alias", None):
-                df = df.selectExpr("*", f"{field.get('alias')} AS {field.get('name')}")
+                df = df.selectExpr(
+                    "*",
+                    f"{self.auto_backtick(field.get('alias'))} AS {self.auto_backtick(field.get('name'))}",
+                )
         return df
 
     def run_assertions(self, df: DataFrame) -> DataFrame:

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client/types/dbui.py

@@ -11,7 +11,8 @@ from dasl_api import (
     DbuiV1TransformRequestTransformsInner,
     DbuiV1TransformRequestTransformsInnerPresetOverrides,
     DbuiV1TransformResponse,
-    DbuiV1TransformResponseStagesInner, ContentV1DatasourcePresetAutoloaderCloudFiles,
+    DbuiV1TransformResponseStagesInner,
+    ContentV1DatasourcePresetAutoloaderCloudFiles,
     DbuiV1TransformRequestAutoloaderInput,
 )
 
@@ -202,7 +203,7 @@ class TransformRequest(BaseModel):
 
             @staticmethod
             def from_api_obj(
-                    obj: Optional[ContentV1DatasourcePresetAutoloaderCloudFiles],
+                obj: Optional[ContentV1DatasourcePresetAutoloaderCloudFiles],
             ) -> "TransformRequest.Autoloader.CloudFiles":
                 if obj is None:
                     return None
@@ -227,7 +228,7 @@ class TransformRequest(BaseModel):
 
         @staticmethod
         def from_api_obj(
-                obj: Optional[DbuiV1TransformRequestAutoloaderInput],
+            obj: Optional[DbuiV1TransformRequestAutoloaderInput],
         ) -> "Optional[TransformRequest.Autoloader]":
             if obj is None:
                 return None
@@ -235,10 +236,13 @@ class TransformRequest(BaseModel):
                 format=obj.format,
                 location=obj.location,
                 schema_file=obj.schema_file,
-                cloud_files=TransformRequest.Autoloader.CloudFiles.from_api_obj(obj.cloud_files),
+                cloud_files=TransformRequest.Autoloader.CloudFiles.from_api_obj(
+                    obj.cloud_files
+                ),
                 row_count=obj.row_count,
                 row_offset=obj.row_offset,
             )
+
         def to_api_obj(self) -> DbuiV1TransformRequestAutoloaderInput:
             return DbuiV1TransformRequestAutoloaderInput(
                 format=self.format,
@@ -247,10 +251,9 @@ class TransformRequest(BaseModel):
                 schema=self.schema_file,
                 cloudFiles=Helpers.maybe(lambda o: o.to_api_obj(), self.cloud_files),
                 rowCount=self.row_count,
-                rowOffset=self.row_offset
+                rowOffset=self.row_offset,
             )
 
-
     class Transform(BaseModel):
         """
         A transform configuration to apply to the data.
@@ -360,7 +363,9 @@ class TransformRequest(BaseModel):
     def from_api_obj(obj: DbuiV1TransformRequest) -> "TransformRequest":
         return TransformRequest(
             input=TransformRequest.Input.from_api_obj(obj.input),
-            autoloader_input=TransformRequest.Autoloader.from_api_obj(obj.autoloader_input),
+            autoloader_input=TransformRequest.Autoloader.from_api_obj(
+                obj.autoloader_input
+            ),
             use_preset=obj.use_preset,
             transforms=[
                 TransformRequest.Transform.from_api_obj(item) for item in obj.transforms

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dasl_client
-Version: 1.0.14
+Version: 1.0.16
 Summary: The DASL client library used for interacting with the DASL workspace
 Home-page: https://github.com/antimatter/asl
 Author: Antimatter Team
@@ -11,8 +11,7 @@ License-File: LICENSE
 Requires-Dist: dasl_api==0.1.17
 Requires-Dist: databricks-sdk>=0.41.0
 Requires-Dist: pydantic>=2
-Requires-Dist: pyspark>=3.5.5
-Requires-Dist: typing_extensions==4.10.0
+Requires-Dist: typing_extensions>=4.10.0
 
 # DASL Client Library
 

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client.egg-info/SOURCES.txt

@@ -31,6 +31,10 @@ dasl_client/types/helpers.py
 dasl_client/types/rule.py
 dasl_client/types/types.py
 dasl_client/types/workspace_config.py
+test/__init__.py
+test/conftest.py
+test/constants.py
 test/test_api_changes.py
 test/test_api_surface.py
+test/test_databricks_secret_auth.py
 test/test_marshaling.py

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client.egg-info/requires.txt

@@ -1,5 +1,4 @@
 dasl_api==0.1.17
 databricks-sdk>=0.41.0
 pydantic>=2
-pyspark>=3.5.5
-typing_extensions==4.10.0
+typing_extensions>=4.10.0

{dasl_client-1.0.14 → dasl_client-1.0.16}/dasl_client.egg-info/top_level.txt

@@ -1 +1,2 @@
 dasl_client
+test

{dasl_client-1.0.14 → dasl_client-1.0.16}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dasl_client"
-version = "1.0.14"
+version = "1.0.16"
 description = "The DASL client library used for interacting with the DASL workspace"
 readme = "README.md"
 authors = [
@@ -16,6 +16,5 @@ dependencies = [
     "dasl_api==0.1.17",
     "databricks-sdk>=0.41.0",
     "pydantic>=2",
-    "pyspark>=3.5.5",
-    "typing_extensions==4.10.0",
+    "typing_extensions>=4.10.0",
 ]

dasl_client-1.0.16/test/conftest.py

@@ -0,0 +1,18 @@
+import pytest
+
+from dasl_client import Client
+
+from .constants import *
+
+
+@pytest.fixture(scope="session")
+def api_client():
+    client = Client.new_workspace(
+        admin_email="test@antimatter.io",
+        app_client_id=app_client_id,
+        service_principal_id=databricks_client_id,
+        service_principal_secret=databricks_client_secret,
+        workspace_url=databricks_host,
+        dasl_host=dasl_host,
+    )
+    yield client

dasl_client-1.0.16/test/constants.py

@@ -0,0 +1,10 @@
+import os
+from urllib.parse import urlparse
+
+dasl_host = os.environ["DASL_API_URL"]
+databricks_host = os.environ["DASL_DATABRICKS_HOST"]
+databricks_client_id = os.environ["DASL_DATABRICKS_CLIENT_ID"]
+databricks_client_secret = os.environ["DASL_DATABRICKS_CLIENT_SECRET"]
+workspace = urlparse(databricks_host).hostname
+app_client_id = "22853b93-68ba-4ae2-8e41-976417f501dd"
+alternate_app_client_id = "335ac0d3-e0ea-4732-ba93-0277423b5029"

{dasl_client-1.0.14 → dasl_client-1.0.16}/test/test_api_surface.py

@@ -1,29 +1,6 @@
-import os
-import pytest
-from urllib.parse import urlparse
-
 from dasl_client import *
 
-dasl_host = os.environ["DASL_API_URL"]
-databricks_host = os.environ["DATABRICKS_HOST"]
-databricks_client_id = os.environ["DATABRICKS_CLIENT_ID"]
-databricks_client_secret = os.environ["DATABRICKS_CLIENT_SECRET"]
-workspace = urlparse(databricks_host).hostname
-app_client_id = "22853b93-68ba-4ae2-8e41-976417f501dd"
-alternate_app_client_id = "335ac0d3-e0ea-4732-ba93-0277423b5029"
-
-
-@pytest.fixture(scope="session")
-def api_client():
-    client = Client.new_workspace(
-        admin_email="test@antimatter.io",
-        app_client_id=app_client_id,
-        service_principal_id=databricks_client_id,
-        service_principal_secret=databricks_client_secret,
-        workspace_url=databricks_host,
-        dasl_host=dasl_host,
-    )
-    yield client
+from .constants import *
 
 
 def test_admin_config(api_client):

dasl_client-1.0.16/test/test_databricks_secret_auth.py

@@ -0,0 +1,116 @@
+import base64
+import datetime
+import os
+import time
+
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service import jobs, workspace as dbworkspace
+
+from .constants import *
+
+pylib_volume_path = os.environ["PYLIB_VOLUME_PATH"]
+pylib_wheel_path = os.environ["PYLIB_WHEEL_PATH"]
+
+
+def test_secret_auth(api_client):
+    # making sure it's even possible to get a config
+    api_client.get_config()
+
+    # need to do an API operation using databricks secret auth.
+    notebook_data = f"""
+    %pip install {pylib_wheel_path}
+    dbutils.library.restartPython()
+    # COMMAND ----------
+    from dasl_client.client import Client
+
+    Client.for_workspace(
+        workspace_url="{databricks_host}",
+        dasl_host="{dasl_host}",
+    ).get_config()
+    # COMMAND ----------
+    dbutils.notebook.exit("SUCCESS")
+    """
+    print(f"notebook_data={notebook_data}")
+
+    wsc = WorkspaceClient()
+    wsc.workspace.mkdirs(path=pylib_volume_path)
+
+    notebook_path = f"{pylib_volume_path}/test_secret_auth_notebook"
+    wsc.workspace.import_(
+        path=notebook_path,
+        format=dbworkspace.ImportFormat.SOURCE,
+        language=dbworkspace.Language.PYTHON,
+        content=base64.b64encode(notebook_data.encode("utf-8")).decode("utf-8"),
+        overwrite=True,
+    )
+
+    job_id = None
+    try:
+        job_id = wsc.jobs.create(
+            name="run test_secret_auth notebook",
+            tasks=[
+                jobs.Task(
+                    task_key="run_notebook",
+                    notebook_task=jobs.NotebookTask(notebook_path=notebook_path),
+                ),
+            ],
+        ).job_id
+
+        wsc.jobs.run_now(job_id=job_id)
+
+        logs = []
+        start = datetime.datetime.now()
+        complete = False
+        while not complete:
+            elapsed = datetime.datetime.now() - start
+            if elapsed > datetime.timedelta(seconds=300):
+                raise Exception(f"timed out waiting for job")
+
+            time.sleep(5)
+
+            status, logs = fetch_latest_run_status_and_logs(wsc, job_id)
+            print(f"logs={logs}")
+
+            if status == jobs.TerminationCodeCode.RUN_EXECUTION_ERROR:
+                raise Exception(f"job terminated with error")
+
+            complete = status == jobs.TerminationCodeCode.SUCCESS
+
+        print(logs)
+        assert len(logs) == 1
+        assert logs[0] == "SUCCESS"
+    finally:
+        wsc.workspace.delete(pylib_volume_path, recursive=True)
+        if job_id is not None:
+            wsc.jobs.delete(job_id=job_id)
+
+
+def fetch_latest_run_status_and_logs(
+    wsc: WorkspaceClient,
+    job_id: str,
+):
+    runs = list(wsc.jobs.list_runs(job_id=job_id, expand_tasks=True))
+    if not runs:
+        return "No runs found", None
+
+    # Find the latest run based on the start time
+    latest_run = max(runs, key=lambda r: r.start_time)
+    if latest_run.status.termination_details is None:
+        return "No runs found", None
+    status = latest_run.status.termination_details.code
+    logs = []
+    for task in latest_run.tasks:
+        output = wsc.jobs.get_run_output(task.run_id)
+        if output.error is not None:
+            logs.append(output.error)
+        elif output.logs is not None:
+            logs.append(output.logs)
+        elif output.notebook_output is not None:
+            logs.append(output.notebook_output.result)
+        elif output.run_job_output is not None:
+            raise Exception("Nested jobs are not supported")
+        elif output.sql_output is not None:
+            raise Exception("SQL jobs are unsupported")
+        else:
+            logs.append("")
+    return status, logs

{dasl_client-1.0.14 → dasl_client-1.0.16}/test/test_marshaling.py

@@ -1,9 +1,6 @@
-import os
-
 from dasl_client import *
 
-databricks_host = os.environ["DATABRICKS_HOST"]
-workspace = databricks_host.split("//")[1]
+from .constants import *
 
 
 def test_workspace_config_marshal_unmarshal():
@@ -701,13 +698,15 @@ def test_transform_request_marshal_unmarshal():
             ],
             data=[{"col1": "1", "col2": "a"}, {"col1": "2", "col2": "b"}],
         ),
-        autoloader_input=DataSource.Autoloader(
+        autoloader_input=TransformRequest.Autoloader(
             format="csv",
             location="s3://bucket/data",
             schema_file="schema.json",
-            cloud_files=DataSource.Autoloader.CloudFiles(
+            cloud_files=TransformRequest.Autoloader.CloudFiles(
                 schema_hints_file="hints_file.csv", schema_hints="hint1, hint2"
             ),
+            row_count=1,
+            row_offset=5,
         ),
         use_preset="preset_value",
         transforms=[