dasl-client 1.0.22__py3-none-any.whl → 1.0.24__py3-none-any.whl

dasl_client/preset_development/preview_parameters.py

@@ -9,11 +9,6 @@ import uuid
 from IPython import get_ipython
 
 
-@udf(StringType())
-def constant_udf(*args):
-    return "<sortable_random_id>"
-
-
 class PreviewParameters:
     """
     This class provides three methods for supplying input records to the preset development environment.
@@ -60,6 +55,78 @@ class PreviewParameters:
         .set_table("system.access.audit")
     ```
 
+    **4. SilverBronze Mode:**
+    "silverbronze" mode, works like a more advanced "table" mode. It allows for joining of multiple
+    tables as input. This mode requires setting bronze table definitions. This mode behaves in 2
+    seperate ways depending on whether an autoloader location is set or not. If an autoloader location
+    is set the first entry in the bronze table definitions is used to name and alias the autoloader's
+    input and these can be used in later join expressions. Used in this way, the autoloader will be
+    loaded as in "autoloader" mode, and run through preBronze stages before being joined with the
+    remainder of the bronze table definitions. This mimics not skipping bronze in a DataSource and
+    joining what was read in silver. If an autoloader location is not set, the behaviour instead
+    attempts to emulate a DataSource set to skip the bronze stage. That is, all preBronze and bronze
+    stages will be skipped, and the name of the first entry in the given bronze table definitions will
+    be read from instead. Any subsequent bronze table definitions will be joined against this table.
+
+    Using no autoloader location (this will read from the first table):
+    ```python
+    bronze_tables = [
+    {
+        "name": "databricks_dev.default.sev_map",
+        "alias": "tab1"
+    },
+    {
+        "name": "databricks_dev.alan_bronze.akamai_waf",
+        "alias": "tab2",
+        "joinExpr": "id::string = tab2.serviceID",
+        "joinType": "left"
+    },
+    {
+        "name": "databricks_dev.alan_silver.cloudflare_hjttp_request",
+        "alias": "tab3",
+        "joinExpr": "tab1.id::string = tab3.ClientRequestsBytes",
+        "joinType": "inner"
+    }
+    ]
+
+    ds_params = (
+        PreviewParameters(spark)
+        .from_silverbronze_tables()
+        .set_bronze_table_definitions(bronze_tables)
+    )
+
+    ps = PreviewEngine(spark, yaml_string, ds_params)
+    ps.evaluate("stage.gold")
+    ```
+
+    Using an autoloader location (this will read from the autoloader and name the df tab1):
+    ```python
+    bronze_tables = [
+    {
+        "name": "tab1"
+    },
+    {
+        "name": "databricks_dev.alan_bronze.akamai_waf",
+        "alias": "tab2",
+        "joinExpr": "id::string = tab2.serviceID",
+        "joinType": "left"
+    },
+    {
+        "name": "databricks_dev.alan_silver.cloudflare_hjttp_request",
+        "alias": "tab3",
+        "joinExpr": "tab1.id::string = tab3.ClientRequestsBytes",
+        "joinType": "inner"
+    }
+    ]
+
+    ds_params = (
+        PreviewParameters(spark)
+        .from_silverbronze_tables()
+        .set_bronze_table_definitions(bronze_tables)
+        .set_autoloader_location("s3://antimatter-dasl-testing/csamples3/mars/area1/")
+    )
+    ```
+
     **Note:**
     When using autoloader mode, this implementation requires a location to store a temporary schema for
     the loaded records. By default, this is set to `"dbfs:/tmp/schemas"`. You can change this using
@@ -94,9 +161,10 @@ class PreviewParameters:
             df (DataFrame): Internal Spark DataFrame loaded using the specified parameters.
         """
         self._spark = spark
-        self._mode = None  # [input, autoloader]
+        self._mode = None  # [input, table, autoloader, silverbronze]
         self._record_limit = 10
         self._autoloader_temp_schema_location = "dbfs:/tmp/schemas"
+        self._gold_test_schemas = []
 
         self._time_column = None
         self._start_time = None
@@ -109,18 +177,97 @@ class PreviewParameters:
         self._cloudfiles_schema_hints = None
         self._cloudfiles_reader_case_sensitive = "true"
         self._cloudfiles_multiline = "true"
+        self._cloudfiles_wholetext = "false"
         self._schema_uuid_str = str(uuid.uuid4())
+        self._single_variant_column = None
 
         self._schema = None
         self._data = None
 
         self._table = None
 
+        self._bronze_tables = None
+
         self._pretransform_name = None
-        self._bronze_pre_transform: Optional[List[str]] = None
 
         self._df = None
 
