clouddq-migration/lib.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import re
import random
import string
import yaml
from yaml.loader import SafeLoader

class LineNumberLoader(SafeLoader):
    def construct_mapping(self, node, deep=False):
        """
            Method to add line number into config
        """
        mapping = super().construct_mapping(node, deep)
        mapping['__line__'] = node.start_mark.line
        return mapping

def validateConfigFields(config) -> None:
    """
        Method to validate the Config Fields

        return: None
    """
    if isinstance(config, dict):
        for key, value in config.items():
            if value is None:
                raise ValueError(f"Field '{key}' is None for the block at line {config.get('__line__')}")
            validateConfigFields(value)
    elif isinstance(config, list):
        for item in config:
            validateConfigFields(item)

def removeLineKeys(config) -> list:
    """
        Method to recursively remove '__line__' keys from config.

        return: config
    """
    if isinstance(config, list):
        return [removeLineKeys(item) for item in config]
    elif isinstance(config, dict):
        return {
            key: removeLineKeys(value)
            for key, value in config.items()
            if key != '__line__'
        }
    else:
        return config

def validateConfigFile(config_path) -> list:
    """
        Method to valide the Config File

        return: configs
    """
    # load the config file
    with open(config_path, 'r') as f:
        config_file = list(yaml.load_all(f, Loader=LineNumberLoader))

    # validate the config file
    for config in config_file:
        if not {'taskId'} <= config.keys():
            raise ValueError(
                "Config file must define the required config field: "
                "'taskId' at line ",config.get('__line__')
            )

        # validate format for taskId
        task = config['taskId']
        validate_task(task)

        # validate nested fields
        validateConfigFields(config)
    configs = [removeLineKeys(config) for config in config_file]
    return configs

def merge_configs(config, new_config) -> dict:
    """"
        Method to merge configs

        return: final_config
    """
    # Copy original config to avoid modifying it directly
    final_config = config.copy()

    # Update the final config with the new config
    for key, value in new_config.items():
        if isinstance(value, dict) and key in final_config:
            # Recursively update dictionary if the key exists and both are dictionaries
            final_config[key] = merge_configs(final_config.get(key, {}), value)
        else:
            # Otherwise, update or add the key
            final_config[key] = value

    return final_config

def validate_task(task)-> None:
    """
        Method to validate the format for tasks ids
    """
    if not re.match(r'^[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+\.[a-zA-Z0-9_-]+$', task):
        raise ValueError(
            f"Error: {task} does not match the expected format 'project_id.location_id.lake_id.task_id'")

def generate_config(yaml_data) -> dict:
    '''
        Method to generate the config of the task
    '''
    config = {
        'dataQualitySpec': {
            'rules': []
        }
    }
    rule_bindings = yaml_data.get('rule_bindings')
    entities = yaml_data.get('entities')
    rules = yaml_data.get('rules')
    row_filters = yaml_data.get('row_filters')
    if not rule_bindings:
        print('No rule binding Id present to migrate. Skipped')
    else:
        for rule_binding in rule_bindings:
            column_id = rule_bindings[rule_binding]['column_id']
            if 'entity_id' in rule_bindings[rule_binding]:
                entity_id = rule_bindings[rule_binding]['entity_id']
                project_id = entities[entity_id]['project_name']
                dataset_id = entities[entity_id]['dataset_name']
                table_id = entities[entity_id]['table_name']
                if column_id in entities[entity_id]['columns'].keys():
                    column = entities[entity_id]['columns'][column_id]['name']
                else:
                    column_id = column_id.upper()
                    column = entities[entity_id]['columns'][column_id]['name']
                config['resource'] =  f'//bigquery.googleapis.com/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}'
            else:
                entity_uri = rule_bindings[rule_binding]['entity_uri']
                if entity_uri.startswith('dataplex'):
                    print('Does not support enitity source type. ')
                    continue
                else:
                    entity = entity_uri.split('/')
                    project_id = entity[3]
                    dataset_id = entity[5]
                    table_id = entity[7]
                    column = column_id
                    config['resource'] =  f'//bigquery.googleapis.com/projects/{project_id}/datasets/{dataset_id}/tables/{table_id}'
            rule_ids = rule_bindings[rule_binding]['rule_ids']
            for rule_id in rule_ids:
                if isinstance(rule_id, str):
                    if rules[rule_id]['rule_type'] == 'NOT_NULL':
                        config['dataQualitySpec']['rules'].append({
                            'dimension': 'COMPLETENESS',
                            'column': column,
                            'non_null_expectation': {},
                        })
                    elif rules[rule_id]['rule_type'] == 'REGEX':
                        config['dataQualitySpec']['rules'].append({
                            'dimension': 'ACCURACY',
                            'column': column,
                            'regex_expectation': {
                                "regex": rules[rule_id]['params']['pattern']
                            },
                        })
                    else:
                        custom_sql_expr = f'NOT($column = "") '
                        custom_sql_expr = custom_sql_expr.replace(f"$column", column)
                        config['dataQualitySpec']['rules'].append({
                            'dimension': 'COMPLETENESS',
                            'column': column,
                            'row_condition_expectation': {
                                "sql_expression": custom_sql_expr
                            },
                        })
                else:
                    for key, value in rule_id.items():
                        rule = key
                        params = value

                    if rules[rule]['rule_type'] == 'CUSTOM_SQL_EXPR':
                        custom_sql_expr = rules[rule]['params']['custom_sql_expr']
                        for key, value in params.items():
                            custom_sql_expr = custom_sql_expr.replace(f"${key}", str(value))
                        custom_sql_expr = custom_sql_expr.replace(f"$column", column)
                        config['dataQualitySpec']['rules'].append({
                            'dimension': 'ACCURACY',
                            'column': column,
                            'row_condition_expectation': {
                                "sql_expression": custom_sql_expr
                            },
                        })
                    else:

                        custom_sql_statement = rules[rule]['params']['custom_sql_statement']

                        for key, value in params.items():
                            custom_sql_statement = custom_sql_statement.replace(f"${key}", str(value))
                        custom_sql_statement = custom_sql_statement.replace(f"$column", column)
                        custom_sql_statement = custom_sql_statement.replace(f"data", f'{project_id}.{dataset_id}.{table_id}')

                        config['dataQualitySpec']['rules'].append({
                            'dimension': 'ACCURACY',
                            'column': column,
                            'sql_assertion': {
                                "sql_statement": custom_sql_statement
                            },
                        })

    return config

def generate_id() -> str:
    '''
        Method to generate the autoDQ task id
    '''
    # generate datascan id
    letters_and_digits = string.ascii_lowercase + string.digits
    random_string = ''.join(random.choice(letters_and_digits) for i in range(28))
    datascan_id = f'auto-dq-{random_string}'
    return datascan_id