Merge pull request #86 from DalgoT4D/arithmetic-op-changes

minor changes in arithmetic operation
DalgoT4D · Mar 17, 2024 · 7df33c2 · 7df33c2
2 parents 10001eb + 7ae61fe
commit 7df33c2
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 22 deletions.
diff --git a/dbt_automation/assets/operations.template.yml b/dbt_automation/assets/operations.template.yml
@@ -126,9 +126,12 @@ operations:
           - <column name>
           - ...
         operands:
-          - <name of the 1st column>
-          - <name of the 2nd column or a const value>
-          - <name of the 3nd column or a const value>
+          - value: <string (column name or const)>
+            is_col: <boolean>
+          - value: <string (column name or const)>
+            is_col: <boolean>
+          - value: <string (column name or const)>
+            is_col: <boolean>
         output_column_name: <output column name>
     - type: dropcolumns
       config:
@@ -312,9 +315,12 @@ operations:
                 - <column name>
                 - ...
               operands:
-                - <operand1>
-                - <operand2>
-                - <operand3>
+                - value: <string (column name or const)>
+                  is_col: <boolean>
+                - value: <string (column name or const)>
+                  is_col: <boolean>
+                - value: <string (column name or const)>
+                  is_col: <boolean>
               output_column_name: <output_column_name>
           - type: coalescecolumns
             config:

diff --git a/dbt_automation/operations/arithmetic.py b/dbt_automation/operations/arithmetic.py
@@ -5,7 +5,7 @@
 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.columnutils import quote_columnname, quote_constvalue
 
 from dbt_automation.utils.tableutils import source_or_ref
 
@@ -20,7 +20,7 @@ def arithmetic_dbt_sql(config: dict, warehouse: WarehouseInterface):
     config["input"] is dict {"source_name": "", "input_name": "", "input_type": ""}
     """
     operator = config["operator"]
-    operands = config["operands"]
+    operands = config["operands"]  # {"is_col": true, "value": "1"}[]
     output_col_name = config["output_column_name"]
     source_columns = config.get("source_columns", [])
 
@@ -39,27 +39,30 @@ def arithmetic_dbt_sql(config: dict, warehouse: WarehouseInterface):
         [quote_columnname(col, warehouse.name) for col in source_columns]
     )
 
+    # dbt_utils function safe_add, safe_subtract, safe_divide
+    # takes input as single quoted fields for column eg. 'field1'. regardless of the warehouse
+
     if operator == "add":
         dbt_code += ","
         dbt_code += "{{dbt_utils.safe_add(["
         for operand in operands:
-            dbt_code += f"'{quote_columnname(str(operand), warehouse.name)}',"
+            dbt_code += f"{quote_constvalue(str(operand['value']), warehouse.name)},"
         dbt_code = dbt_code[:-1]
         dbt_code += "])}}"
         dbt_code += f" AS {output_col_name} "
 
     if operator == "mul":
         dbt_code += ","
         for operand in operands:
-            dbt_code += f"{operand} * "
+            dbt_code += f"{quote_columnname(str(operand['value']), warehouse.name) if operand['is_col'] else quote_constvalue(str(operand['value']), warehouse.name)} * "
         dbt_code = dbt_code[:-2]
         dbt_code += f" AS {output_col_name} "
 
     if operator == "sub":
         dbt_code += ","
         dbt_code += "{{dbt_utils.safe_subtract(["
         for operand in operands:
-            dbt_code += f"'{quote_columnname(str(operand), warehouse.name)}',"
+            dbt_code += f"{quote_constvalue(str(operand['value']), warehouse.name)},"
         dbt_code = dbt_code[:-1]
         dbt_code += "])}}"
         dbt_code += f" AS {output_col_name} "
@@ -68,7 +71,7 @@ def arithmetic_dbt_sql(config: dict, warehouse: WarehouseInterface):
         dbt_code += ","
         dbt_code += "{{dbt_utils.safe_divide("
         for operand in operands:
-            dbt_code += f"'{quote_columnname(str(operand), warehouse.name)}',"
+            dbt_code += f"{quote_constvalue(str(operand['value']), warehouse.name)},"
         dbt_code += ")}}"
         dbt_code += f" AS {output_col_name} "
 

diff --git a/tests/warehouse/test_bigquery_ops.py b/tests/warehouse/test_bigquery_ops.py
@@ -351,7 +351,10 @@ def test_arithmetic_add(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "add",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "output_column_name": "add_col",
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
         }
@@ -391,7 +394,10 @@ def test_arithmetic_sub(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "sub",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "output_column_name": "sub_col",
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
         }
@@ -431,7 +437,10 @@ def test_arithmetic_mul(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "mul",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "output_column_name": "mul_col",
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
         }
@@ -471,7 +480,10 @@ def test_arithmetic_div(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "div",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "output_column_name": "div_col",
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
         }
@@ -632,7 +644,10 @@ def test_merge_operation(self):
                     "type": "arithmetic",
                     "config": {
                         "operator": "add",
-                        "operands": ["measure1", "measure2"],
+                        "operands": [
+                            {"value": "measure1", "is_col": True},
+                            {"value": "measure2", "is_col": True},
+                        ],
                         "output_column_name": "add_col",
                         "source_columns": [
                             "NGO",

diff --git a/tests/warehouse/test_postgres_ops.py b/tests/warehouse/test_postgres_ops.py
@@ -359,7 +359,10 @@ def test_arithmetic_add(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "add",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
             "output_column_name": "add_col",
         }
@@ -399,7 +402,10 @@ def test_arithmetic_sub(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "sub",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
             "output_column_name": "sub_col",
         }
@@ -439,7 +445,10 @@ def test_arithmetic_mul(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "mul",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
             "output_column_name": "mul_col",
         }
@@ -479,7 +488,10 @@ def test_arithmetic_div(self):
             "dest_schema": "pytest_intermediate",
             "output_name": output_name,
             "operator": "div",
-            "operands": ["measure1", "measure2"],
+            "operands": [
+                {"value": "measure1", "is_col": True},
+                {"value": "measure2", "is_col": True},
+            ],
             "source_columns": ["NGO", "Month", "measure1", "measure2", "Indicator"],
             "output_column_name": "div_col",
         }
@@ -647,7 +659,10 @@ def test_merge_operation(self):
                     "type": "arithmetic",
                     "config": {
                         "operator": "add",
-                        "operands": ["measure1", "measure2"],
+                        "operands": [
+                            {"value": "measure1", "is_col": True},
+                            {"value": "measure2", "is_col": True},
+                        ],
                         "output_column_name": "add_col",
                         "source_columns": [
                             "NGO",