+    def __create_from_autoloader(self) -> DataFrame:
+        stream_df = (
+            self._spark.readStream.format("cloudFiles")
+            .option("cloudFiles.format", self._autoloader_format)
+            .option("readerCaseSensitive", self._cloudfiles_reader_case_sensitive)
+        )
+
+        # text and wholetext needs to be handled seperately.
+        stream_df = (
+            stream_df.option("multiline", self._cloudfiles_multiline)
+            if self._autoloader_format != "text"
+            else stream_df.option("wholetext", self._cloudfiles_wholetext)
+        )
+
+        if self._single_variant_column:
+            stream_df = stream_df.option(
+                "singleVariantColumn", self._single_variant_column
+            )
+
+        if self._schema_file:
+            with open(self._schema_file, "r") as f:
+                stream_df = stream_df.schema(f.read().strip())
+        else:
+            stream_df = (
+                stream_df.option("inferSchema", "true")
+                .option("cloudFiles.inferColumnTypes", "true")
+                .option(
+                    "cloudFiles.schemaLocation",
+                    f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
+                )
+            )
+
+        if self._cloudfiles_schema_hints:
+            stream_df = stream_df.option(
+                "cloudFiles.schemaHints", self._cloudfiles_schema_hints
+            )
+        elif self._clouldfiles_schema_hints_file:
+            stream_df = stream_df.option(
+                "cloudFiles.schemaHintsFile", self._clouldfiles_schema_hints_file
+            )
+
+        stream_df = stream_df.load(self._autoloader_location).limit(self._record_limit)
+
+        query = (
+            stream_df.writeStream.format("memory")
+            .queryName("batch_data")
+            .trigger(availableNow=True)
+            .start()
+        )
+
+        query.awaitTermination()
+
+    def __create_from_silverbronze_tables_join(self) -> DataFrame:
+        if not self._bronze_tables or not len(self._bronze_tables):
+            raise MissingBronzeTablesError()
+
+        # Validate name and joinExpr are set.
+        for i in range(len(self._bronze_tables)):
+            if not self._bronze_tables[i].get("name", None):
+                raise MissingBronzeTableFieldError("name")
+            if i > 0 and not self._bronze_tables[i].get("joinExpr", None):
+                raise MissingBronzeTableFieldError("joinExpr")
+
+        # If there is an autoloader location given, we create the df now and
+        # then allow preBronze stage to run. Otherwise we skip preBronze stages
+        # and as part of running the silverbronze joins we create the df from
+        # the first entry in the bronze tables list.
+        df = None
+        if self._autoloader_location:
+            self.__create_from_autoloader()
+            df = self._spark.table("batch_data").alias(
+                self._bronze_tables[0].get("name", "")
+            )  # Use first's name.
+
+        return df
+
     def __enter__(self):
         """
         Creates a DataFrame with data using the method specified. In the case of "autoloader",
@@ -136,59 +283,10 @@ class PreviewParameters:
         elif self._mode == "table":
             self._df = self._spark.table(self._table).limit(self._record_limit)
         elif self._mode == "autoloader":
-            stream_df = (
-                self._spark.readStream.format("cloudFiles")
-                .option("cloudFiles.format", self._autoloader_format)
-                .option("multiline", self._cloudfiles_multiline)
-                .option("readerCaseSensitive", self._cloudfiles_reader_case_sensitive)
-            )
-
-            if self._schema_file:
-                with open(self._schema_file, "r") as f:
-                    stream_df = stream_df.schema(f.read().strip())
-            else:
-                stream_df = (
-                    stream_df.option("inferSchema", "true")
-                    .option("cloudFiles.inferColumnTypes", "true")
-                    .option(
-                        "cloudFiles.schemaLocation",
-                        f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
-                    )
-                )
-
-            if self._cloudfiles_schema_hints:
-                stream_df = stream_df.option(
-                    "cloudFiles.schemaHints", self._cloudfiles_schema_hints
-                )
-            elif self._clouldfiles_schema_hints_file:
-                stream_df = stream_df.option(
-                    "cloudFiles.schemaHintsFile", self._clouldfiles_schema_hints_file
-                )
-
-            stream_df = stream_df.load(self._autoloader_location).limit(
-                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(availableNow=True)
-                .start()
-            )
-
-            query.awaitTermination()
-
+            self.__create_from_autoloader()
             self._df = self._spark.table("batch_data")
-
-        if self._time_column:
-            self._df = self._df.filter(
-                f"timestamp({self._time_column}) >= timestamp('{self._start_time}') AND timestamp({self._time_column}) < timestamp('{self._end_time}')"
-            )
-
-        self._df = self._df.withColumn("dasl_id", constant_udf())
+        elif self._mode == "silverbronze":
+            self._df = self.__create_from_silverbronze_tables_join()
 
         return self._df
 
@@ -206,10 +304,21 @@ class PreviewParameters:
                 f"{self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
                 recurse=True,
             )
+            for gold_test_schema in self._gold_test_schemas:
+                dbutils.fs.rm(
+                    f"{self._autoloader_temp_schema_location}/{gold_test_schema}",
+                    recurse=True,
+                )
         else:
-            print(
-                f"FYI, we are leaking temp data {self._autoloader_temp_schema_location}/{self._schema_uuid_str}"
-            )
+            leaked_lines = [
+                f"FYI, we are leaking temp data {self._autoloader_temp_schema_location}/{self._schema_uuid_str}",
+                *[
+                    f"{self._autoloader_temp_schema_location}/{x}"
+                    for x in self._gold_test_schemas
+                ],
+            ]
+            print(", ".join(leaked_lines))
+        self._gold_test_schemas = []
 
     def from_input(self):
         """
