58 combining operations using cte

dbt_automation/assets/operations.template.yml

@@ -36,6 +36,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema: <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns_to_copy:
           - column1
           - column2
@@ -52,6 +57,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema: <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns:
           - columnname: <column name>
             columntype: <type to cast column to>
@@ -63,6 +73,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema: <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns:
           - columnname: <first column>
           - columnname: <second column>
@@ -76,6 +91,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema: <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns:
           - name: <string (column or const)>
             is_col: <yes or no>
@@ -94,6 +114,11 @@ operations:
         output_name: <enter the name of the output model>
         dest_schema: <enter your destination/output schema>
         operator: <can be "add", "sub", "mul", "div">
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         operands:
           - <name of the 1st column>
           - <name of the 2nd column or a const value>
@@ -107,6 +132,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema:  <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns:
           - <column name>
     - type: renamecolumns
@@ -117,6 +147,11 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema:  <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+              - ...
         columns:
           <old column name>: <new column name>
     - type: regexextraction
@@ -127,6 +162,157 @@ operations:
           source_name: <name of the source defined in source.yml; will be null for type "model">
         dest_schema:  <destination schema>
         output_name: <name of the output model>
+        source_columns:
+          - <column name>
+          - <column name>
+          - <column name>
+          - ...
         columns:
           column1: <regex operation>
           column2: <regex operation>
+
+    - type: mergeoperations
+      config:
+        dest_schema: <destination_schema>
+        output_name: mergeoperation
+        source_name: <name of the source defined in source.yml; will be null for type "model">
+        operations:
+          - type: castdatatypes
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: cast_output
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns:
+                - columnname: <column_name>
+                  columntype: <column_type>
+          - type: arithmetic
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              output_name: arithmetic_output
+              dest_schema: <intermediate_schema>
+              operator: <operator>
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              operands:
+                - <operand1>
+                - <operand2>
+                - <operand3>
+              output_column_name: <output_column_name>
+          - type: coalescecolumns
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: coalesce_columns
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns:
+                - columnname: <column_name>
+                - columnname: <column_name>
+              output_column_name: <output_column_name>
+          - type: concat
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: concat
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns:
+                - name: <column_name>
+                  is_col: <yes_or_no>
+                - name: <column_name>
+                  is_col: <yes_or_no>
+                - name: <column_name>
+                  is_col: <yes_or_no>
+              output_column_name: <output_column_name>
+          - type: dropcolumns
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: drop_column
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+          - ...
+              columns:
+                - <column_name>
+          - type: renamecolumns
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: rename_column
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns:
+                <old_column_name>: <new_column_name>
+          - type: flattenjson
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: flatten_json
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns_to_copy:
+                - <column_name>
+                - <column_name>
+              json_column: <json_column>
+              json_columns_to_copy:
+                - <json_column_name>
+                - <json_column_name>
+          - type: regexextraction
+            config:
+              input:
+                input_type: <source_type>
+                input_name: <source_name>
+                source_name: <source_name>
+              dest_schema: <intermediate_schema>
+              output_name: regex_extraction
+              source_columns:
+                - <column name>
+                - <column name>
+                - <column name>
+                - ...
+              columns:
+                <column_name>: <regex_pattern>
+                <column_name>: <regex_pattern>

dbt_automation/operations/arithmetic.py

@@ -5,41 +5,42 @@
 from logging import basicConfig, getLogger, INFO
 from dbt_automation.utils.dbtproject import dbtProject
 from dbt_automation.utils.interfaces.warehouse_interface import WarehouseInterface
+from dbt_automation.utils.columnutils import quote_columnname
+
 from dbt_automation.utils.tableutils import source_or_ref
 
 basicConfig(level=INFO)
 logger = getLogger()
 
 
 # pylint:disable=unused-argument,logging-fstring-interpolation
-def arithmetic(config: dict, warehouse: WarehouseInterface, project_dir: str):
+def arithmetic_dbt_sql(config: dict, warehouse: WarehouseInterface):
     """
     performs arithmetic operations: +/-/*//
     config["input"] is dict {"source_name": "", "input_name": "", "input_type": ""}
     """
