dasl-client 1.0.7__py3-none-any.whl → 1.0.11__py3-none-any.whl

dasl_client/auth/auth.py

@@ -1,4 +1,5 @@
 import abc
+import base64
 import time
 from datetime import datetime
 
@@ -176,15 +177,15 @@ class DatabricksSecretAuth(Authorization):
             principalName=self._principal,
         )
         handler = api.WorkspaceV1Api(api_client=self._client)
-        resp = handler.workspace_v1_begin_authentication(
-            workspace=self._workspace, workspace_v1_begin_authenticate_request=req
+        resp = handler.workspace_v1_request_secret(
+            workspace=self._workspace, workspace_v1_request_secret_request=req
         )
         secret_name = resp.secret_name
         secret_value = ""
         for tries in range(3):
             try:
-                secret_value = WorkspaceClient().secrets.get_secret(
-                    secret_name, "token"
+                secret_value = (
+                    WorkspaceClient().secrets.get_secret(secret_name, "token").value
                 )
                 break
             except ResourceDoesNotExist:
@@ -194,7 +195,7 @@ class DatabricksSecretAuth(Authorization):
             raise RuntimeError(f"failed to complete secret auth")
 
         req = WorkspaceV1AuthenticateRequest(
-            databricks_secret=secret_value,
+            databricks_secret=base64.b64decode(secret_value).decode("utf-8"),
         )
         handler = api.WorkspaceV1Api(api_client=self._client)
 

dasl_client/client.py

@@ -3,8 +3,8 @@ from typing import Any, Callable, Iterator, List, Optional, TypeVar
 
 from dasl_api import (
     CoreV1Api,
-    CoreV1DataSourceAutoloaderSpec,
     DbuiV1Api,
+    ContentV1Api,
     WorkspaceV1Api,
     WorkspaceV1CreateWorkspaceRequest,
     api,
@@ -13,6 +13,7 @@ from pydantic import Field
 
 from dasl_client.auth.auth import (
     Authorization,
+    DatabricksSecretAuth,
     DatabricksTokenAuth,
     ServiceAccountKeyAuth,
 )
@@ -29,6 +30,8 @@ from .types import (
     WorkspaceConfig,
     TransformRequest,
     TransformResponse,
+    DataSourcePresetsList,
+    DataSourcePreset,
 )
 
 
@@ -51,7 +54,7 @@ class Client:
 
         :param auth: Authorization instance for authorizing requests to
             the dasl control plane.
-        :returns Client
+        :returns: Client
         """
         self.auth = auth
 
@@ -83,7 +86,7 @@ class Client:
         :param dasl_host: The URL of the DASL server. This value should
             not generally be specified unless you are testing against
             an alternative environment.
-        :returns Client
+        :returns: Client for the newly created workspace.
         """
         with error_handler():
             if workspace_url is None:
@@ -131,7 +134,7 @@ class Client:
         :param dasl_host: The URL of the DASL server. This value should
             not generally be specified unless you are testing against
             an alternative environment.
-        :returns Client
+        :returns: Client for the existing workspace.
         """
         with error_handler():
             if workspace_url is None:
@@ -195,7 +198,7 @@ class Client:
         :param dasl_host: The URL of the DASL server. This value should
             not generally be specified unless you are testing against
             an alternative environment.
-        :returns Client
+        :returns: Client for the newly created or existing workspace.
         """
         try:
             return Client.new_workspace(
@@ -229,6 +232,9 @@ class Client:
     def _dbui_client(self) -> DbuiV1Api:
         return DbuiV1Api(self.auth.client())
 
+    def _content_client(self) -> ContentV1Api:
+        return ContentV1Api(self.auth.client())
+
     def _workspace(self) -> str:
         return self.auth.workspace()
 
@@ -274,7 +280,7 @@ class Client:
         you will need to repopulate the service_principal_secret correctly
         before passing the result back to put_admin_config.
 
-        :returns AdminConfig
+        :returns: AdminConfig containing the current settings.
         """
         with error_handler():
             return AdminConfig.from_api_obj(
@@ -306,6 +312,8 @@ class Client:
         Retrieve the WorkspaceConfig from the DASL server. The returned
         value can be updated directly and passed to put_config in order
         to make changes.
+
+        :returns: WorkspaceConfig containing the current configuration.
         """
         with error_handler():
             return WorkspaceConfig.from_api_obj(
@@ -348,7 +356,7 @@ class Client:
         in order to make changes.
 
         :param name: The unique name of the DataSource within this workspace
-        :returns DataSource
+        :returns: DataSource
         """
         with error_handler():
             return DataSource.from_api_obj(
@@ -459,7 +467,7 @@ class Client:
         in order to make changes.
 
         :param name: The unique name of the Rule within this workspace
-        :returns Rule
+        :returns: Rule
         """
         with error_handler():
             return Rule.from_api_obj(
@@ -601,7 +609,7 @@ class Client:
         :param value: The observable value
         :param cursor: A cursor to be used when paginating results
         :param limit: A limit of the number of results to return
-        :returns EventsList
+        :returns: EventsList
         """
         with error_handler():
             return Dbui.ObservableEvents.EventsList.from_api_obj(
@@ -614,3 +622,31 @@ class Client:
                     limit=limit,
                 )
             )
+
+    def list_presets(self) -> DataSourcePresetsList:
+        """
+        List the Presets in this workspace. This will include any user defined
+        presets if a custom presets path has been configured in the workspace.
+
+        :returns: DataSourcePresetsList
+        """
+        with error_handler():
+            return DataSourcePresetsList.from_api_obj(
+                self._content_client().content_v1_get_preset_data_sources(
+                    self._workspace(),
+                )
+            )
+
+    def get_preset(self, name: str) -> DataSourcePreset:
+        """
+        Get the preset with the argument name from the DASL server. If the preset name
+        begins with 'internal_' it will instead be collected from the user catalog,
+        provided a preset path is set in the workspace config.
+
+        :param name: The unique name of the DataSource preset within this workspace.
+        :returns: DataSourcePreset
+        """
+        with error_handler():
+            return DataSourcePreset.from_api_obj(
+                self._content_client().content_v1_get_preset_datasource(self._workspace(), name)
+            )

dasl_client/preset_development/__init__.py

@@ -0,0 +1,4 @@
+# dasl_client/preset_development/__init__.py
+from .preview_parameters import *
+from .errors import *
+from .preview_engine import *

dasl_client/preset_development/errors.py

@@ -0,0 +1,159 @@
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import *
+from pyspark.sql.dataframe import DataFrame
+from typing import Dict, Any, List, Mapping, Tuple
+from IPython import get_ipython
+
+
+class PresetError(Exception):
+    pass
+
+
+class InvalidGoldTableSchemaError(PresetError):
+    def __init__(self, schema: str, additional_message: str = ""):
+        self.schema = schema
+        message = (
+            f"Malformed gold schema provided {schema}. {additional_message}".strip()
+        )
+        super().__init__(message)
+
+
+class NoSilverStageProvdedError(PresetError):
+    def __init__(self, additional_msg: str = ""):
+        message = f"No silver stage provided{additional_msg}."
+        super().__init__(message)
+
+
+class NoSilverTransformStageProvdedError(PresetError):
+    def __init__(
+        self,
+        message: str = "No silver transform stage provided, but gold stage is present.",
+    ):
+        super().__init__(message)
+
+
+class PreTransformNotFound(PresetError):
+    def __init__(
+        self,
+        message: str = "Requested silver pretransform name not found in preset's silver pretransforms.",
+    ):
+        super().__init__(message)
+
+
+class NoSilverPreTransformStageProvdedError(PresetError):
+    def __init__(
+        self,
+        message: str = "No silver transform stage provided, but prestransform name provided.",
+    ):
+        super().__init__(message)
+
+
+class MissingTableFieldError(PresetError):
+    def __init__(self, layer: str, table_name: str, field_name: str):
+        self.layer = layer
+        self.table_name = table_name
+        self.field_name = field_name
+        message = f"{layer} stage {table_name} is missing {field_name} field."
+        super().__init__(message)
+
+
+class DuplicateFieldNameError(PresetError):
+    def __init__(self, stage: str, stage_name: str, field_name: str):
+        self.stage = stage
+        self.stage_name = stage_name
+        self.field_name = field_name
+        message = f"Duplicate field specification name found in {stage} stage {stage_name} named {field_name}."
+        super().__init__(message)
+
+
+class MalformedFieldError(PresetError):
+    def __init__(self, stage: str, stage_name: str, field_name: str):
+        self.stage = stage
+        self.stage_name = stage_name
+        self.field_name = field_name
+        message = f"Please provide 1 operation only in {stage} stage {stage_name}'s field specification named {field_name}."
+        super().__init__(message)
+
+
+class MissingFieldNameError(PresetError):
+    def __init__(self, stage: str, stage_name: str):
+        self.stage = stage
+        self.stage_name = stage_name
+        message = (
+            f"Field specification in {stage} stage {stage_name} missing name field."
+        )
+        super().__init__(message)
+
+
+class MissingSilverKeysError(PresetError):
+    def __init__(self, missing_keys: str):
+        self.missing_keys = missing_keys
+        message = f"Gold table/s have no corresponding input from silver table/s: {missing_keys}"
+        super().__init__(message)
+
+
+class MissingAutoloaderConfigError(PresetError):
+    def __init__(
+        self,
+        message: str = "Autoloader mode selected, but no autoloader configuration found in preset.autoloader.",
+    ):
+        super().__init__(message)
+
+
+class AutoloaderMissingFieldError(PresetError):
+    def __init__(self, field_name: str):
+        self.field_name = field_name
+        message = f"Autoloader mode selected, but missing field {field_name} in preset."
+        super().__init__(message)
+
+
+class UnknownGoldTableError(PresetError):
+    def __init__(self, table_name: str, schema: str):
+        self.table_name = table_name
+        self.schema = schema
+        message = (
+            f"The referenced Gold table name {table_name} does not exist in {schema}."
+        )
+        super().__init__(message)
+
+
+class GoldTableCompatibilityError(PresetError):
+    def __init__(self, message: str):
+        super().__init__(message)
+
+
+class ReferencedColumnMissingError(PresetError):
+    def __init__(self, operation: str, column_name: str):
+        self.operation = operation
+        self.column_name = column_name
+        message = f"The referenced column {column_name} was not found in the dataframe during {operation} operation."
+        super().__init__(message)
+
+
+class MissingJoinFieldError(PresetError):
+    def __init__(self, field_name: str):
+        self.field_name = field_name
+        message = f"Join operation is missing required field {field_name}."
+        super().__init__(message)
+
+
+class MissingUtilityConfigurationFieldError(PresetError):
+    def __init__(self, operation: str, field_name: str):
+        self.operation = operation
+        self.field_name = field_name
+        message = f"The required configuration field {field_name} was not suppled in the {operation} operation."
+        super().__init__(message)
+
+
+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"]
+
+        self.expr = expr
+        self.assertion_message = assertion_message
+        self.df = df
+        message = f"The above rows failed the assertion expression {expr} with reason: {assertion_message}\n"
+        display(df)
+        super().__init__(message)

dasl_client/preset_development/preview_engine.py

@@ -0,0 +1,344 @@
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.types import *
+from pyspark.sql.dataframe import DataFrame
+from pyspark.sql.functions import lit, col as col_, sum as sum_, when
+from dasl_client.preset_development.preview_parameters import *
+from dasl_client.preset_development.stage import *
+from dasl_client.preset_development.errors import *
+import yaml
+from IPython import get_ipython
+
+
+class PreviewEngine:
+    """
+    This class deserializes the in-development preset's YAML and performs a series of
+    validation steps before attempting to compile each stage's table and execute them
+    based on the provided PreviewParameters.
+
+    Upon successful execution, output is generated for each successfully executed
+    stage's table operations. Additionally, if Gold stages are computed, their outputs
+    are validated against the provided Gold stage tables to ensure compatibility on a
+    per-table-name basis with the Unity Catalog.
+
+    For example, a preset Gold stage table named "http_activity" will be checked against
+    the corresponding table in the Unity Catalog schema—also named "http_activity" to
+    confirm that inserting into the Unity Catalog most likely not cause errors.
+    """
+
+    def __init__(
+        self, spark: SparkSession, preset_yaml_str: str, ds_params: PreviewParameters
+    ):
+        """
+        Creates the PreviewEngine using the given preset YAML and datasource parameters.
+        The YAML is deserialized here and checked to verify whether the requested
+        pretransform name, if provided, exists in the preset.
+
+        Instance Attributes:
+            ds_params (PreviewParameters): The input datasource's configuration.
+            preset (Dict[str, Any]): The deserialized preset YAML.
+            pretransform_name (str): The name of the requested pretransform. Defaults to None.
+            pre (Stage): Stores the pretransform Stage object internally.
+            silver (List[Stage]): Stores the Silver Stage objects internally.
+            gold (List[Stage]): Stores the Gold Stage objects internally.
+        """
+        self._spark = spark
+        self._ds_params = ds_params
+        self._preset = yaml.safe_load(preset_yaml_str)
+        self._pretransform_name = ds_params._pretransform_name
+
+        self._validate_gold_inputs(
+            self._preset.get("silver", None), self._preset.get("gold", None)
+        )
+        if self._pretransform_name:
+            self._validate_pretransform_name(
+                self._preset.get("silver", None), self._pretransform_name
+            )
+
+        self._pre = None
+        self._silver = []
+        self._gold = []
+        self._result_df_map = {}
+
+    def _validate_pretransform_name(
+        self, silver: Dict[str, str], pretransform_name: str
+    ) -> None:
+        """
+        Validates the given pretransform name exists in the provided preset's Silver
+        PreTransform stages.
+        """
+        if not silver:
+            raise NoSilverStageProvdedError(", but pretransform name provided")
+        if not (silver_pre_transform := silver.get("preTransform", None)):
+            raise NoSilverPreTransformStageProvdedError()
+        silver_pre_output_names = []
+        for table in silver_pre_transform:
+            if not (name := table.get("name", None)):
+                raise MissingTableFieldError(
+                    "Silver pretransform",
+                    table.get("name", "<stage missing name>"),
+                    "name",
+                )
+            silver_pre_output_names += [name]
+        if pretransform_name not in silver_pre_output_names:
+            raise PreTransformNotFound()
+
+    def _validate_gold_inputs(
+        self, silver: Dict[str, str], gold: Dict[str, str]
+    ) -> None:
+        """
+        Validate gold tables all have a silver table to input from.
+        """
+        if not gold:
+            return
+
+        if not len(gold):
+            return
+
+        if not silver:
+            raise NoSilverStageProvdedError(", but gold stage is present")
+
+        gold_input_names = []
+        for table in gold:
+            if not (input := table.get("input", None)):
+                raise MissingTableFieldError(
+                    "Gold", table.get("name", "<stage missing name>"), "input"
+                )
+            gold_input_names += [input]
+
+        if not (silver_transform := silver.get("transform", None)):
+            raise NoSilverTransformStageProvdedError()
+        silver_output_names = []
+        for table in silver_transform:
+            if not (name := table.get("name", None)):
+                raise MissingTableFieldError(
+                    "Silver transform", table.get("name", ""), "name"
+                )
+            silver_output_names += [name]
+
+        missing_keys = set(gold_input_names) - set(silver_output_names)
+        if missing_keys:
+            raise MissingSilverKeysError(missing_keys)
+
+    def _compile_stages(self) -> None:
+        """
+        Creates Stage objects, setting pretransform to None if not provided.
+        """
+        pretransform = None
+        if self._pretransform_name:
+            for table in self._preset["silver"]["preTransform"]:
+                if table["name"] == self._pretransform_name:
+                    self._pre = Stage(self._spark, "silver pretransform", table)
+                    break
+
+        self._silver = [
+            Stage(self._spark, "silver transform", table)
+            for table in self._preset.get("silver", {}).get("transform", [])
+        ]
+        self._gold = [
+            Stage(self._spark, "gold", table) for table in self._preset.get("gold", [])
+        ]
+
+    def _run(
+        self, df: DataFrame
+    ) -> Tuple[DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]]:
+        """
+        Runs all stages, in medallion stage order. This allows prior stage outputs to feed
+        into later stage inputs.
+
+        Returns:
+            Dataframes containing the output from each run Stage.
+        """
+        if self._pre:
+            df = self._pre.run(df)
+
+        silver_output_map = {}
+        for table in self._silver:
+            silver_output_map[table._name] = table.run(df)
+
+        gold_output_map = {}
+        for table in self._gold:
+            # We store as gold_name/silver_input to prevent clobbering on duplicate gold table use.
+            gold_output_map[f"{table._name}/{table._input}"] = table.run(
+                silver_output_map[table._input]
+            )
+
+        return (
+            (df, silver_output_map, gold_output_map)
+            if self._pre
+            else (None, silver_output_map, gold_output_map)
+        )
+
+    def _render_output(
+        self,
+        input_df: DataFrame,
+        stage_dataframes: Tuple[DataFrame, Dict[str, DataFrame], Dict[str, DataFrame]],
+        gold_table_schema: str,
+    ) -> None:
+        """
+        Displays formatted HTML output from executed Stages' DataFrames.
+        """
+        # TODO: Investigate further into using Databricks's style sheets here.
+
+        # Get the Databricks built-in functions out the namespace.
+        ipython = get_ipython()
+        displayHTML = ipython.user_ns["displayHTML"]
+        display = ipython.user_ns["display"]
+
+        def d(txt, lvl) -> None:
+            displayHTML(
+                f"""
+                <div style="background-color: 
+                background-color: rgb(18, 23, 26); padding: 0; margin: 0;">
+                    <h{lvl} style="margin: 0; background-color: rgb(244, 234, 229);">{txt}</h{lvl}>
+                </div>
+                """
+            )
+
+        def check_struct_compatibility(
+            target_field: StructField, df_field: StructField, prefix=""
+        ):
+            if not (
+                isinstance(target_field.dataType, StructType)
+                and isinstance(df_field.dataType, StructType)
+            ):
+                return
+
+            target_fields = {
+                field.name: field for field in target_field.dataType.fields
+            }
+            for field in df_field.dataType.fields:
+                if field.name not in target_fields:
+                    raise GoldTableCompatibilityError(
+                        f"Extra field found in gold stage output STRUCT column {prefix}{target_field.name}: {field.name}"
+                    )
+                else:
+                    if isinstance(field.dataType, StructType):
+                        check_struct_compatibility(
+                            target_fields[field.name],
+                            field,
+                            prefix=prefix + target_field.name + ".",
+                        )
+
+        (pre_df, silver, gold) = stage_dataframes
+        d("Input", 1)
+        display(input_df)
+        d("Silver Pre-Transform", 1)
+        if pre_df:
+            display(pre_df)
+        else:
+            d("Skipped", 2)
+        d("Silver Transform", 1)
+        for name, df in silver.items():
+            d(f"{name}", 2)
+            display(df)
+        d("Gold", 1)
+        for name, df in gold.items():
+            d(f"{name}", 2)
+            d("Stage output", 3)
+            display(df)
+
+            # NOTE: Name is stored as Gold_name/Silver_input. So we need to get just the Gold table
+            # name that we are comparing the dataframe metadata to.
+            name = name.split("/")[0]
+
+            if not self._spark.catalog.tableExists(f"{gold_table_schema}.{name}"):
+                raise UnknownGoldTableError(name, gold_table_schema)
+
+            # Performs the type check.
+            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.
+            if not set(df.columns).issubset(delta_df.columns):
+                raise GoldTableCompatibilityError(
+                    f"Extra columns provided: {', '.join([col for col in df.columns if col not in delta_df.columns])}"
+                )
+
+            # Now we check no new fields in STRUCT columns.
+            for field in delta_df.schema.fields:
+                if isinstance(field.dataType, StructType) and field.name in df.columns:
+                    # Retrieve the corresponding field from the DataFrame's schema.
+                    df_field = next(f for f in df.schema.fields if f.name == field.name)
+                    check_struct_compatibility(field, df_field)
+
+            # Check nullable columns exist, and data what we are inserting is set.
+            non_nullable_cols = [
+                field.name for field in delta_df.schema.fields if not field.nullable
+            ]
+            null_checks = [
+                sum_(when(col_(col).isNull(), 1).otherwise(0)).alias(col)
+                for col in non_nullable_cols
+            ]
+            null_counts = df.select(null_checks).collect()[0].asDict()
+            cols_with_nulls = []
+            try:
+                cols_with_nulls = [
+                    col_name for col_name, count in null_counts.items() if count > 0
+                ]
+            except TypeError:
+                # There were no records returned and so null_counts == None.
+                pass
+            if cols_with_nulls:
+                raise GoldTableCompatibilityError(
+                    f"Record with null data found for non-nullable columns: {', '.join([col for col in cols_with_nulls])}"
+                )
+
+            d("Resultant gold table preview", 3)
+            display(unioned_df)
+
+    def evaluate(self, gold_table_schema: str) -> 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
+        the Unity Catalog Gold tables.
+        """
+        s = gold_table_schema.split(".")
+        if len(s) != 2:
+            raise InvalidGoldTableSchemaError(gold_table_schema)
+        catalog_name = s[0]
+        schema_name = s[1]
+        if any(
+            row.catalog == catalog_name
+            for row in self._spark.sql("SHOW CATALOGS").collect()
+        ):
+            if not any(
+                row.databaseName == schema_name
+                for row in self._spark.sql(f"SHOW SCHEMAS IN {catalog_name}").collect()
+            ):
+                raise InvalidGoldTableSchemaError(
+                    gold_table_schema,
+                    f"Schema {schema_name} not found in catalog {catalog_name} or insufficient permissions.",
+                )
+        else:
+            raise InvalidGoldTableSchemaError(
+                gold_table_schema,
+                f"Catalog {catalog_name} not found or insufficient permissions.",
+            )
+
+        # If we are using the autoloader, fetch format from preset and others.
+        if self._ds_params._mode == "autoloader":
+            if not (autoloader_conf := self._preset.get("autoloader", None)):
+                raise MissingAutoloaderConfigError()
+            if not (file_format := autoloader_conf.get("format", None)):
+                raise AutoloaderMissingFieldError("format")
+            self._ds_params.set_autoloader_format(file_format)
+            if schemaFile := autoloader_conf.get("schemaFile", None):
+                self._ds_params.set_autoloader_schema_file(schemaFile)
+            if multiline := autoloader_conf.get("multiline", None):
+                if multiline == "true":
+                    self._ds_params.set_multiline(True)
+                else:
+                    self._ds_params.set_multiline(False)
+            if cloudFiles := autoloader_conf.get("cloudFiles", None):
+                if schema_hints := cloudFiles.get("schemaHints", None):
+                    self._ds_params.set_autoloader_cloudfiles_schema_hints(schema_hints)
+                if schema_hints_file := cloudFiles.get("schemaHintsFile", None):
+                    self._ds_params.set_autoloader_cloudfiles_schema_hint_file(
+                        schema_hints_file
+                    )
+
+        self._compile_stages()
+
+        with self._ds_params as df:
+            self._result_df_map = self._run(df)
+            self._render_output(df, self._result_df_map, gold_table_schema)