@@ -242,6 +351,36 @@ class PreviewParameters:
         self._mode = "table"
         return self
 
+    def from_silverbronze_tables(self):
+        """
+        Set the data source loader to "bronze tables" mode. Requires a list of bronze table
+        definitions to be provided.
+
+        Returns:
+            PreviewParameters: The current instance with updated configuration.
+        """
+        self._mode = "silverbronze"
+        return self
+
+    def set_bronze_table_definitions(self, definitions: List[Dict[str, str]]):
+        """
+        Set the bronze table definitions for bronze tables mode. `name` and `joinExpr` are
+        required. If `alias` is not provided, one can use the `name` to refer to the table.
+        If `joinType` is not provided, "left" is used as a default value. If pr
+
+        [
+            {
+                "name": "name",
+                "alias": "alias1",
+                "joinType": "inner",
+                "joinExpr": "base_table.col1 = alias1.col1
+            },
+            ...
+        ]
+        """
+        self._bronze_tables = definitions
+        return self
+
     def set_autoloader_temp_schema_location(self, path: str):
         """
         Set the location for the autoloader's streaming mode schema to be created. This is
@@ -253,6 +392,15 @@ class PreviewParameters:
         self._autoloader_temp_schema_location = path
         return self
 
+    def get_autoloader_temp_schema_location(self) -> str:
+        """
+        Get the location for the autoloader's streaming mode schema to be created.
+
+        Returns:
+             str: The location for the autoloader's streaming mode schema to be created.
+        """
+        return self._autoloader_temp_schema_location
+
     def set_data_schema(self, schema: StructType):
         """
         Set the input schema for "input" mode. For example:
@@ -290,7 +438,7 @@ class PreviewParameters:
         self._autoloader_location = location
         return self
 
-    def set_autoloader_format(self, file_format: str):
+    def _set_autoloader_format(self, file_format: str):
         """
         Used internally to set the autoloader format.
 
@@ -299,12 +447,16 @@ class PreviewParameters:
         """
         if file_format.lower() == "jsonl":
             self._autoloader_format = "json"
-            self.set_autoloader_multiline(False)
+            self._cloudfiles_multiline = "false"
+            return self
+        if file_format.lower() == "wholetext":
+            self._autoloader_format = "text"
+            self._cloudfiles_wholetext = "true"
             return self
         self._autoloader_format = file_format
         return self
 
-    def set_autoloader_schema_file(self, path: str):
+    def _set_autoloader_schema_file(self, path: str):
         """
         Set the schema file path for "autoloader" mode.
 
@@ -314,7 +466,7 @@ class PreviewParameters:
         self._schema_file = path
         return self
 
-    def set_autoloader_cloudfiles_schema_hint_file(self, path: str):
+    def _set_autoloader_cloudfiles_schema_hint_file(self, path: str):
         """
         Set the cloudFiles schema hints file path for "autoloader" mode.
 
@@ -324,7 +476,7 @@ class PreviewParameters:
         self._clouldfiles_schema_hints_file = path
         return self
 
