def load_configs(configs_path: Path, configs_type: DqConfigType) -> typing.Dict:
def load_rule_bindings_config(configs_path: Path) -> typing.Dict:
def load_rule_dimensions_config(configs_path: Path) -> list:
def load_entities_config(configs_path: Path) -> typing.Dict:
def load_rules_config(configs_path: Path) -> typing.Dict:
def load_row_filters_config(configs_path: Path) -> typing.Dict:
def load_metadata_registry_default_configs(
    configs_path: Path,
) -> MetadataRegistryDefaults:
def create_rule_binding_view_model(
    rule_binding_id: str,
    rule_binding_configs: typing.Dict,
    dq_summary_table_name: str,
    environment: str,
    configs_cache: DqConfigsCache,
    dq_summary_table_exists: bool = False,
    metadata: typing.Optional[typing.Dict] = None,
    debug: bool = False,
    progress_watermark: bool = True,
    default_configs: typing.Optional[typing.Dict] = None,
) -> str:
def create_entity_summary_model(
    entity_table_id: str,
    entity_target_rule_binding_configs: dict,
    gcp_project_id: str,
    gcp_bq_dataset_id: str,
    debug: bool = False,
) -> str:
def write_sql_string_as_dbt_model(
    model_id: str, sql_string: str, dbt_model_path: Path
) -> None:
def prepare_configs_from_rule_binding_id(
    rule_binding_id: str,
    rule_binding_configs: typing.Dict,
    dq_summary_table_name: str,
    environment: typing.Optional[str],
    configs_cache: DqConfigsCache,
    dq_summary_table_exists: bool = False,
    metadata: typing.Optional[typing.Dict] = None,
    progress_watermark: bool = True,
    default_configs: typing.Optional[typing.Dict] = None,
) -> typing.Dict:
def prepare_configs_cache(configs_path: Path) -> DqConfigsCache: class MetadataRegistryDefaults:
    """ """

    default_configs: dict

    @classmethod
    def from_dict(
        cls: MetadataRegistryDefaults, kwargs: dict
    ) -> MetadataRegistryDefaults:
        default_configs = {}
        logger.debug(f"Parsing input 'metadata_registry_defaults':\n {pformat(kwargs)}")
        for registry_scheme, registry_defaults in kwargs.items():
            scheme = EntityUriScheme.from_scheme(registry_scheme)
            default_configs[scheme.value] = {}
            expected_uri_fields = None
            if scheme == EntityUriScheme.DATAPLEX:
                expected_uri_fields = DATAPLEX_URI_FIELDS
            else:
                raise NotImplementedError(
                    f"'metadata_registry_default' not implemented for registry '{scheme}' "
                )
            for key, value in registry_defaults.items():
                if key not in expected_uri_fields:
                    raise ValueError(
                        f"'metadata_registry_default' for scheme '{scheme}' "
                        f"contains unexpected field '{key}'."
                    )
                default_configs[scheme][key] = value
        return MetadataRegistryDefaults(default_configs=default_configs)

    def to_dict(self: MetadataRegistryDefaults) -> dict:
        return self.default_configs

    def get_dataplex_registry_defaults(self, key: str | None = None) -> str | None:
        if self.default_configs.get("DATAPLEX", None):
            if key:
                return self.default_configs["DATAPLEX"].get(key, None)
            else:
                return self.default_configs["DATAPLEX"]
        else:
            return None class BigQueryClient:
    _gcp_credentials: GcpCredentials
    _client: bigquery.client.Client = None
    target_audience = "https://bigquery.googleapis.com"

    def __init__(
        self,
        gcp_credentials: GcpCredentials = None,
        gcp_project_id: str = None,
        gcp_service_account_key_path: Path = None,
        gcp_impersonation_credentials: str = None,
    ) -> None:
        if gcp_credentials:
            self._gcp_credentials = gcp_credentials
        else:
            self._gcp_credentials = GcpCredentials(
                gcp_project_id=gcp_project_id,
                gcp_service_account_key_path=gcp_service_account_key_path,
                gcp_impersonation_credentials=gcp_impersonation_credentials,
            )

    def get_connection(
        self, new: bool = False, project_id: str = None
    ) -> bigquery.client.Client:
        """Creates return new Singleton database connection"""
        if project_id or self._client is None or new:
            job_config = bigquery.QueryJobConfig(use_legacy_sql=False)
            client_project_id = (
                project_id if project_id else self._gcp_credentials.project_id
            )
            self._client = bigquery.Client(
                default_query_job_config=job_config,
                credentials=self._gcp_credentials.credentials,
                project=client_project_id,
                client_info=ClientInfo(user_agent=USER_AGENT_TAG),
            )
            return self._client
        else:
            return self._client

    def close_connection(self) -> None:
        if self._client:
            self._client.close()

    def get_dataset_region(self, dataset: str, project_id: str) -> str:
        client = self.get_connection(project_id=project_id)
        try:
            dataset_info = client.get_dataset(dataset)
        except KeyError as error:
            raise KeyError(f"\n\nInput dataset `{dataset}` is not valid.\n{error}")
        return dataset_info.location

    def check_query_dry_run(self, query_string: str, project_id: str = None) -> None:
        """check whether query is valid."""
        dry_run_job_config = bigquery.QueryJobConfig(
            dry_run=True, use_query_cache=False, use_legacy_sql=False
        )
        try:
            client = self.get_connection(project_id=project_id)
            query = CHECK_QUERY.safe_substitute(query_string=query_string.strip())
            # Start the query, passing in the extra configuration.
            query_job = client.query(
                query=query, timeout=10, job_config=dry_run_job_config
            )
            # A dry run query completes immediately.
            logger.debug(
                "This query will process {} bytes.".format(
                    query_job.total_bytes_processed
                )
            )
        except NotFound as e:
            table_name = re.search(RE_EXTRACT_TABLE_NAME, str(e))
            if table_name:
                table_name = table_name.group(1).replace(":", ".")
                raise AssertionError(f"Table name `{table_name}` does not exist.")
        except Forbidden as e:
            logger.error("User has insufficient permissions.")
            raise e

    def is_table_exists(self, table: str, project_id: str = None) -> bool:
        try:
            client = self.get_connection(project_id=project_id)
            client.get_table(table)
            return True
        except (NotFound, KeyError):
            return False

    def assert_required_columns_exist_in_table(
        self, table: str, project_id: str = None
    ) -> None:
        try:
            client = self.get_connection(project_id=project_id)
            table_ref = client.get_table(table)
            column_names = {column.name for column in table_ref.schema}
            failures = {}
            for column_name, column_type in REQUIRED_COLUMN_TYPES.items():
                if column_name not in column_names:
                    failures[
                        column_name
                    ] = f"ALTER TABLE `{table}` ADD COLUMN {column_name} {column_type};\n"
            if failures:
                raise ValueError(
                    f"Cannot find required column '{list(failures.keys())}' in BigQuery table '{table}'.\n"
                    f"Ensure they are created by running the following SQL script and retry:\n"
                    "```\n" + "\n".join(failures.values()) + "```"
                )
        except NotFound:
            logger.warning(
                f"Table {table} does not yet exist. It will be created in this run."
            )
        except KeyError as error:
            logger.fatal(f"Input table `{table}` is not valid.", exc_info=True)
            raise KeyError(f"\n\nInput table `{table}` is not valid.\n{error}")

    def table_from_string(self, full_table_id: str) -> bigquery.table.Table:
        return bigquery.table.Table.from_string(full_table_id)

    def is_dataset_exists(self, dataset: str, project_id: str = None) -> bool:
        try:
            client = self.get_connection(project_id=project_id)
            client.get_dataset(dataset)
            return True
        except (NotFound, KeyError):
            return False

    def execute_query(
        self,
        query_string: str,
        job_config: bigquery.job.QueryJobConfig = None,
        project_id: str = None,
    ) -> bigquery.job.QueryJob:
        """
        The method is used to execute the sql query
        Parameters:
        query_string (str) : sql query to be executed
        Returns:
            result of the sql execution is returned
        """

        client = self.get_connection(project_id=project_id)
        logger.debug(f"Query string is:\n```\n{query_string}\n```")
        job_id_prefix = "clouddq-"

        if not job_config:
            job_config = bigquery.QueryJobConfig(
                use_query_cache=False, use_legacy_sql=False
            )

        query_job = client.query(
            query_string, job_config=job_config, job_id_prefix=job_id_prefix
        )

        return query_job

    def get_table_schema(self, table: str, project_id: str = None) -> dict:

        client = self.get_connection(project_id=project_id)
        try:
            table_ref = client.get_table(table)
        except KeyError as error:
            raise KeyError(f"\n\nInput table `{table}` is not valid.\n{error}")
        columns = {}
        for column in table_ref.schema:
            column_configs = {
                "name": column.name,
                "type": column.field_type,
                "mode": column.mode,
                "data_type": column.field_type,
            }
            columns[column.name.upper()] = column_configs
        table_type = table_ref.table_type
        logger.debug(f"Table Type is : {table_type}")
        table_partitioning_type = table_ref.partitioning_type
        logger.debug(f"Table Partitioning Type is : {table_partitioning_type}")
        if table_partitioning_type:
            partition_fields = []
            time_partitioning = table_ref.time_partitioning
            if time_partitioning:
                partition_field = {}
                time_partitioning_type = time_partitioning.type_
                time_partitioning_field = time_partitioning.field
                if not time_partitioning_field:
                    partition_fields = "_PARTITIONTIME"
                else:
                    partition_field["name"] = time_partitioning_field
                    partition_field["type"] = columns[time_partitioning_field.upper()][
                        "type"
                    ]
                    partition_field["partitioning_type"] = time_partitioning_type
                    partition_fields.append(partition_field)

            range_partitioning = table_ref.range_partitioning
            if range_partitioning:
                raise NotImplementedError(
                    f"Bigquery Range Partitioning for table '{table}' is "
                    f"currently not supported."
                )
            columns_dict = {
                "columns": columns,
                "partition_fields": partition_fields,
            }
        else:
            columns_dict = {
                "columns": columns,
                "partition_fields": None,
            }

        logger.debug(f"Schema for table {table} is: {columns_dict}")
        return columns_dict class TargetTable:

    invocation_id: str = None
    bigquery_client: BigQueryClient = None

    def __init__(self, invocation_id: str, bigquery_client: BigQueryClient):
        self.invocation_id = invocation_id
        self.bigquery_client = bigquery_client

    def write_to_target_bq_table(
        self,
        partition_date: date,
        target_bigquery_summary_table: str,
        dq_summary_table_name: str,
        summary_to_stdout: bool = False,
    ) -> int:
        try:

            num_rows = load_target_table_from_bigquery(
                bigquery_client=self.bigquery_client,
                invocation_id=self.invocation_id,
                partition_date=partition_date,
                target_bigquery_summary_table=target_bigquery_summary_table,
                dq_summary_table_name=dq_summary_table_name,
                summary_to_stdout=summary_to_stdout,
            )
            return num_rows

        except Exception as error:

            raise error class CloudDqDataplexClient:
    _client: DataplexClient
    gcs_bucket_name: str

    def __init__(
        self,
        gcp_project_id: str | None = None,
        gcp_dataplex_lake_name: str | None = None,
        gcp_dataplex_region: str | None = None,
        gcs_bucket_name: str | None = None,
        gcp_credentials: GcpCredentials | None = None,
        dataplex_endpoint: str = "https://dataplex.googleapis.com",
    ) -> None:
        if gcs_bucket_name:
            self.gcs_bucket_name = gcs_bucket_name
        else:
            self.gcs_bucket_name = DEFAULT_GCS_BUCKET_NAME.format(
                gcp_dataplex_region=gcp_dataplex_region
            )
        self._client = DataplexClient(
            gcp_credentials=gcp_credentials,
            gcp_project_id=gcp_project_id,
            gcp_dataplex_lake_name=gcp_dataplex_lake_name,
            gcp_dataplex_region=gcp_dataplex_region,
            dataplex_endpoint=dataplex_endpoint,
        )

    def create_clouddq_task(  # noqa: C901
        self,
        task_id: str,
        clouddq_yaml_spec_file_path: str,
        clouddq_run_project_id: str,
        clouddq_run_bq_region: str,
        clouddq_run_bq_dataset: str,
        task_service_account: str,
        target_bq_result_project_name: str,
        target_bq_result_dataset_name: str,
        target_bq_result_table_name: str,
        task_trigger_spec_type: DATAPLEX_TASK_TRIGGER_TYPE = DATAPLEX_TASK_TRIGGER_TYPE.ON_DEMAND,  # noqa: E501
        task_description: str | None = None,
        task_labels: dict | None = None,
        clouddq_pyspark_driver_path: str | None = None,
        clouddq_executable_path: str | None = None,
        clouddq_executable_checksum_path: str | None = None,
        validate_only: bool = False,
        clouddq_pyspark_driver_filename: str = "clouddq_pyspark_driver.py",
        enable_experimental_bigquery_entity_uris: bool = True,
        enable_experimental_dataplex_gcs_validation: bool = True,
    ) -> Response:
        # Set default CloudDQ PySpark driver path if not manually overridden
        clouddq_pyspark_driver_path = self._validate_clouddq_artifact_path(
            clouddq_pyspark_driver_path, clouddq_pyspark_driver_filename
        )
        # Set default CloudDQ executable path if not manually overridden
        clouddq_executable_path = self._validate_clouddq_artifact_path(
            clouddq_executable_path, "clouddq-executable.zip"
        )
        # Set default CloudDQ executable checksum path if not manually overridden
        clouddq_executable_checksum_path = self._validate_clouddq_artifact_path(
            clouddq_executable_checksum_path, "clouddq-executable.zip.hashsum"
        )
        # Prepare input CloudDQ YAML specs path
        clouddq_yaml_spec_file_path = str(clouddq_yaml_spec_file_path)
        if clouddq_yaml_spec_file_path[:5] == "gs://":
            clouddq_configs_gcs_path = clouddq_yaml_spec_file_path
        else:
            clouddq_yaml_spec_file_path = Path(clouddq_yaml_spec_file_path)
            if clouddq_yaml_spec_file_path.is_file():
                upload_blob(
                    self.gcs_bucket_name,
                    clouddq_yaml_spec_file_path.name,
                    str(clouddq_yaml_spec_file_path.name),
                )
                gcs_uri = (
                    f"gs://{self.gcs_bucket_name}/{clouddq_yaml_spec_file_path.name}"
                )
                clouddq_configs_gcs_path = gcs_uri
            else:
                raise ValueError(
                    "'clouddq_yaml_spec_file_path' argument "
                    f"{clouddq_yaml_spec_file_path} "
                    "must either be a single file (`.yml` or `.zip`) "
                    "or a GCS path to the `.yml` or `.zip` configs file."
                )
        # Add user-agent tag as Task label
        allowed_user_agent_label = re.sub("[^0-9a-zA-Z]+", "-", USER_AGENT_TAG.lower())
        if task_labels:
            task_labels["user-agent"] = allowed_user_agent_label
        else:
            task_labels = {"user-agent": allowed_user_agent_label}
        # Prepare CloudDQ execution argumnets
        execution_arguments = (
            f"clouddq-executable.zip, "
            "ALL, "
            f"{clouddq_configs_gcs_path}, "
            f'--gcp_project_id="{clouddq_run_project_id}", '
            f'--gcp_region_id="{clouddq_run_bq_region}", '
            f'--gcp_bq_dataset_id="{clouddq_run_bq_dataset}", '
            f"--target_bigquery_summary_table="
            f'"{target_bq_result_project_name}.'
            f"{target_bq_result_dataset_name}."
            f'{target_bq_result_table_name}",'
        )
        # Set experimental flags
        if enable_experimental_bigquery_entity_uris:
            execution_arguments += " --enable_experimental_bigquery_entity_uris,"
        if enable_experimental_dataplex_gcs_validation:
            execution_arguments += "--enable_experimental_dataplex_gcs_validation,"
        # Prepare Dataplex Task message body for CloudDQ Job
        clouddq_post_body = {
            "spark": {
                "python_script_file": clouddq_pyspark_driver_path,
                "file_uris": [
                    f"{clouddq_executable_path}",
                    f"{clouddq_executable_checksum_path}",
                    f"{clouddq_configs_gcs_path}",
                ],
            },
            "execution_spec": {
                "args": {"TASK_ARGS": execution_arguments},
                "service_account": f"{task_service_account}",
            },
            "trigger_spec": {"type": task_trigger_spec_type},
            "description": task_description,
            "labels": task_labels,
        }
        # Set trigger_spec for RECURRING trigger type
        if task_trigger_spec_type == DATAPLEX_TASK_TRIGGER_TYPE.RECURRING:
            raise NotImplementedError(
                f"task_trigger_spec_type {task_trigger_spec_type} not yet supported."
            )
        response = self._client.create_dataplex_task(
            task_id=task_id,
            post_body=clouddq_post_body,
            validate_only=validate_only,
        )
        return response

    def get_clouddq_task_status(self, task_id: str) -> str:

        """
        Get the dataplex task status
        :param task_id: dataplex task id
        :return: Task status
        """
        res = self._client.get_dataplex_task_jobs(task_id)
        logger.info(f"Response status code is {res.status_code}")
        logger.info(f"Response text is {res.text}")
        resp_obj = json.loads(res.text)

        if res.status_code == 200:

            if (
                "jobs" in resp_obj
                and len(resp_obj["jobs"]) > 0  # noqa: W503
                and "state" in resp_obj["jobs"][0]  # noqa: W503
            ):
                task_status = resp_obj["jobs"][0]["state"]
                return task_status
        else:
            return res

    def delete_clouddq_task_if_exists(self, task_id: str) -> Response:

        """
        List the dataplex task jobs
        :param task_id: task id for dataplex task
        :return: Response object
        """
        get_task_response = self._client.get_dataplex_task(
            task_id=task_id,
        )
        if get_task_response.status_code == 200:
            delete_task_response = self._client.delete_dataplex_task(
                task_id=task_id,
            )
            if delete_task_response.status_code == 200:
                retry_iteration = 0
                get_task_response = self._client.get_dataplex_task(
                    task_id=task_id,
                )
                try:
                    while get_task_response.status_code != 404:
                        exponential_backoff(retry_iteration)
                        retry_iteration += 1
                        get_task_response = self._client.get_dataplex_task(
                            task_id=task_id,
                        )
                    logger.info(f"Successfully deleted Task ID: {task_id}")
                    return delete_task_response
                except RuntimeError as e:
                    logger.error(
                        f"Failed to delete Task ID: {task_id} with error: {e}",
                        exc_info=True,
                    )
            else:
                return delete_task_response
        else:
            return get_task_response

    def get_dataplex_lake(self, lake_name: str) -> Response:
        return self._client.get_dataplex_lake(lake_name)

    def _validate_clouddq_artifact_path(
        self, clouddq_artifact_path: str | None, artifact_name: str
    ) -> str:
        if not clouddq_artifact_path:
            clouddq_artifact_gcs_path = f"gs://{self.gcs_bucket_name}/{artifact_name}"
        else:
            clouddq_artifact_path = str(clouddq_artifact_path)
            clouddq_artifact_name = clouddq_artifact_path.split("/")[-1]
            if not clouddq_artifact_path[:5] == "gs://":
                raise ValueError(
                    f"Artifact path argument for {artifact_name}: "
                    f"{clouddq_artifact_path} must be a GCS path."
                )
            elif clouddq_artifact_name != artifact_name:
                raise ValueError(
                    f"Artifact path argument for {artifact_name}: "
                    f"{clouddq_artifact_path} must end with '{artifact_name}'."
                )
            else:
                clouddq_artifact_gcs_path = clouddq_artifact_path
        return clouddq_artifact_gcs_path

    def get_dataplex_entity(
        self,
        zone_id: str,
        entity_id: str,
        gcp_project_id: str = None,
        location_id: str = None,
        lake_name: str = None,
    ) -> DataplexEntity:
        logger.debug(f"CloudDqDataplex.get_dataplex_entity() arguments: {locals()}")
        params = {"view": "FULL"}
        response = self._client.get_entity(
            zone_id=zone_id,
            entity_id=entity_id,
            gcp_project_id=gcp_project_id,
            location_id=location_id,
            lake_name=lake_name,
            params=params,
        )
        if response.status_code == 200:
            return DataplexEntity.from_dict(
                entity_id=entity_id, kwargs=json.loads(response.text)
            )
        else:
            raise RuntimeError(
                f"Failed to retrieve Dataplex entity: "
                f"'/projects/{gcp_project_id}/locations/{location_id}"
                f"/lakes/{lake_name}/zones/{zone_id}/entities/{entity_id}':\n {response.text}"
            )

    def list_dataplex_entities(
        self,
        zone_id: str,
        prefix: str = None,
        data_path: str = None,
        gcp_project_id: str = None,
        location_id: str = None,
        lake_name: str = None,
    ) -> list[DataplexEntity]:
        params = {"page_size": 10}

        if prefix and data_path:
            raise ValueError("Either prefix or datapath should be passed but not both")
        if prefix:
            params.update({"filter": f"id=starts_with({prefix})"})
        if data_path:
            params.update({"filter": f"data_path=starts_with({data_path})"})

        response_dict = {}
        response = self._client.list_entities(
            zone_id=zone_id,
            params=params,
            gcp_project_id=gcp_project_id,
            location_id=location_id,
            lake_name=lake_name,
        )
        response_dict.update(response.json())

        while "nextPageToken" in response_dict:
            time.sleep(3)  # to avoid api limit exceed error of 4 calls per 10 sec
            next_page_token = response_dict["nextPageToken"]
            logger.debug("Getting next page...")
            page_token = {"page_token": f"{next_page_token}"}
            params.update(page_token)
            next_page_response = self._client.list_entities(
                zone_id=zone_id,
                params=params,
                gcp_project_id=gcp_project_id,
                location_id=location_id,
                lake_name=lake_name,
            ).json()
            logger.debug(f"Next page response {next_page_response}")

            if "nextPageToken" not in next_page_response:
                del response_dict["nextPageToken"]
                response_dict = update_dict(response_dict, next_page_response)
            else:
                response_dict = update_dict(response_dict, next_page_response)
                response_dict["nextPageToken"] = next_page_response["nextPageToken"]

        dataplex_entities = []
        if "entities" in response_dict:
            for entity in response_dict["entities"]:
                entity_with_schema = self.get_dataplex_entity(
                    entity_id=entity["id"],
                    zone_id=zone_id,
                    gcp_project_id=gcp_project_id,
                    location_id=location_id,
                    lake_name=lake_name,
                )
                dataplex_entities.append(entity_with_schema)
        return dataplex_entities class GcpCredentials:
    credentials: Credentials = None
    project_id: str = None
    user_id: str = None

    def __init__(
        self,
        credentials: Credentials = None,
        gcp_project_id: str = None,
        gcp_service_account_key_path: Path = None,
        gcp_impersonation_credentials: str = None,
    ) -> None:
        # Use Credentials object directly if provided
        if credentials:
            source_credentials = credentials
        # Use service account json key if provided
        elif gcp_service_account_key_path:
            source_credentials = service_account.Credentials.from_service_account_file(
                filename=gcp_service_account_key_path,
                scopes=TARGET_SCOPES,
                quota_project_id=gcp_project_id,
            )
        # Otherwise, use Application Default Credentials
        else:
            application_credentials_file = os.environ.get(
                "GOOGLE_APPLICATION_CREDENTIALS", None
            )
            if application_credentials_file:
                logger.info(
                    "Using environment variable GOOGLE_APPLICATION_CREDENTIALS "
                    f"path: {application_credentials_file}"
                )
            source_credentials, _ = google.auth.default(
                scopes=TARGET_SCOPES, quota_project_id=gcp_project_id
            )
        if not source_credentials.valid:
            self.__refresh_credentials(source_credentials)
        # Attempt service account impersonation if requested
        if gcp_impersonation_credentials:
            target_credentials = impersonated_credentials.Credentials(
                source_credentials=source_credentials,
                target_principal=gcp_impersonation_credentials,
                target_scopes=TARGET_SCOPES,
                lifetime=3600,
            )
            self.credentials = target_credentials
        else:
            # Otherwise use source_credentials
            self.credentials = source_credentials
        self.project_id = self.__resolve_project_id(
            credentials=self.credentials, project_id=gcp_project_id
        )
        self.user_id = self.__resolve_credentials_username(credentials=self.credentials)
        if self.user_id:
            logger.info("Successfully created GCP Client.")
        else:
            logger.warning(
                "Encountered error while retrieving user from GCP credentials.",
            )

    def __refresh_credentials(self, credentials: Credentials) -> str:
        # Attempt to refresh token if not currently valid
        try:
            auth_req = google.auth.transport.requests.Request()
            credentials.refresh(auth_req)
        except RefreshError as err:
            logger.error(
                "Could not get refreshed credentials for GCP. "
                "Reauthentication Required."
            )
            raise err

    def __resolve_credentials_username(self, credentials: Credentials) -> str:
        # Attempt to refresh token if not currently valid
        if not credentials.valid:
            self.__refresh_credentials(credentials=credentials)
        # Try to get service account credentials user_id
        if credentials.__dict__.get("_service_account_email"):
            user_id = credentials.service_account_email
        elif credentials.__dict__.get("_target_principal"):
            user_id = credentials.service_account_email
        else:
            # Otherwise try to get ADC credentials user_id
            request = google.auth.transport.requests.Request()
            token = credentials.id_token
            id_info = id_token.verify_oauth2_token(token, request)
            user_id = id_info["email"]
        return user_id

    def __resolve_project_id(
        self, credentials: Credentials, project_id: str = None
    ) -> str:
        """Get project ID from local configs"""
        if project_id:
            _project_id = project_id
        elif credentials.__dict__.get("_project_id"):
            _project_id = credentials.project_id
        else:
            _project_id = None
            logger.warning(
                "Could not retrieve project_id from GCP credentials.", exc_info=True
            )
        return _project_id class JsonEncoderDatetime(JsonEncoderStrFallback):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
        else:
            return super().default(obj) def add_cloud_logging_handler(logger: Logger):
    client = google.cloud.logging.Client()
    handler = CloudLoggingHandler(
        client=client,
        name="clouddq",
        labels={
            "name": APP_NAME,
            "releaseId": APP_VERSION,
        },
    )
    handler.setFormatter(JSONFormatter())
    logger.addHandler(handler) def get_json_logger():
    json_logger = logging.getLogger("clouddq-json-logger")
    if not len(json_logger.handlers):
        json_logger.setLevel(LOG_LEVEL)
        logging_stream_handler = logging.StreamHandler(sys.stdout)
        logging_stream_handler.setFormatter(JSONFormatter())
        json_logger.addHandler(logging_stream_handler)
    return json_logger def get_logger():
    logger = logging.getLogger("clouddq")
    if not len(logger.handlers):
        logger.setLevel(LOG_LEVEL)
        logging_stream_handler = logging.StreamHandler(sys.stderr)
        stream_formatter = logging.Formatter(
            "{asctime} {name} {levelname:8s} {message}", style="{"
        )
        logging_stream_handler.setFormatter(stream_formatter)
        logger.addHandler(logging_stream_handler)
    return logger class DbtRunner:
    dbt_path: Path
    dbt_profiles_dir: Path
    environment_target: str
    connection_config: DbtConnectionConfig
    dbt_rule_binding_views_path: Path
    dbt_entity_summary_path: Path

    def __init__(
        self,
        dbt_path: Optional[Path],
        dbt_profiles_dir: Optional[str],
        environment_target: Optional[str],
        gcp_project_id: Optional[str],
        gcp_region_id: Optional[str],
        gcp_bq_dataset_id: Optional[str],
        gcp_service_account_key_path: Optional[Path],
        gcp_impersonation_credentials: Optional[str],
        bigquery_client: Optional[BigQueryClient] = None,
        create_paths_if_not_exists: bool = True,
    ):
        # Prepare local dbt environment
        self.dbt_path = self._resolve_dbt_path(
            dbt_path=dbt_path,
            create_paths_if_not_exists=create_paths_if_not_exists,
            write_log=True,
        )
        self._prepare_dbt_project_path()
        self._prepare_dbt_main_path()
        self._prepare_rule_binding_view_path(write_log=True)
        self._prepare_entity_summary_path(write_log=True)
        # Prepare connection configurations
        self._resolve_connection_configs(
            dbt_profiles_dir=dbt_profiles_dir,
            environment_target=environment_target,
            gcp_project_id=gcp_project_id,
            gcp_bq_dataset_id=gcp_bq_dataset_id,
            bigquery_client=bigquery_client,
            gcp_region_id=gcp_region_id,
            gcp_service_account_key_path=gcp_service_account_key_path,
            gcp_impersonation_credentials=gcp_impersonation_credentials,
        )
        logger.debug(f"Using 'dbt_profiles_dir': {self.dbt_profiles_dir}")

    def run(
        self, configs: Dict, debug: bool = False, dry_run: bool = False
    ) -> None:
        logger.debug(f"Running dbt in path: {self.dbt_path}")
        if debug:
            self.test_dbt_connection()
        run_dbt(
            dbt_path=self.dbt_path,
            dbt_profile_dir=self.dbt_profiles_dir,
            configs=configs,
            environment=self.environment_target,
            debug=False,
            dry_run=dry_run,
        )

    def test_dbt_connection(self):
        run_dbt(
            dbt_path=self.dbt_path,
            dbt_profile_dir=self.dbt_profiles_dir,
            environment=self.environment_target,
            debug=True,
            dry_run=True,
        )

    def get_dbt_path(self) -> Path:
        self._resolve_dbt_path(self.dbt_path)
        return Path(self.dbt_path)

    def get_rule_binding_view_path(self) -> Path:
        self._prepare_rule_binding_view_path()
        return Path(self.dbt_rule_binding_views_path)

    def get_entity_summary_path(self) -> Path:
        self._prepare_entity_summary_path()
        return Path(self.dbt_entity_summary_path)

    def get_dbt_profiles_dir(self) -> Path:
        self._resolve_connection_configs(
            dbt_profiles_dir=self.dbt_profiles_dir,
            environment_target=self.environment_target,
        )
        return Path(self.dbt_profiles_dir)

    def get_dbt_environment_target(self) -> str:
        self._resolve_connection_configs(
            dbt_profiles_dir=self.dbt_profiles_dir,
            environment_target=self.environment_target,
        )
        return self.environment_target

    def get_dq_summary_dataset_region(self) -> str:
        self._resolve_connection_configs(
            dbt_profiles_dir=self.dbt_profiles_dir,
            environment_target=self.environment_target,
        )
        Path(self.dbt_profiles_dir)
        return self.environment_target

    def _resolve_connection_configs(
        self,
        dbt_profiles_dir: Optional[str],
        environment_target: Optional[str],
        bigquery_client: Optional[BigQueryClient] = None,
        gcp_project_id: Optional[str] = None,
        gcp_region_id: Optional[str] = None,
        gcp_bq_dataset_id: Optional[str] = None,
        gcp_service_account_key_path: Optional[Path] = None,
        gcp_impersonation_credentials: Optional[str] = None,
    ) -> None:
        if dbt_profiles_dir:
            dbt_profiles_dir = Path(dbt_profiles_dir).absolute()
            if not dbt_profiles_dir.joinpath("profiles.yml").is_file():
                raise ValueError(
                    f"Cannot find connection `profiles.yml` configurations at "
                    f"`dbt_profiles_dir` path: {dbt_profiles_dir}"
                )
            self.dbt_profiles_dir = dbt_profiles_dir
            self.environment_target = environment_target
        else:
            # create GcpDbtConnectionConfig
            connection_config = GcpDbtConnectionConfig(
                gcp_project_id=gcp_project_id,
                gcp_bq_dataset_id=gcp_bq_dataset_id,
                bigquery_client=bigquery_client,
                gcp_region_id=gcp_region_id,
                gcp_service_account_key_path=gcp_service_account_key_path,
                gcp_impersonation_credentials=gcp_impersonation_credentials,
            )
            self.connection_config = connection_config
            self.dbt_profiles_dir = Path(self.dbt_path)
            logger.debug(
                "Using dbt profiles.yml path: {self.dbt_profiles_dir}",
            )
            if environment_target:
                logger.debug(f"Using `environment_target`: {environment_target}")
                self.environment_target = environment_target
                self.connection_config.to_dbt_profiles_yml(
                    target_directory=self.dbt_profiles_dir,
                    environment_target=self.environment_target,
                )
            else:
                self.environment_target = DEFAULT_DBT_ENVIRONMENT_TARGET
                self.connection_config.to_dbt_profiles_yml(
                    target_directory=self.dbt_profiles_dir
                )

    def _resolve_dbt_path(
        self,
        dbt_path: str,
        create_paths_if_not_exists: bool = False,
        write_log: bool = False,
    ) -> Path:
        logger.debug(f"Current working directory: {Path().cwd()}")
        if not dbt_path:
            dbt_path = Path().cwd().joinpath("dbt").absolute()
            logger.debug(
                "No argument 'dbt_path' provided. Defaulting to use "
                f"'dbt' directory in current working directory at: {dbt_path}"
            )
        else:
            dbt_path = Path(dbt_path).absolute()
            if dbt_path.name != "dbt":
                dbt_path = dbt_path / "dbt"
        if not dbt_path.is_dir():
            if create_paths_if_not_exists:
                logger.debug(f"Creating a new dbt directory at 'dbt_path': {dbt_path}")
                dbt_path.mkdir(parents=True, exist_ok=True)
            else:
                raise ValueError(f"Provided 'dbt_path' does not exists: {dbt_path}")
        if write_log:
            logger.debug(f"Using 'dbt_path': {dbt_path}")
        return dbt_path

    def _prepare_dbt_project_path(self) -> None:
        dbt_project_path = self.dbt_path.absolute().joinpath("dbt_project.yml")
        if not dbt_project_path.is_file():
            logger.debug(
                f"Cannot find `dbt_project.yml` in path: {dbt_project_path} \n"
                f"Writing templated file to: {dbt_project_path}/dbt_project.yml"
            )
            write_templated_file_to_path(dbt_project_path, DBT_TEMPLATED_FILE_LOCATIONS)
        logger.debug(f"Using 'dbt_project_path': {dbt_project_path}")

    def _prepare_dbt_main_path(self) -> None:
        assert self.dbt_path.is_dir()
        dbt_main_path = self.dbt_path / "models" / "data_quality_engine"
        dbt_main_path.mkdir(parents=True, exist_ok=True)
        write_templated_file_to_path(
            dbt_main_path.joinpath("main.sql"), DBT_TEMPLATED_FILE_LOCATIONS
        )
        write_templated_file_to_path(
            dbt_main_path.joinpath("dq_summary.sql"), DBT_TEMPLATED_FILE_LOCATIONS
        )

    def _prepare_rule_binding_view_path(self, write_log: bool = False) -> None:
        assert self.dbt_path.is_dir()
        self.dbt_rule_binding_views_path = (
            self.dbt_path / "models" / "rule_binding_views"
        )
        self.dbt_rule_binding_views_path.mkdir(parents=True, exist_ok=True)
        if write_log:
            logger.debug(
                "Using rule_binding_views path: "
                f"{self.dbt_rule_binding_views_path.absolute()}/",
            )

    def _prepare_entity_summary_path(self, write_log: bool = False) -> None:
        assert self.dbt_path.is_dir()
        self.dbt_entity_summary_path = (
            self.dbt_path / "models" / "entity_dq_statistics"
        )
        self.dbt_entity_summary_path.mkdir(parents=True, exist_ok=True)
        if write_log:
            logger.debug(
                "Using entity_dq_statistics path: "
                f"{self.dbt_entity_summary_path.absolute()}/",
            ) def get_bigquery_dq_summary_table_name(
    dbt_path: Path,
    dbt_profiles_dir: Path,
    environment_target: str,
    table_name: str = "dq_summary",
) -> str:
    # Get bigquery project and dataset for dq_summary table names
    dbt_project_path = dbt_path.joinpath("dbt_project.yml")
    if not dbt_project_path.is_file():
        raise ValueError(
            "Not able to find 'dbt_project.yml' config file at "
            "input path {dbt_project_path}."
        )
    dbt_profiles_key = load_yaml(dbt_project_path, "profile")
    dbt_profiles_config = load_yaml(dbt_profiles_dir / "profiles.yml", dbt_profiles_key)
    logger.debug(f"Content of dbt_profiles.yml:\n {pformat(dbt_profiles_config)}")
    dbt_profile = dbt_profiles_config["OUTPUTS"][environment_target]
    dbt_project = dbt_profile["project"]
    if "{{" in dbt_project:
        dbt_project = extract_dbt_env_var(dbt_project)
    dbt_dataset = dbt_profile["dataset"]
    if "{{" in dbt_dataset:
        dbt_dataset = extract_dbt_env_var(dbt_dataset)
    dq_summary_table_name = f"{dbt_project}.{dbt_dataset}.{table_name}"
    return dq_summary_table_name def get_dbt_invocation_id(dbt_path: Path) -> str:

    try:
        logger.debug(f"Attempting to get dbt invocation_id from path {dbt_path}")
        # JSON file
        with open(dbt_path.joinpath("target", "manifest.json")) as f:
            # Reading from file
            data = json.loads(f.read())
            return data["metadata"]["invocation_id"]
    except (FileNotFoundError, KeyError) as e:
        logger.error(
            f"dbt invocation_id is not found, dbt hasn't run yet or "
            f"the execution has failed {e}\n",
            exc_info=True,
        ) def assert_not_none_or_empty(value: typing.Any, error_msg: str) -> None:
    """

    Args:
      value: typing.Any:
      error_msg: str:

    Returns:

    """
    if not value:
        raise ValueError(error_msg)
