Merge pull request #30 from databrickslabs/feature/dlt-meta-uc-cli

ravi-databricks · web-flow · commit 2214076a1477 · 2023-12-20T11:10:42.000-08:00
Added unit tests and readme
diff --git a/README.md b/README.md
@@ -69,11 +69,14 @@ With this framework you need to record the source and target metadata in an onbo
 ## Getting Started
 Refer to the [Getting Started](https://databrickslabs.github.io/dlt-meta/getting_started)
 ### Databricks Labs DLT-META CLI lets you run onboard and deploy in interactive python terminal
+#### pre-requisites:
+- [Databricks CLI](https://docs.databricks.com/en/dev-tools/cli/tutorial.html)
+- Python 3.8.0 +
+#### Steps:
 - ``` git clone dlt-meta ```
 - ``` cd dlt-meta ```
 - ``` python -m venv .venv ```
 - ```source .venv/bin/activate ```
-- ``` pip install databricks ```
 - ``` pip install databricks-sdk ```
 - ```databricks labs dlt-meta onboard``` 
 - - Above command will prompt you to provide onboarding details. If you have cloned dlt-meta git repo then accept defaults which will launch config from demo folder.
diff --git a/src/__main__.py b/src/__main__.py
@@ -35,6 +35,10 @@ def parse_args():
 def main():
     """Whl file entry point."""
     args = parse_args()
+    onboard_dataflowspecs(args)
+
+
+def onboard_dataflowspecs(args):
     onboard_layer = args.onboard_layer
     uc_enabled = True if args.uc_enabled and args.uc_enabled.lower() == "true" else False
     onboarding_args_dict = args.__dict__
diff --git a/src/dataflow_pipeline.py b/src/dataflow_pipeline.py
@@ -360,7 +360,6 @@ def invoke_dlt_pipeline(spark, layer):
             dataflowspec_list = DataflowSpecUtils.get_bronze_dataflow_spec(spark)
         elif "silver" == layer.lower():
             dataflowspec_list = DataflowSpecUtils.get_silver_dataflow_spec(spark)
-
         logger.info(f"Length of Dataflow Spec {len(dataflowspec_list)}")
         for dataflowSpec in dataflowspec_list:
             logger.info("Printing Dataflow Spec")
diff --git a/src/pipeline_readers.py b/src/pipeline_readers.py
@@ -105,7 +105,6 @@ def read_kafka(spark, bronze_dataflow_spec, schema_json) -> DataFrame:
             # add date, hour, and minute columns derived from eventhub enqueued timestamp
             .selectExpr("*", "to_date(timestamp) as date", "hour(timestamp) as hour", "minute(timestamp) as minute")
         )
-
         if schema_json:
             schema = StructType.fromJson(schema_json)
             return (
diff --git a/tests/resources/onboarding.json b/tests/resources/onboarding.json
@@ -1,154 +1 @@
-[
-   {
-      "data_flow_id": "100",
-      "data_flow_group": "A1",
-      "source_system": "MYSQL",
-      "source_format": "cloudFiles",
-      "source_details": {
-         "source_database": "APP",
-         "source_table": "CUSTOMERS",
-         "source_path_dev": "tests/resources/data/customers",
-         "source_schema_path": "tests/resources/schema/customer_schema.ddl"
-      },
-      "bronze_database_dev": "bronze",
-      "bronze_database_staging": "bronze",
-      "bronze_database_prd": "bronze",
-      "bronze_table": "customers_cdc",
-      "bronze_reader_options": {
-         "cloudFiles.format": "json",
-         "cloudFiles.inferColumnTypes": "true",
-         "cloudFiles.rescuedDataColumn": "_rescued_data"
-      },
-      "bronze_table_path_dev": "tests/resources/delta/customers",
-      "bronze_table_properties": {
-         "pipelines.autoOptimize.managed": "false",
-         "pipelines.reset.allowed": "false"
-      },
-      "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/customers/bronze_data_quality_expectations.json",
-      "silver_database_dev": "silver",
-      "silver_database_staging": "silver",
-      "silver_database_prd": "silver",
-      "silver_table": "customers",
-      "silver_cdc_apply_changes": {
-         "keys": [
-            "id"
-         ],
-         "sequence_by": "operation_date",
-         "scd_type": "1",
-         "apply_as_deletes": "operation = 'DELETE'",
-         "except_column_list": [
-            "operation",
-            "operation_date",
-            "_rescued_data"
-         ]
-      },
-      "silver_table_path_dev": "tests/resources/data/silver/customers",
-      "silver_table_properties": {
-         "pipelines.autoOptimize.managed": "false",
-         "pipelines.reset.allowed": "false",
-         "pipelines.autoOptimize.zOrderCols": "id,email"
-      },
-      "silver_transformation_json_dev": "tests/resources/silver_transformations.json"
-   },
-   {
-      "data_flow_id": "101",
-      "data_flow_group": "A1",
-      "source_system": "MYSQL",
-      "source_format": "cloudFiles",
-      "source_details": {
-         "source_database": "APP",
-         "source_table": "TRANSACTIONS",
-         "source_path_prd": "tests/resources/data/transactions",
-         "source_path_dev": "tests/resources/data/transactions"
-      },
-      "bronze_database_dev": "bronze",
-      "bronze_database_staging": "bronze",
-      "bronze_database_prd": "bronze",
-      "bronze_table": "transactions_cdc",
-      "bronze_reader_options": {
-         "cloudFiles.format": "json",
-         "cloudFiles.inferColumnTypes": "true",
-         "cloudFiles.rescuedDataColumn": "_rescued_data"
-      },
-      "bronze_table_path_dev": "tests/resources/delta/transactions",
-      "bronze_table_path_staging": "s3://db-dlt-meta-staging/demo/data/bronze/transactions",
-      "bronze_table_path_prd": "s3://db-dlt-meta-prod/demo/data/bronze/transactions",
-      "bronze_table_properties": {
-         "pipelines.reset.allowed": "false"
-      },
-      "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/transactions/bronze_data_quality_expectations.json",
-      "bronze_database_quarantine_dev": "bronze",
-      "bronze_database_quarantine_staging": "bronze",
-      "bronze_database_quarantine_prd": "bronze",
-      "bronze_quarantine_table": "transactions_cdc_quarantine",
-      "bronze_quarantine_table_path_dev": "tests/resources/data/bronze/transactions_quarantine",
-      "silver_database_dev": "silver",
-      "silver_database_preprd": "silver",
-      "silver_database_prd": "silver",
-      "silver_table": "transactions",
-      "silver_cdc_apply_changes": {
-         "keys": [
-            "id"
-         ],
-         "sequence_by": "operation_date",
-         "scd_type": "1",
-         "apply_as_deletes": "operation = 'DELETE'",
-         "except_column_list": [
-            "operation",
-            "operation_date",
-            "_rescued_data"
-         ]
-      },
-      "silver_partition_columns": "transaction_date",
-      "silver_table_path_dev": "tests/resources/data/silver/transactions",
-      "silver_transformation_json_dev": "tests/resources/silver_transformations.json",
-      "silver_table_properties": {
-         "pipelines.reset.allowed": "false",
-         "pipelines.autoOptimize.zOrderCols": "id, customer_id"
-      }
-   },
-   {
-      "data_flow_id": "103",
-      "data_flow_group": "A2",
-      "source_system": "MYSQL",
-      "source_format": "eventhub",
-      "source_details": {
-         "source_schema_path": "tests/resources/schema/eventhub_iot_schema.ddl",
-         "eventhub.accessKeyName": "iotIngestionAccessKey",
-         "eventhub.name": "iot",
-         "eventhub.accessKeySecretName": "iotIngestionAccessKey",
-         "eventhub.secretsScopeName": "eventhubs_creds",
-         "kafka.sasl.mechanism": "PLAIN",
-         "kafka.security.protocol": "SASL_SSL",
-         "kafka.bootstrap.servers": "ganesh-standard.servicebus.windows.net:9093"
-      },
-      "bronze_database_dev": "bronze",
-      "bronze_database_staging": "bronze",
-      "bronze_database_prd": "bronze",
-      "bronze_table": "iot_cdc",
-      "bronze_reader_options": {
-         "maxOffsetsPerTrigger": "50000",
-         "startingOffsets": "latest",
-         "failOnDataLoss": "false",
-         "kafka.request.timeout.ms": "60000",
-         "kafka.session.timeout.ms": "60000"
-      },
-      "bronze_table_path_dev": "tests/resources/delta/iot_cdc",
-      "bronze_table_path_staging": "s3://db-dlt-meta-staging/demo/data/bronze/iot_cdc",
-      "bronze_table_path_prd": "s3://db-dlt-meta-prod/demo/data/bronze/iot_cdc",
-      "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/iot_cdc/bronze_data_quality_expectations.json",
-      "silver_database_dev": "silver",
-      "silver_table": "iot_cdc",
-      "silver_cdc_apply_changes": {
-         "keys": [
-            "device_id"
-         ],
-         "sequence_by": "timestamp",
-         "scd_type": "1",
-         "apply_as_deletes": "operation = 'DELETE'",
-         "except_column_list": []
-      },
-      "silver_table_path_dev": "tests/resources/data/silver/iot_cdc",
-      "silver_transformation_json_dev": "tests/resources/silver_transformations.json"
-   }
-]
+[{"data_flow_id": "100", "data_flow_group": "A1", "source_system": "MYSQL", "source_format": "cloudFiles", "source_details": {"source_database": "APP", "source_table": "CUSTOMERS", "source_path_dev": "tests/resources/data/customers", "source_schema_path": "tests/resources/schema/customer_schema.ddl"}, "bronze_database_dev": "bronze", "bronze_database_staging": "bronze", "bronze_database_prd": "bronze", "bronze_table": "customers_cdc", "bronze_reader_options": {"cloudFiles.format": "json", "cloudFiles.inferColumnTypes": "true", "cloudFiles.rescuedDataColumn": "_rescued_data"}, "bronze_table_path_dev": "tests/resources/delta/customers", "bronze_table_properties": {"pipelines.autoOptimize.managed": "false", "pipelines.reset.allowed": "false"}, "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/customers/bronze_data_quality_expectations.json", "silver_database_dev": "silver", "silver_database_staging": "silver", "silver_database_prd": "silver", "silver_table": "customers", "silver_cdc_apply_changes": {"keys": ["id"], "sequence_by": "operation_date", "scd_type": "1", "apply_as_deletes": "operation = 'DELETE'", "except_column_list": ["operation", "operation_date", "_rescued_data"]}, "silver_table_path_dev": "tests/resources/data/silver/customers", "silver_table_properties": {"pipelines.autoOptimize.managed": "false", "pipelines.reset.allowed": "false", "pipelines.autoOptimize.zOrderCols": "id,email"}, "silver_transformation_json_dev": "tests/resources/silver_transformations.json"}, {"data_flow_id": "101", "data_flow_group": "A1", "source_system": "MYSQL", "source_format": "cloudFiles", "source_details": {"source_database": "APP", "source_table": "TRANSACTIONS", "source_path_prd": "tests/resources/data/transactions", "source_path_dev": "tests/resources/data/transactions"}, "bronze_database_dev": "bronze", "bronze_database_staging": "bronze", "bronze_database_prd": "bronze", "bronze_table": "transactions_cdc", "bronze_reader_options": {"cloudFiles.format": "json", "cloudFiles.inferColumnTypes": "true", "cloudFiles.rescuedDataColumn": "_rescued_data"}, "bronze_table_path_dev": "tests/resources/delta/transactions", "bronze_table_path_staging": "s3://db-dlt-meta-staging/demo/data/bronze/transactions", "bronze_table_path_prd": "s3://db-dlt-meta-prod/demo/data/bronze/transactions", "bronze_table_properties": {"pipelines.reset.allowed": "false"}, "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/transactions/bronze_data_quality_expectations.json", "bronze_database_quarantine_dev": "bronze", "bronze_database_quarantine_staging": "bronze", "bronze_database_quarantine_prd": "bronze", "bronze_quarantine_table": "transactions_cdc_quarantine", "bronze_quarantine_table_path_dev": "tests/resources/data/bronze/transactions_quarantine", "silver_database_dev": "silver", "silver_database_preprd": "silver", "silver_database_prd": "silver", "silver_table": "transactions", "silver_cdc_apply_changes": {"keys": ["id"], "sequence_by": "operation_date", "scd_type": "1", "apply_as_deletes": "operation = 'DELETE'", "except_column_list": ["operation", "operation_date", "_rescued_data"]}, "silver_partition_columns": "transaction_date", "silver_table_path_dev": "tests/resources/data/silver/transactions", "silver_transformation_json_dev": "tests/resources/silver_transformations.json", "silver_table_properties": {"pipelines.reset.allowed": "false", "pipelines.autoOptimize.zOrderCols": "id, customer_id"}}, {"data_flow_id": "103", "data_flow_group": "A2", "source_system": "MYSQL", "source_format": "eventhub", "source_details": {"source_schema_path": "tests/resources/schema/eventhub_iot_schema.ddl", "eventhub.accessKeyName": "iotIngestionAccessKey", "eventhub.name": "iot", "eventhub.accessKeySecretName": "iotIngestionAccessKey", "eventhub.secretsScopeName": "eventhubs_creds", "kafka.sasl.mechanism": "PLAIN", "kafka.security.protocol": "SASL_SSL", "kafka.bootstrap.servers": "ganesh-standard.servicebus.windows.net:9093"}, "bronze_database_dev": "bronze", "bronze_database_staging": "bronze", "bronze_database_prd": "bronze", "bronze_table": "iot_cdc", "bronze_reader_options": {"maxOffsetsPerTrigger": "50000", "startingOffsets": "latest", "failOnDataLoss": "false", "kafka.request.timeout.ms": "60000", "kafka.session.timeout.ms": "60000"}, "bronze_table_path_dev": "tests/resources/delta/iot_cdc", "bronze_table_path_staging": "s3://db-dlt-meta-staging/demo/data/bronze/iot_cdc", "bronze_table_path_prd": "s3://db-dlt-meta-prod/demo/data/bronze/iot_cdc", "bronze_data_quality_expectations_json_dev": "tests/resources/dqe/iot_cdc/bronze_data_quality_expectations.json", "silver_database_dev": "silver", "silver_table": "iot_cdc", "silver_cdc_apply_changes": {"keys": ["device_id"], "sequence_by": "timestamp", "scd_type": "1", "apply_as_deletes": "operation = 'DELETE'", "except_column_list": []}, "silver_table_path_dev": "tests/resources/data/silver/iot_cdc", "silver_transformation_json_dev": "tests/resources/silver_transformations.json"}]
diff --git a/tests/resources/onboarding_ac_type2.json b/tests/resources/onboarding_ac_type2.json
@@ -31,6 +31,7 @@
            "sequence_by": "dmsTimestamp",
            "scd_type": "2",
            "apply_as_deletes": "Op = 'D'",
+           "apply_as_truncates": "Op = 'T'",
            "except_column_list": [
               "Op",
               "dmsTimestamp",
@@ -72,6 +73,7 @@
            "sequence_by": "dmsTimestamp",
            "scd_type": "2",
            "apply_as_deletes": "Op = 'D'",
+           "apply_as_truncates": "Op = 'T'",
            "except_column_list": [
               "Op",
               "dmsTimestamp",
diff --git a/tests/test_main.py b/tests/test_main.py
@@ -3,6 +3,10 @@
 import sys
 import copy
 from src import __main__
+from unittest.mock import MagicMock
+
+spark = MagicMock()
+OnboardDataflowspec = MagicMock()
 
 
 class MainTests(DLTFrameworkTestCase):
@@ -101,3 +105,73 @@ def test_main_layer_missing(self):
         sys.argv = list
         with self.assertRaises(Exception):
             __main__.main()
+
+    def test_main_bronze_silver_uc(self):
+        """Test bronze and silver onboarding for uc."""
+        OnboardDataflowspec.return_value = None
+        spark_mock = MagicMock("SparkSession")
+        spark.builder.appName("DLT-META_Onboarding_Task").getOrCreate().return_value = spark_mock
+        param_map = copy.deepcopy(self.onboarding_bronze_silver_params_uc_map)
+        param_map["onboard_layer"] = "bronze_silver"
+        list = ["dummy_test"]
+        for key in param_map:
+            list.append(f"--{key}={param_map[key]}")
+        sys.argv = list
+        __main__.main()
+        bronze_dataflowSpec_df = (self.spark.read.format("delta").table(
+            f"{param_map['database']}.{param_map['bronze_dataflowspec_table']}")
+        )
+        self.assertEqual(bronze_dataflowSpec_df.count(), 3)
+        silver_dataflowSpec_df = (self.spark.read.format("delta") .table(
+            f"{param_map['database']}.{param_map['silver_dataflowspec_table']}")
+        )
+        self.assertEqual(silver_dataflowSpec_df.count(), 3)
+        del param_map['onboard_layer']
+        del param_map['uc_enabled']
+        del param_map['bronze_dataflowspec_path']
+        del param_map['silver_dataflowspec_path']
+        OnboardDataflowspec.called_once_with(spark_mock, param_map, uc_enabled=True)
+
+    def test_onboarding(self):
+        mock_onboard_dataflowspec = OnboardDataflowspec
+        mock_args = MagicMock()
+        mock_args.onboard_layer = 'bronze_silver'
+        mock_args.uc_enabled = 'true'
+        mock_args.__dict__ = {
+            'onboard_layer': 'bronze_silver',
+            'uc_enabled': 'true',
+            'bronze_dataflowspec_path': 'path/to/bronze_dataflowspec',
+            'silver_dataflowspec_path': 'path/to/silver_dataflowspec'
+        }
+
+        spark_mock = MagicMock("SparkSession")
+        spark.builder.appName("DLT-META_Onboarding_Task").getOrCreate().return_value = spark_mock
+
+        mock_onboard_obj = MagicMock()
+        mock_onboard_dataflowspec.return_value = mock_onboard_obj
+
+        # Act
+        onboard_layer = mock_args.onboard_layer
+        uc_enabled = True if mock_args.uc_enabled and mock_args.uc_enabled.lower() == "true" else False
+        onboarding_args_dict = mock_args.__dict__
+        del onboarding_args_dict['onboard_layer']
+        del onboarding_args_dict['uc_enabled']
+        if uc_enabled:
+            if 'bronze_dataflowspec_path' in onboarding_args_dict:
+                del onboarding_args_dict['bronze_dataflowspec_path']
+            if 'silver_dataflowspec_path' in onboarding_args_dict:
+                del onboarding_args_dict['silver_dataflowspec_path']
+        onboard_obj = mock_onboard_dataflowspec(spark, onboarding_args_dict, uc_enabled=uc_enabled)
+
+        if onboard_layer.lower() == "bronze_silver":
+            onboard_obj.onboard_dataflow_specs()
+        elif onboard_layer.lower() == "bronze":
+            onboard_obj.onboard_bronze_dataflow_spec()
+        elif onboard_layer.lower() == "silver":
+            onboard_obj.onboard_silver_dataflow_spec()
+        else:
+            raise Exception("onboard_layer argument missing in commandline")
+
+        # Assert
+        mock_onboard_dataflowspec.assert_called_once_with(spark, {}, uc_enabled=True)
+        mock_onboard_obj.onboard_dataflow_specs.assert_called_once()
diff --git a/tests/test_onboard_dataflowspec.py b/tests/test_onboard_dataflowspec.py
@@ -121,6 +121,41 @@ def test_onboardDataFlowSpecs_with_merge(self):
             if bronze_row.dataFlowId == "103":
                 self.assertEqual(bronze_row.readerConfigOptions.get("maxOffsetsPerTrigger"), "60000")
 
+    def test_onboardDataFlowSpecs_with_merge_uc(self):
+        """Test for onboardDataflowspec with merge scenario."""
+        local_params = copy.deepcopy(self.onboarding_bronze_silver_params_uc_map)
+        local_params["onboarding_file_path"] = self.onboarding_json_file
+        del local_params["uc_enabled"]
+        onboardDataFlowSpecs = OnboardDataflowspec(self.spark, local_params, uc_enabled=True)
+        onboardDataFlowSpecs.onboard_dataflow_specs()
+        bronze_dataflowSpec_df = self.read_dataflowspec(
+            self.onboarding_bronze_silver_params_map['database'],
+            self.onboarding_bronze_silver_params_map['bronze_dataflowspec_table'])
+        silver_dataflowSpec_df = self.read_dataflowspec(
+            self.onboarding_bronze_silver_params_map['database'],
+            self.onboarding_bronze_silver_params_map['silver_dataflowspec_table'])
+        self.assertEqual(bronze_dataflowSpec_df.count(), 3)
+        self.assertEqual(silver_dataflowSpec_df.count(), 3)
+        local_params["overwrite"] = "False"
+        local_params["onboarding_file_path"] = self.onboarding_v2_json_file
+        onboardDataFlowSpecs = OnboardDataflowspec(self.spark, local_params, uc_enabled=True)
+        onboardDataFlowSpecs.onboard_dataflow_specs()
+        bronze_dataflowSpec_df = self.read_dataflowspec(
+            self.onboarding_bronze_silver_params_map['database'],
+            self.onboarding_bronze_silver_params_map['bronze_dataflowspec_table'])
+        silver_dataflowSpec_df = self.read_dataflowspec(
+            self.onboarding_bronze_silver_params_map['database'],
+            self.onboarding_bronze_silver_params_map['silver_dataflowspec_table'])
+        self.assertEqual(bronze_dataflowSpec_df.count(), 3)
+        self.assertEqual(silver_dataflowSpec_df.count(), 3)
+        bronze_df_rows = bronze_dataflowSpec_df.collect()
+        for bronze_df_row in bronze_df_rows:
+            bronze_row = BronzeDataflowSpec(**bronze_df_row.asDict())
+            if bronze_row.dataFlowId in ["101", "102"]:
+                self.assertIsNone(bronze_row.readerConfigOptions.get("cloudFiles.rescuedDataColumn"))
+            if bronze_row.dataFlowId == "103":
+                self.assertEqual(bronze_row.readerConfigOptions.get("maxOffsetsPerTrigger"), "60000")
+
     def test_onboardBronzeDataflowSpec_positive(self):
         """Test for onboardDataflowspec."""
         onboarding_params_map = copy.deepcopy(self.onboarding_bronze_silver_params_map)
diff --git a/tests/test_pipeline_readers.py b/tests/test_pipeline_readers.py

Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,6 @@ def read_kafka(spark, bronze_dataflow_spec, schema_json) -> DataFrame:`
`105`	`105`	`# add date, hour, and minute columns derived from eventhub enqueued timestamp`
`106`	`106`	`.selectExpr("*", "to_date(timestamp) as date", "hour(timestamp) as hour", "minute(timestamp) as minute")`
`107`	`107`	`)`
`108`		`-`
`109`	`108`	`if schema_json:`
`110`	`109`	`schema = StructType.fromJson(schema_json)`
`111`	`110`	`return (`