Merge pull request #11 from tanukifk/sandbox

Hugo Rialan · web-flow · commit fc5cdfdf4c4d · 2024-06-13T11:23:23.000+02:00
add config block
diff --git a/definitions/example.js b/definitions/example.js
@@ -11,6 +11,11 @@ const commonAssertionsResult = commonAssertions({
     // Set the 'dataform.projectConfig.vars.env' var in 'dataform.json' for this to work.
     // "disabledInEnvs": ["dv", "qa"]
   },
+  config: {
+    "first_table": {
+      "where": "updated_date >= CURRENT_DATE() - 7"
+    },
+  },
   rowConditions: {
     "first_table": {
       "id_not_null": "id IS NOT NULL",
diff --git a/includes/data_completeness_assertions.js b/includes/data_completeness_assertions.js
@@ -11,13 +11,14 @@
 
 /**
  * @param {Object} globalParams - See index.js for details.
+ * @param {string} filter - The condition to filter the data.
  * @param {string} tableName - The name of the table to check for data completeness.
  * @param {Object} columnConditions - An object mapping column names to their allowed percentage of null values. If a value is an object, it should have an `allowedPercentageNull` property.
  */
 
 const assertions = [];
 
-const createDataCompletenessAssertion = (globalParams, tableName, columnConditions) => {
+const createDataCompletenessAssertion = (globalParams, filter, tableName, columnConditions) => {
 
   for (let columnName in columnConditions) {
     const allowedPercentageNull = columnConditions[columnName];
@@ -27,10 +28,21 @@ const createDataCompletenessAssertion = (globalParams, tableName, columnConditio
       .schema(globalParams.schema)
       .description(`Check data completeness for ${tableName}.${columnName}, allowed percentage of null values: ${allowedPercentageNull}`)
       .tags("assert-data-completeness")
-      .query(ctx => `SELECT COUNT(*) AS total_rows,
+      .query(ctx => `
+                WITH
+                    filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(tableName)}
+                        WHERE
+                            ${filter}
+                    )      
+                    SELECT COUNT(*) AS total_rows,
                         SUM(CASE WHEN ${columnName} IS NULL THEN 1 ELSE 0 END) AS null_count
-                        FROM ${ctx.ref(tableName)}
-                        HAVING SAFE_DIVIDE(null_count, total_rows) > ${allowedPercentageNull / 100} AND null_count > 0 AND total_rows > 0`);
+                        FROM filtering
+                        HAVING SAFE_DIVIDE(null_count, total_rows) > ${allowedPercentageNull / 100} AND null_count > 0 AND total_rows > 0
+                    `);
 
     (globalParams.tags && globalParams.tags.forEach((tag) => assertion.tags(tag)));
 
@@ -41,11 +53,12 @@ const createDataCompletenessAssertion = (globalParams, tableName, columnConditio
 
 };
 
-module.exports = (globalParams, dataCompletenessConditions) => {
+module.exports = (globalParams, config, dataCompletenessConditions) => {
   // Loop through dataCompletenessConditions to create data completeness check assertions.
   for (let tableName in dataCompletenessConditions) {
     const columnConditions = dataCompletenessConditions[tableName];
-    createDataCompletenessAssertion(globalParams, tableName, columnConditions);
+    const filter = config[tableName]?.where ?? true;
+    createDataCompletenessAssertion(globalParams, filter, tableName, columnConditions);
   }
 
   return assertions;
diff --git a/includes/data_freshness_assertions.js b/includes/data_freshness_assertions.js
@@ -9,29 +9,41 @@
 
 /**
  * @param {Object} globalParams - See index.js for details.
+ * @param {string} filter - The condition to filter the data.
  * @param {string} tableName - The name of the table to check for data freshness.
  * @param {number} delayCondition - The maximum allowed delay (in units specified by `timeUnit`) for the data to be considered fresh.
  * @param {string} timeUnit - The unit of time to use for the delay condition. This should be a string that is valid in a SQL `DATE_DIFF` function, such as 'DAY', 'HOUR', etc.
  * @param {string} dateColumn - The name of the date column to check for data freshness.
+ * @param {string} timeZone - The name of the time zone for the date column.
  */
 
 const assertions = [];
 
-const createDataFreshnessAssertion = (globalParams, tableName, delayCondition, timeUnit, dateColumn, timeZone = "UTC") => {
+const createDataFreshnessAssertion = (globalParams, filter, tableName, delayCondition, timeUnit, dateColumn, timeZone = "UTC") => {
+  
   const assertion = assert(`assert_freshness_${tableName}`)
     .database(globalParams.database)
     .schema(globalParams.schema)
     .description(`Assert that data in ${tableName} is fresh with a delay less than ${delayCondition} ${timeUnit}`)
     .tags("assert-data-freshness")
     .query(ctx => `
                 WITH
+                    filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(tableName)}
+                        WHERE
+                            ${filter}
+                    ),      
+                    
                     freshness AS (
                         SELECT
                           ${["DAY", "WEEK", "MONTH", "QUARTER", "YEAR"].includes(timeUnit)
                               ? `DATE_DIFF(CURRENT_DATE("${timeZone}"), MAX(${dateColumn}), ${timeUnit})`
                               : `TIMESTAMP_DIFF(CURRENT_TIMESTAMP(), MAX(${dateColumn}), ${timeUnit})`} AS delay
                         FROM
-                            ${ctx.ref(tableName)}
+                            filtering
                     )
                 SELECT
                     *
@@ -48,7 +60,7 @@ const createDataFreshnessAssertion = (globalParams, tableName, delayCondition, t
   assertions.push(assertion);
 };
 
-module.exports = (globalParams, freshnessConditions) => {
+module.exports = (globalParams, config, freshnessConditions) => {
 
   // Loop through freshnessConditions to create assertions.
   for (let tableName in freshnessConditions) {
@@ -58,7 +70,8 @@ module.exports = (globalParams, freshnessConditions) => {
       dateColumn,
       timeZone
     } = freshnessConditions[tableName];
-    createDataFreshnessAssertion(globalParams, tableName, delayCondition, timeUnit, dateColumn, timeZone);
+    const filter = config[tableName]?.where ?? true;
+    createDataFreshnessAssertion(globalParams, filter, tableName, delayCondition, timeUnit, dateColumn, timeZone);
   }
 
   return assertions;
diff --git a/includes/referential_integrity_assertions.js b/includes/referential_integrity_assertions.js
@@ -11,6 +11,8 @@
 
 /**
  * @param {Object} globalParams - See index.js for details.
+ * @param {string} parentFilter - The condition to filter the data of parent table.
+ * @param {string} childFilter - The condition to filter the data of child table.
  * @param {Object} parentTable - The name of the parent table in the foreign key relationship.
  * @param {Object} parentKey - The name of the column in the parent table that is the primary key.
  * @param {Object} childTable - The name of the child table in the foreign key relationship.
@@ -19,18 +21,37 @@
 
 const assertions = [];
 
-const createReferentialIntegrityAssertion = (globalParams, parentTable, parentKey, childTable, childKey) => {
+const createReferentialIntegrityAssertion = (globalParams, parentFilter, childFilter, parentTable, parentKey, childTable, childKey) => {
 
   const assertion = assert(`assert_referential_integrity_${parentTable}_${childTable}`)
     .database(globalParams.database)
     .schema(globalParams.schema)
     .description(`Check referential integrity for ${childTable}.${childKey} referencing ${parentTable}.${parentKey}`)
     .tags("assert-referential-integrity")
     .query(ctx => `
-          SELECT pt.${parentKey}
-          FROM ${ctx.ref(parentTable)} AS pt
-          LEFT JOIN ${ctx.ref(childTable)} AS t ON t.${childKey} = pt.${parentKey}
-          WHERE t.${childKey} IS NULL
+                WITH
+                    parent_filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(parentTable)}
+                        WHERE
+                            ${parentFilter}
+                    ),
+
+                    child_filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(childTable)}
+                        WHERE
+                            ${childFilter}
+                    )
+
+                    SELECT pt.${parentKey}
+                    FROM parent_filtering AS pt
+                    LEFT JOIN child_filtering AS t ON t.${childKey} = pt.${parentKey}
+                    WHERE t.${childKey} IS NULL
         `);
 
   (globalParams.tags && globalParams.tags.forEach((tag) => assertion.tags(tag)));
@@ -40,16 +61,18 @@ const createReferentialIntegrityAssertion = (globalParams, parentTable, parentKe
   assertions.push(assertion);
 };
 
-module.exports = (globalParams, referentialIntegrityConditions) => {
+module.exports = (globalParams, config, referentialIntegrityConditions) => {
   for (let parentTable in referentialIntegrityConditions) {
     const relationships = referentialIntegrityConditions[parentTable];
+    const parentFilter = config[parentTable]?.where ?? true;
 
     relationships.forEach(({
       parentKey,
       childTable,
       childKey
     }) => {
-      createReferentialIntegrityAssertion(globalParams, parentTable, parentKey, childTable, childKey);
+      const childFilter = config[childTable]?.where ?? true;
+      createReferentialIntegrityAssertion(globalParams, parentFilter, childFilter, parentTable, parentKey, childTable, childKey);
     })
   }
   return assertions;
diff --git a/includes/row_condition_assertions.js b/includes/row_condition_assertions.js
@@ -11,22 +11,34 @@
 
 /**
  * @param {Object} globalParams - See index.js for details.
+ * @param {string} filter - The condition to filter the data.
  * @param {string} tableName - The name of the table to check for row conditions.
  * @param {string} conditionName - The name of the condition to check.
  * @param {string} conditionQuery - The SQL query that defines the condition to check.
  */
 
 const assertions = [];
 
-const createRowConditionAssertion = (globalParams, tableName, conditionName, conditionQuery) => {
+const createRowConditionAssertion = (globalParams, filter, tableName, conditionName, conditionQuery) => {
   const assertion = assert(`assert_${conditionName.replace(/-/g , "_")}_${tableName}`)
     .database(globalParams.database)
     .schema(globalParams.schema)
     .description(`Assert that rows in ${tableName} meet ${conditionName}`)
     .tags("assert-row-condition")
-    .query(ctx => `SELECT "Condition not met: ${conditionQuery}, Table: ${ctx.ref(tableName)}" AS assertion_description
-                       FROM ${ctx.ref(tableName)}
-                       WHERE NOT (${conditionQuery})`);
+    .query(ctx => `
+                WITH
+                    filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(tableName)}
+                        WHERE
+                            ${filter}
+                    )
+                    SELECT "Condition not met: ${conditionQuery}, Table: ${ctx.ref(tableName)}" AS assertion_description
+                        FROM filtering
+                        WHERE NOT (${conditionQuery})
+                    `);
 
   (globalParams.tags && globalParams.tags.forEach((tag) => assertion.tags(tag)));
 
@@ -35,13 +47,14 @@ const createRowConditionAssertion = (globalParams, tableName, conditionName, con
   assertions.push(assertion);
 };
 
-module.exports = (globalParams, rowConditions) => {
+module.exports = (globalParams, config, rowConditions) => {
 
   // Loop through rowConditions to create assertions.
   for (let tableName in rowConditions) {
     for (let conditionName in rowConditions[tableName]) {
       const conditionQuery = rowConditions[tableName][conditionName];
-      createRowConditionAssertion(globalParams, tableName, conditionName, conditionQuery);
+      const filter = config[tableName]?.where ?? true;
+      createRowConditionAssertion(globalParams, filter, tableName, conditionName, conditionQuery);
     }
   }
 
diff --git a/includes/unique_key_assertions.js b/includes/unique_key_assertions.js
@@ -11,24 +11,36 @@
 
 /**
  * @param {Object} globalParams - See index.js for details.
+ * @param {string} filter - The condition to filter the data.
  * @param {string} tableName - The name of the table to check for unique keys.
  * @param {Array} columns - An array of column names that should form a unique key.
  */
 
 const assertions = [];
 
-const createUniqueKeyAssertion = (globalParams, tableName, columns) => {
+const createUniqueKeyAssertion = (globalParams, filter, tableName, columns) => {
   const uniqueColumns = columns.join(', ');
 
   const assertion = assert(`assert_unique_key_${tableName}`)
     .database(globalParams.database)
     .schema(globalParams.schema)
     .description(`Check that values in columns (${uniqueColumns}) in ${tableName} form a unique key`)
     .tags("assert-unique-key")
-    .query(ctx => `SELECT ${uniqueColumns}
-                       FROM ${ctx.ref(tableName)}
+    .query(ctx => `
+                WITH
+                    filtering AS (
+                        SELECT
+                            *
+                        FROM
+                            ${ctx.ref(tableName)}
+                        WHERE
+                            ${filter}
+                    )
+                SELECT ${uniqueColumns}
+                       FROM filtering
                        GROUP BY ${uniqueColumns}
-                       HAVING COUNT(*) > 1`);
+                       HAVING COUNT(*) > 1
+                `);
 
   (globalParams.tags && globalParams.tags.forEach((tag) => assertion.tags(tag)));
 
@@ -37,12 +49,13 @@ const createUniqueKeyAssertion = (globalParams, tableName, columns) => {
   assertions.push(assertion);
 };
 
-module.exports = (globalParams, uniqueKeyConditions) => {
+module.exports = (globalParams, config, uniqueKeyConditions) => {
 
   // Loop through uniqueKeyConditions to create unique key check assertions.
   for (let tableName in uniqueKeyConditions) {
     const columns = uniqueKeyConditions[tableName];
-    createUniqueKeyAssertion(globalParams, tableName, columns);
+    const filter = config[tableName]?.where ?? true;
+    createUniqueKeyAssertion(globalParams, filter, tableName, columns);
   }
 
   return assertions;
diff --git a/index.js b/index.js
@@ -34,17 +34,18 @@ module.exports = ({
         tags: [],
         disabledInEnvs: []
     },
+    config = {},
     rowConditions = {},
     uniqueKeyConditions = {},
     dataFreshnessConditions = {},
     dataCompletenessConditions = {},
     referentialIntegrityConditions = {}
 }) => {
-    const rowConditionAssertionsResult = row_condition_assertions(globalAssertionsParams, rowConditions);
-    const uniqueKeyAssertionsResult = unique_key_assertions(globalAssertionsParams, uniqueKeyConditions);
-    const dataFreshnessAssertionsResult = data_freshness_assertions(globalAssertionsParams, dataFreshnessConditions);
-    const dataCompletenessAssertionsResult = data_completeness_assertions(globalAssertionsParams, dataCompletenessConditions);
-    const referentialIntegrityAssertionsResult = referential_integrity_assertions(globalAssertionsParams, referentialIntegrityConditions); // New assertion
+    const rowConditionAssertionsResult = row_condition_assertions(globalAssertionsParams, config, rowConditions);
+    const uniqueKeyAssertionsResult = unique_key_assertions(globalAssertionsParams, config, uniqueKeyConditions);
+    const dataFreshnessAssertionsResult = data_freshness_assertions(globalAssertionsParams, config, dataFreshnessConditions);
+    const dataCompletenessAssertionsResult = data_completeness_assertions(globalAssertionsParams, config, dataCompletenessConditions);
+    const referentialIntegrityAssertionsResult = referential_integrity_assertions(globalAssertionsParams, config, referentialIntegrityConditions); // New assertion
 
     return {
         rowConditionAssertions: rowConditionAssertionsResult,