-    output_name = config["output_name"]
-    dest_schema = config["dest_schema"]
     operator = config["operator"]
     operands = config["operands"]
     output_col_name = config["output_column_name"]
+    source_columns = config.get("source_columns", [])
 
     if operator not in ["add", "sub", "mul", "div"]:
-        raise ValueError("unknown operation")
+        raise ValueError("Unknown operation")
 
-    if len(operands) <= 1:
-        raise ValueError("need atleast two operands to perform operations")
+    if len(operands) < 2:
+        raise ValueError("At least two operands are required to perform operations")
 
-    if operator == "div" and (len(operands) != 2):
-        raise ValueError("division needs exactly two operands")
+    if operator == "div" and len(operands) != 2:
+        raise ValueError("Division requires exactly two operands")
 
-    # setup the dbt project
-    dbtproject = dbtProject(project_dir)
-    dbtproject.ensure_models_dir(dest_schema)
+    dbt_code = "SELECT "
 
-    dbt_code = "{{ config(materialized='table',schema='" + dest_schema + "') }}\n"
+    dbt_code += ", ".join(
+        [quote_columnname(col, warehouse.name) for col in source_columns]
+    )
 
     if operator == "add":
-        dbt_code += "SELECT *,"
+        dbt_code += ","
         dbt_code += "{{dbt_utils.safe_add(["
         for operand in operands:
             dbt_code += f"'{str(operand)}',"
@@ -48,14 +49,14 @@ def arithmetic(config: dict, warehouse: WarehouseInterface, project_dir: str):
         dbt_code += f" AS {output_col_name} "
 
     if operator == "mul":
-        dbt_code += "SELECT *,"
+        dbt_code += ","
         for operand in operands:
             dbt_code += f"{operand} * "
         dbt_code = dbt_code[:-2]
         dbt_code += f" AS {output_col_name} "
 
     if operator == "sub":
-        dbt_code += "SELECT *,"
+        dbt_code += ","
         dbt_code += "{{dbt_utils.safe_subtract(["
         for operand in operands:
             dbt_code += f"'{str(operand)}',"
@@ -64,17 +65,39 @@ def arithmetic(config: dict, warehouse: WarehouseInterface, project_dir: str):
         dbt_code += f" AS {output_col_name} "
 
     if operator == "div":
-        dbt_code += "SELECT *,"
+        dbt_code += ","
         dbt_code += "{{dbt_utils.safe_divide("
         for operand in operands:
             dbt_code += f"'{str(operand)}',"
         dbt_code += ")}}"
         dbt_code += f" AS {output_col_name} "
 
-    dbt_code += " FROM " + "{{" + source_or_ref(**config["input"]) + "}}" + "\n"
+    select_from = source_or_ref(**config["input"])
+    if config["input"]["input_type"] == "cte":
+        dbt_code += "\n FROM " + select_from + "\n"
+    else:
+        dbt_code += "\n FROM " + "{{" + select_from + "}}" + "\n"
+
+    return dbt_code
+
+
+def arithmetic(config: dict, warehouse: WarehouseInterface, project_dir: str):
+    """
+    Perform arithmetic operations and generate a DBT model.
+    """
+    config_sql = ""
+    if config["input"]["input_type"] != "cte":
+        config_sql = (
+            "{{ config(materialized='table', schema='" + config["dest_schema"] + "') }}"
+        )
+
+    sql = config_sql + "\n" + arithmetic_dbt_sql(config, warehouse)
+
+    dbtproject = dbtProject(project_dir)
+    dbtproject.ensure_models_dir(config["dest_schema"])
 
-    logger.info(f"writing dbt model {dbt_code}")
-    model_sql_path = dbtproject.write_model(dest_schema, output_name, dbt_code)
-    logger.info(f"dbt model {output_name} successfully created")
+    model_sql_path = dbtproject.write_model(
+        config["dest_schema"], config["output_name"], sql
+    )
 
     return model_sql_path