-    def set_autoloader_cloudfiles_schema_hints(self, cloudfiles_schema_hints: str):
+    def _set_autoloader_cloudfiles_schema_hints(self, cloudfiles_schema_hints: str):
         """
         Set the cloudFiles schema hints string for "autoloader" mode.
 
@@ -334,26 +486,6 @@ class PreviewParameters:
         self._cloudfiles_schema_hints = cloudfiles_schema_hints
         return self
 
-    def set_autoloader_reader_case_sensitive(self, b: bool):
-        """
-        Set the cloudFiles reader case-sensitive boolean for "autoloader" mode.
-
-        Returns:
-            PreviewParameters: The current instance with updated configuration.
-        """
-        self._cloudfiles_reader_case_sensitive = "true" if b else "false"
-        return self
-
-    def set_autoloader_multiline(self, b: bool):
-        """
-        Set the cloudFiles multiline boolean for "autoloader" mode.
-
-        Returns:
-            PreviewParameters: The current instance with updated configuration.
-        """
-        self._cloudfiles_multiline = "true" if b else "false"
-        return self
-
     def set_pretransform_name(self, pretransform_name: str):
         """
         Set the pretransform name to use, if desired. If not set, Silver PreTransform
@@ -365,16 +497,6 @@ 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
@@ -409,3 +531,21 @@ class PreviewParameters:
         """
         self._table = table_name
         return self
+
+    def _set_load_as_single_variant(self, col_name: Optional[str] = None):
+        """
+        Enable loadAsSingleVariant mode. This will ingest data into a single VARIANT-typed column.
+        The default name of that column is `data`.
+
+        Returns:
+            PreviewParameters: The current instance with updated configuration.
+        """
+        self._single_variant_column = col_name if col_name is not None else "data"
+        return self
+
+    def add_gold_schema_table(self, gold_schema_table_name: str):
+        """
+        Add a gold schema temporary table name that will need to be cleaned
+        up at the end of the run.
+        """
+        self._gold_test_schemas.append(gold_schema_table_name)

dasl_client/preset_development/stage.py

@@ -4,6 +4,8 @@ from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.functions import col, lit
 from dasl_client.preset_development.errors import *
 
+import re
+
 FieldSpec = Dict[str, Any]
 
 
@@ -45,6 +47,57 @@ class Stage:
 
     __op_list = ["assert", "literal", "from", "alias", "expr", "join"]
 
+    # From fieldSpec valid characters.
+    __invalid_char_pattern = re.compile(r"[\s,;{}\(\)\n\t=]")
+
+    def __validate_field_spec(self, fields: List[Dict[str, str]], stage: str):
+        names = []
+        for field in self._fields:
+            # Check for name. If no name check for assert.
+            if not (name := field.get("name", None)):
+                if not field.get("assert", None):  # Can't walrus em all :/
+                    raise MissingFieldNameError(self._stage, self._name)
+
+            # Check this new name does not duplicate an existing.
+            if name in names:
+                raise DuplicateFieldNameError(self._stage, self._name, name)
+            names += [name] if name != None else []
+
+            # Check for only 1 defined operation.
+            missing_op_count = [
+                spec for spec in [field.get(op, None) for op in self.__op_list]
+            ].count(None)
+            if (missing_op_count == len(self.__op_list)) or (
+                len(self.__op_list) - missing_op_count > 1
+            ):
+                raise MalformedFieldError(
+                    self._stage, self._name, field.get("name", None)
+                )
+
+            # Literal must be a string.
+            if lit := field.get("literal", None):
+                if type(lit) != str:
+                    raise InvalidLiteralError(
+                        self._stage, self._name, field.get("name", None)
+                    )
+
+            # Validate from (makes sure its not an expression, etc.). This mirrors Scala code's validation.
+            if frm := field.get("from", None):
+                if len(frm) >= 256:
+                    raise InvalidFromError(
+                        self._stage,
+                        self._name,
+                        field.get("name", None),
+                        "Column name too long",
+                    )
+                if frm.strip() == "" or self.__invalid_char_pattern.search(frm):
+                    raise InvalidFromError(
+                        self._stage,
+                        self._name,
+                        field.get("name", None),
+                        "Malformed column name referenced",
+                    )
+
     def __init__(self, spark: SparkSession, stage: str, table: Dict[str, any]):
         """
         Initializes a Stage object that encapsulates all operations required for a single
@@ -68,9 +121,10 @@ class Stage:
         self._utils = table.get("utils", {})
         self._input = table.get("input", None)
 
+        # The dasl_id does not exist before bronze or when dealing with temp fields.
         fields = (
             [{"name": "dasl_id", "from": "dasl_id"}] + table.get("fields", [])
-            if self._stage != "temp_fields"
+            if self._stage not in ["temp_fields", "bronze pretransform"]
             else table.get("fields", [])
         )
         self._fields = [
@@ -78,24 +132,7 @@ class Stage:
         ]
         self._assertions = [f for f in fields if f.get("assert", None)]
 