from datetime import datetime
from datetime import timezone
from pathlib import Path
from pprint import pformat
from typing import Optional
from clouddq import lib
from clouddq.classes.metadata_registry_defaults import MetadataRegistryDefaults
from clouddq.integration.bigquery.bigquery_client import BigQueryClient
from clouddq.integration.bigquery.dq_target_table_utils import TargetTable
from clouddq.integration.dataplex.clouddq_dataplex import CloudDqDataplexClient
from clouddq.integration.gcp_credentials import GcpCredentials
from clouddq.log import JsonEncoderDatetime
from clouddq.log import add_cloud_logging_handler
from clouddq.log import get_json_logger
from clouddq.log import get_logger
from clouddq.runners.dbt.dbt_runner import DbtRunner
from clouddq.runners.dbt.dbt_utils import get_bigquery_dq_summary_table_name
from clouddq.runners.dbt.dbt_utils import get_dbt_invocation_id
from clouddq.utils import assert_not_none_or_empty
import json
import logging
import logging.config
import click
import coloredlogs
            )
        # Load metadata
        metadata = json.loads(metadata)
        # Load Rule Bindings
        configs_path = Path(rule_binding_config_path)
        logger.debug(f"Loading rule bindings from: {configs_path.absolute()}")
        all_rule_bindings = lib.load_rule_bindings_config(Path(configs_path))
        # Prepare list of Rule Bindings in-scope for run
        target_rule_binding_ids = [
            r.strip().upper() for r in rule_binding_ids.split(",")
        ]
        if len(target_rule_binding_ids) == 1 and target_rule_binding_ids[0] == "ALL":
            target_rule_binding_ids = [
                rule_binding.upper() for rule_binding in all_rule_bindings.keys()
            ]
        logger.info(f"Preparing SQL for rule bindings: {target_rule_binding_ids}")
        # Load default configs for metadata registries
        registry_defaults: MetadataRegistryDefaults = (
            lib.load_metadata_registry_default_configs(Path(configs_path))
        )
        default_dataplex_projects = registry_defaults.get_dataplex_registry_defaults(
            "projects"
        )
        default_dataplex_locations = registry_defaults.get_dataplex_registry_defaults(
            "locations"
        )
        default_dataplex_lakes = registry_defaults.get_dataplex_registry_defaults(
            "lakes"
        )
        dataplex_registry_defaults = registry_defaults.get_dataplex_registry_defaults()
        # Prepare Dataplex Client from metadata registry defaults
        dataplex_client = CloudDqDataplexClient(
            gcp_credentials=gcp_credentials,
            gcp_project_id=default_dataplex_projects,
            gcp_dataplex_lake_name=default_dataplex_lakes,
            gcp_dataplex_region=default_dataplex_locations,
        )
        logger.debug(
            "Created CloudDqDataplexClient with arguments: "
            f"{gcp_credentials}, "
            f"{default_dataplex_projects}, "
            f"{default_dataplex_lakes}, "
            f"{default_dataplex_locations}, "
        )
        # Load all configs into a local cache
        configs_cache = lib.prepare_configs_cache(configs_path=Path(configs_path))
        configs_cache.resolve_dataplex_entity_uris(
            client=dataplex_client,
            bigquery_client=bigquery_client,
            default_configs=dataplex_registry_defaults,
            target_rule_binding_ids=target_rule_binding_ids,
            enable_experimental_bigquery_entity_uris=enable_experimental_bigquery_entity_uris,
            enable_experimental_dataplex_gcs_validation=enable_experimental_dataplex_gcs_validation,
        )
        # Get Entities for entity-level summary views
        target_entity_summary_configs: dict = (
            configs_cache.get_entities_configs_from_rule_bindings(
                target_rule_binding_ids=target_rule_binding_ids,
            )
        )
        # Create Rule_binding views
        for rule_binding_id in target_rule_binding_ids:
            rule_binding_configs = all_rule_bindings.get(rule_binding_id, None)
            assert_not_none_or_empty(
                rule_binding_configs,
                f"Target Rule Binding Id: {rule_binding_id} not found "
                f"in config path {configs_path.absolute()}.",
            )
            if debug:
                logger.debug(
                    f"Creating sql string from configs for rule binding: "
                    f"{rule_binding_id}"
                )
                logger.debug(
                    f"Rule binding config json:\n{pformat(rule_binding_configs)}"
                )
            sql_string = lib.create_rule_binding_view_model(
                rule_binding_id=rule_binding_id,
                rule_binding_configs=rule_binding_configs,
                dq_summary_table_name=dq_summary_table_name,
                configs_cache=configs_cache,
                environment=environment_target,
                metadata=metadata,
                debug=print_sql_queries,
                progress_watermark=progress_watermark,
                default_configs=dataplex_registry_defaults,
                dq_summary_table_exists=dq_summary_table_exists,
            )
            if not skip_sql_validation:
                logger.debug(
                    f"Validating generated SQL code for rule binding "
                    f"{rule_binding_id} using BigQuery dry-run client.",
                )
                bigquery_client.check_query_dry_run(query_string=sql_string)
            logger.debug(
                f"*** Writing sql to {dbt_rule_binding_views_path.absolute()}/"
                f"{rule_binding_id}.sql",
            )
            lib.write_sql_string_as_dbt_model(
                model_id=rule_binding_id,
                sql_string=sql_string,
                dbt_model_path=dbt_rule_binding_views_path,
            )
        # clean up old rule_bindings
        for view in dbt_rule_binding_views_path.glob("*.sql"):
            if view.stem.upper() not in target_rule_binding_ids:
                view.unlink()
        logger.info(
            f"target_entity_summary_configs:\n{pformat(target_entity_summary_configs)}"
        )
        # create entity-level summary table models
        for (
            entity_table_id,
            entity_configs_dict,
        ) in target_entity_summary_configs.items():
            rule_binding_ids_list = entity_configs_dict.get("rule_binding_ids_list")
            assert_not_none_or_empty(
                rule_binding_ids_list,
                f"Internal Error: no rule_binding_id found for entity_table_id {entity_table_id}.",
            )

            sql_string = lib.create_entity_summary_model(