-        names = []
-        for field in self._fields:
-            if not (name := field.get("name", None)):
-                if not field.get("assert", None):  # Can't walrus em all :/
-                    raise MissingFieldNameError(self._stage, self._name)
-            if name in names:
-                raise DuplicateFieldNameError(self._stage, self._name, name)
-            names += [name]
-
-            missing_op_count = [
-                spec for spec in [field.get(op, None) for op in self.__op_list]
-            ].count(None)
-            if (missing_op_count == len(self.__op_list)) or (
-                len(self.__op_list) - missing_op_count > 1
-            ):
-                raise MalformedFieldError(
-                    self._stage, self._name, field.get("name", None)
-                )
+        self.__validate_field_spec(self._fields, self._stage)
 
     def _referenced_columns(self) -> List[str]:
         """
@@ -311,9 +348,15 @@ class Stage:
             # check that the from column exists in the df?
             return f"{self.auto_backtick(field['from'])} AS {self.auto_backtick(name)}"
         elif field.get("literal", None):
-            return f"'{field['literal']}' AS {self.auto_backtick(name)}"
+            return f"{repr(field['literal'])} AS {self.auto_backtick(name)}"
         elif field.get("expr", None) is not None:
-            return f"{field['expr']} AS {self.auto_backtick(name)}"
+            expr = field["expr"].replace("\\", "\\\\")
+            # If we are in a bronze pretransform, we do not want the fieldSpec.name.
+            return (
+                f"{expr} AS {self.auto_backtick(name)}"
+                if self._stage != "bronze pretransform"
+                else expr
+            )
         else:
             return ""
 
@@ -426,6 +469,8 @@ class Stage:
         select_fields = self.render_fields(self._fields)
 
         if preserve := self._utils.get("unreferencedColumns", None):
+            if self._stage == "gold":  # No utils run in gold.
+                raise DisallowedUtilityConfigurationError("unreferencedColumns", "gold")
             should_preserve = preserve.get("preserve", None)
             if type(should_preserve) != bool:
                 raise MissingUtilityConfigurationFieldError(
@@ -438,9 +483,23 @@ class Stage:
                     # 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(list(map(lambda x: self.force_apply_backticks(x), preserved_columns)))}) AS {self.auto_backtick(embed_col)}"
-                    ]
+                    colType = preserve.get("embedColumnType", "struct")
+                    if colType == "struct":
+                        select_fields += [
+                            f"struct({', '.join(list(map(lambda x: self.force_apply_backticks(x), preserved_columns)))}) AS {self.auto_backtick(embed_col)}"
+                        ]
+                    elif colType == "json":
+                        select_fields += [
+                            f"to_json(struct({', '.join(list(map(lambda x: self.force_apply_backticks(x), preserved_columns)))})) AS {self.auto_backtick(embed_col)}"
+                        ]
+                    elif colType == "variant":
+                        select_fields += [
+                            f"parse_json(to_json(struct({', '.join(list(map(lambda x: self.force_apply_backticks(x), preserved_columns)))}))) AS {self.auto_backtick(embed_col)}"
+                        ]
+                    else:
+                        raise UnknownUtilityConfigurationFieldError(
+                            "embedColumnType", "unreferencedColumns"
+                        )
                 else:
                     (
                         preserved_columns,
@@ -476,6 +535,8 @@ class Stage:
             A DataFrame with the resultant operation's records.
         """
         if json_extracts := self._utils.get("jsonExtract", None):
+            if self._stage == "gold":  # No utils run in gold.
+                raise DisallowedUtilityConfigurationError("jsonExtract", "gold")
             for json_extract in json_extracts:
                 source = json_extract.get("source")
                 if not source:
@@ -610,6 +671,8 @@ class Stage:
             A DataFrame with the resultant operation's records.
         """
         if temp_fields := self._utils.get("temporaryFields", None):
+            if self._stage == "gold":  # No utils run in gold.
+                raise DisallowedUtilityConfigurationError("temporaryFields", "gold")
             df = Stage(self._spark, "temp_fields", {"fields": temp_fields}).run(df)
         return df
 

dasl_client/regions.json

@@ -0,0 +1,3 @@
+{
+	"us-east-1": "https://api.prod.sl.antimatter.io"
+}

dasl_client/regions.py

@@ -0,0 +1,18 @@
+import json
+from importlib import resources
+from typing import List
+
+_data = json.loads(resources.files(__package__).joinpath("regions.json").read_text())
+
+
+class Regions:
+    @staticmethod
+    def lookup(name: str) -> str:
+        try:
+            return _data(name)
+        except KeyError as e:
+            raise ValueError(f"unknown region {name}") from e
+
+    @staticmethod
+    def list() -> List[str]:
+        return list(_data.keys())