pycarlo 0.12.24__py3-none-any.whl

pycarlo/features/dbt/dbt_importer.py

@@ -0,0 +1,208 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Dict, Optional, Union, cast
+from uuid import UUID
+
+from dataclasses_json import LetterCase, dataclass_json
+
+from pycarlo.common import get_logger, http
+from pycarlo.common.files import BytesFileReader, JsonFileReader, to_path
+from pycarlo.common.settings import (
+    HEADER_MCD_TELEMETRY_REASON,
+    HEADER_MCD_TELEMETRY_SERVICE,
+    RequestReason,
+)
+from pycarlo.core import Client, Query
+from pycarlo.features.dbt.queries import (
+    GET_DBT_UPLOAD_URL,
+    SEND_DBT_ARTIFACTS_EVENT,
+)
+from pycarlo.features.pii import PiiFilterer, PiiService
+from pycarlo.features.user import UserService
+
+logger = get_logger(__name__)
+
+
+class InvalidArtifactsException(Exception):
+    pass
+
+
+class InvalidFileFormatException(Exception):
+    pass
+
+
+@dataclass_json(letter_case=LetterCase.CAMEL)  # type: ignore
+@dataclass
+class DbtArtifacts:
+    manifest: str
+    run_results: str
+    logs: Optional[str]
+
+
+class DbtImporter:
+    """
+    Import dbt run artifacts to Monte Carlo
+    """
+
+    DEFAULT_PROJECT_NAME = "default-project"
+    DEFAULT_JOB_NAME = "default-job"
+
+    def __init__(
+        self,
+        mc_client: Optional[Client] = None,
+        user_service: Optional[UserService] = None,
+        pii_service: Optional[PiiService] = None,
+        print_func: Callable = logger.info,
+    ):
+        self._mc_client = mc_client or Client()
+        self._user_service = user_service or UserService(mc_client=self._mc_client)
+        self._pii_service = pii_service or PiiService(mc_client=self._mc_client)
+        self._print_func = print_func
+        self._pii_filterer = self._init_pii_filterer()
+
+    def import_run(
+        self,
+        manifest_path: Union[Path, str],
+        run_results_path: Union[Path, str],
+        logs_path: Optional[Union[Path, str]] = None,
+        project_name: str = DEFAULT_PROJECT_NAME,
+        job_name: str = DEFAULT_JOB_NAME,
+        resource_id: Optional[Union[str, UUID]] = None,
+    ):
+        """
+        Import artifacts from a single dbt command execution.
+
+        :param manifest_path: local path to the dbt manifest file (manifest.json)
+        :param run_results_path: local path to the dbt run results file (run_results.json)
+        :param logs_path: local path to a file containing dbt run logs
+        :param project_name: Project name (perhaps a logical group of dbt models, analogous to a
+                             project in dbt Cloud)
+        :param job_name: Job name (perhaps a logical sequence of dbt commands, analogous to a
+                         job in dbt Cloud)
+        :param resource_id: identifier of a Monte Carlo resource (warehouse or lake) to use to
+                            resolve dbt models to tables, this will be required if you have more
+                            than one
+        """
+        # get resource
+        resource = self._user_service.get_resource(resource_id)
+
+        # read local artifacts
+        manifest = JsonFileReader(manifest_path).read()
+        run_results = JsonFileReader(run_results_path).read()
+        logs = BytesFileReader(logs_path).read() if logs_path else None
+
+        # extract dbt invocation id (and verify it is the same for each artifact)
+        invocation_id = self._get_invocation_id(
+            manifest_path=manifest_path,
+            manifest=manifest,
+            run_results_path=run_results_path,
+            run_results=run_results,
+        )
+
+        # upload artifacts to S3 (using pre-signed URLs)
+        artifacts = DbtArtifacts(
+            manifest=self._upload_artifact(
+                project_name=project_name,
+                invocation_id=invocation_id,
+                file_path=to_path(manifest_path),
+                content=manifest,
+            ),
+            run_results=self._upload_artifact(
+                project_name=project_name,
+                invocation_id=invocation_id,
+                file_path=to_path(run_results_path),
+                content=run_results,
+            ),
+            logs=self._upload_artifact(
+                project_name=project_name,
+                invocation_id=invocation_id,
+                file_path=to_path(logs_path),  # type: ignore
+                content=logs,
+            )
+            if logs
+            else None,
+        )
+
+        # publish event indicating run artifacts are ready for processing
+        self._mc_client(
+            query=SEND_DBT_ARTIFACTS_EVENT,
+            variables=dict(
+                projectName=project_name,
+                jobName=job_name,
+                invocationId=invocation_id,
+                artifacts=artifacts.to_dict(),  # type: ignore
+                resourceId=str(resource.id),
+            ),
+            additional_headers={
+                HEADER_MCD_TELEMETRY_REASON: RequestReason.SERVICE.value,
+                HEADER_MCD_TELEMETRY_SERVICE: "dbt_importer",
+            },
+        )
+
+        self._print_func("Finished sending run artifacts to Monte Carlo")
+
+    def _get_invocation_id(
+        self,
+        manifest_path: Union[Path, str],
+        manifest: Dict,
+        run_results_path: Union[Path, str],
+        run_results: Dict,
+    ) -> str:
+        manifest_invocation_id = self._extract_invocation_id(path=manifest_path, data=manifest)
+        run_results_invocation_id = self._extract_invocation_id(
+            path=run_results_path, data=run_results
+        )
+
+        if manifest_invocation_id != run_results_invocation_id:
+            raise InvalidArtifactsException(
+                "dbt invocation ids do not match between manifest and run results files"
+            )
+
+        return manifest_invocation_id
+
+    @staticmethod
+    def _extract_invocation_id(path: Union[Path, str], data: Dict) -> str:
+        try:
+            return data["metadata"]["invocation_id"]
+        except KeyError:
+            raise InvalidArtifactsException(
+                f"Unable to get dbt invocation id from '{path}'. Unexpected file format"
+            )
+
+    def _upload_artifact(
+        self,
+        project_name: str,
+        invocation_id: str,
+        file_path: Path,
+        content: Union[bytes, str, Dict],
+    ) -> str:
+        self._print_func(f"Uploading {file_path.name}...")
+        http.upload(
+            method="put",
+            url=self._get_presigned_url(
+                project_name=project_name, invocation_id=invocation_id, file_name=file_path.name
+            ),
+            content=self._pii_filterer.filter_content(content),
+        )
+        return file_path.name
+
+    def _get_presigned_url(self, project_name: str, invocation_id: str, file_name: str) -> str:
+        response = cast(
+            Query,
+            self._mc_client(
+                query=GET_DBT_UPLOAD_URL,
+                variables=dict(
+                    projectName=project_name, invocationId=invocation_id, fileName=file_name
+                ),
+                additional_headers={
+                    HEADER_MCD_TELEMETRY_REASON: RequestReason.SERVICE.value,
+                    HEADER_MCD_TELEMETRY_SERVICE: "dbt_importer",
+                },
+            ),
+        )
+
+        return cast(str, response.get_dbt_upload_url)
+
+    def _init_pii_filterer(self):
+        pii_filters = self._pii_service.get_pii_filters_config()
+        return PiiFilterer(filters_config=pii_filters)

pycarlo/features/dbt/queries.py

@@ -0,0 +1,31 @@
+GET_DBT_UPLOAD_URL = """
+query getDbtUploadUrl(
+  $projectName: String!,
+  $invocationId: String!,
+  $fileName: String!) {
+  getDbtUploadUrl(
+    projectName: $projectName,
+    invocationId: $invocationId,
+    fileName: $fileName
+  )
+}
+"""
+
+SEND_DBT_ARTIFACTS_EVENT = """
+mutation sendDbtArtifactsEvent(
+  $projectName: String!,
+  $jobName: String!,
+  $invocationId: UUID!,
+  $artifacts: DbtArtifactsInput!,
+  $resourceId: UUID) {
+  sendDbtArtifactsEvent(
+    projectName: $projectName,
+    jobName: $jobName,
+    invocationId: $invocationId,
+    artifacts: $artifacts,
+    resourceId: $resourceId
+  ) {
+    ok
+  }
+}
+"""

pycarlo/features/exceptions.py

@@ -0,0 +1,18 @@
+from pycarlo.features.circuit_breakers.exceptions import (
+    CircuitBreakerPipelineException,
+    CircuitBreakerPollException,
+)
+from pycarlo.features.user.exceptions import (
+    MultipleResourcesFoundException,
+    ResourceNotFoundException,
+    UserServiceException,
+)
+
+__all__ = [
+    "CircuitBreakerPipelineException",
+    "CircuitBreakerPipelineException",
+    "CircuitBreakerPollException",
+    "MultipleResourcesFoundException",
+    "ResourceNotFoundException",
+    "UserServiceException",
+]

pycarlo/features/metadata/__init__.py

@@ -0,0 +1,32 @@
+from pycarlo.features.metadata.asset_allow_block_list import AssetAllowBlockList
+from pycarlo.features.metadata.asset_filters_container import AssetFiltersContainer
+from pycarlo.features.metadata.base_allow_block_list import (
+    BaseAllowBlockList,
+    ComparisonType,
+    FilterEffectType,
+    FilterRule,
+    FilterType,
+    RuleEffect,
+)
+from pycarlo.features.metadata.metadata_allow_block_list import (
+    MetadataAllowBlockList,
+    MetadataFilter,
+)
+from pycarlo.features.metadata.metadata_filters_container import MetadataFiltersContainer
+
+__all__ = [
+    # Base classes
+    "FilterRule",
+    "BaseAllowBlockList",
+    "FilterEffectType",
+    "RuleEffect",
+    "FilterType",
+    "ComparisonType",
+    # Metadata filtering classes
+    "MetadataFilter",
+    "MetadataAllowBlockList",
+    "MetadataFiltersContainer",
+    # Asset filtering classes
+    "AssetAllowBlockList",
+    "AssetFiltersContainer",
+]

pycarlo/features/metadata/asset_allow_block_list.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from dataclasses_json import DataClassJsonMixin
+
+from pycarlo.common import get_logger
+from pycarlo.features.metadata.base_allow_block_list import BaseAllowBlockList, FilterRule
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class AssetAllowBlockList(BaseAllowBlockList[FilterRule], DataClassJsonMixin):
+    # JSON deserialization fails without this ugly override
+    rules: Optional[List[FilterRule]] = field(default_factory=list)
+
+    asset_type: Optional[str] = None
+
+    def __post_init__(self):
+        # We can't remove the default value because of properties with defaults in the parent class.
+        if not self.asset_type:
+            raise ValueError("asset_type is required")

pycarlo/features/metadata/asset_filters_container.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+from dataclasses_json import DataClassJsonMixin
+
+from .asset_allow_block_list import AssetAllowBlockList
+from .base_allow_block_list import FilterEffectType
+
+# Mapping of resource types to their supported asset types for collection preferences.
+# This is used for validating asset collection preferences.
+# When support for filtering an asset type is implemented in the DC, it should be added here.
+# The reason it is here instead of in Monolith, is so that it can be referenced by the CLI.
+# The pycarlo version in CLI and monolith should be updated after updating this and releasing a
+# new version.
+ASSET_TYPE_ATTRIBUTES = {"tableau": {"project": ["name"], "workbook": ["name", "luid"]}}
+
+
+@dataclass
+class AssetFiltersContainer(DataClassJsonMixin):
+    """
+    Simple container for asset filtering that focuses on in-memory filtering for REST APIs.
+
+    This class provides basic asset filtering functionality without SQL generation complexity.
+    It's designed for the initial phase where assets are collected via REST APIs rather than
+    SQL queries.
+
+    Example usage:
+        # Block all external assets
+        filters = AssetAllowBlockList(
+            filters=[AssetFilter(asset_type="external", effect=FilterEffectType.BLOCK)]
+        )
+        container = AssetFiltersContainer(asset_filters=filters)
+
+        # Check if an asset is blocked
+        is_blocked = container.is_asset_blocked("external", "my_table")  # True
+        is_blocked = container.is_asset_blocked("table", "users")        # False
+    """
+
+    asset_filters: List[AssetAllowBlockList] = field(default_factory=list)
+
+    def is_asset_type_filtered(self, asset_type: str) -> bool:
+        """Returns True if any filters are configured for the given asset type."""
+        return bool(self._get_asset_filters(asset_type))
+
+    def is_asset_blocked(self, asset_type: str, attributes: Dict[str, str]) -> bool:
+        """
+        Returns True if the specified asset is blocked by the current filters.
+
+        Args:
+          asset_type: The type of asset (e.g., 'tableau_workbook_v2', 'jobs', 'power_bi_workspace')
+          attributes: A dictionary representing the attributes of the asset
+
+        Returns:
+          True if the asset is blocked, False if it's allowed
+        """
+        asset_filters = self._get_asset_filters(asset_type)
+
+        is_blocked = False
+
+        for asset_filter in asset_filters:
+            default_effect_matches = asset_filter.get_default_effect_rules(
+                lambda f: f.matches(force_regexp=False, **attributes)
+            )
+            if default_effect_matches:
+                is_blocked = asset_filter.default_effect == FilterEffectType.BLOCK
+            else:
+                other_effect_matches = asset_filter.get_other_effect_rules(
+                    lambda f: f.matches(force_regexp=False, **attributes)
+                )
+                if other_effect_matches:
+                    is_blocked = asset_filter.other_effect == FilterEffectType.BLOCK
+                else:
+                    # No matches, use default effect
+                    is_blocked = asset_filter.default_effect == FilterEffectType.BLOCK
+
+        return is_blocked
+
+    def _get_asset_filters(self, asset_type: str) -> List[AssetAllowBlockList]:
+        return [f for f in self.asset_filters if f.asset_type == asset_type]

pycarlo/features/metadata/base_allow_block_list.py

@@ -0,0 +1,137 @@
+import enum
+import re
+from dataclasses import dataclass, field
+from typing import Any, Callable, Generic, List, Optional, TypeVar
+
+from dataclasses_json import DataClassJsonMixin
+
+from pycarlo.common import get_logger
+
+logger = get_logger(__name__)
+
+# For documentation and samples check the link below:
+# https://www.notion.so/montecarlodata/Catalog-Schema-Filtering-59edd6eff7f74c94ab6bfca75d2e3ff1
+
+
+def _exclude_none_values(value: Any) -> bool:
+    return value is None
+
+
+class FilterEffectType(enum.Enum):
+    BLOCK = "block"
+    ALLOW = "allow"
+
+
+RuleEffect = FilterEffectType
+
+
+class FilterType(enum.Enum):
+    EXACT_MATCH = "exact_match"
+    PREFIX = "prefix"
+    SUFFIX = "suffix"
+    SUBSTRING = "substring"
+    REGEXP = "regexp"
+
+
+ComparisonType = FilterType
+
+# Type variable for the filter class
+FilterRuleT = TypeVar("FilterRuleT", bound="FilterRule")
+
+
+@dataclass
+class RuleCondition(DataClassJsonMixin):
+    attribute_name: str
+    value: str
+    comparison_type: ComparisonType = ComparisonType.EXACT_MATCH
+
+
+@dataclass
+class FilterRule(DataClassJsonMixin):
+    """
+    Base class for all filter types. Provides common filtering logic that can be
+    shared between different filter implementations (e.g., metadata filters, asset filters).
+    """
+
+    conditions: Optional[List[RuleCondition]] = field(default_factory=list)
+    effect: RuleEffect = RuleEffect.BLOCK
+
+    def matches(self, force_regexp: bool = False, **kwargs: Any) -> bool:
+        """
+        Returns True if all properties specified in kwargs match the conditions specified in
+        properties of the same name in this object.
+        If any of the conditions (for example self.field) is None, that condition will be matched.
+        """
+        if not kwargs:
+            raise ValueError("At least one field needs to be specified for matching")
+
+        # kwargs must match the field names in this class, if any of them do not,
+        # invalidate the filter.
+        try:
+            return all(
+                condition.attribute_name not in kwargs
+                or self._match(
+                    condition=condition,
+                    value=kwargs.get(condition.attribute_name),
+                    force_regexp=force_regexp,
+                )
+                for condition in self.conditions or []
+            )
+        except AttributeError:
+            return False
+
+    @classmethod
+    def _match(cls, condition: RuleCondition, value: Optional[str], force_regexp: bool) -> bool:
+        # Field not specified on this object, e.g. self.field=None, which matches everything
+        if value is None:
+            return False
+
+        # The comparison is performed case-insensitive (check BaseFilter._safe_match)
+        # We can use LOWER here since it is part of standard SQL (like AND/OR/NOT), so including it
+        # here is a way to make sure that all comparisons are case-insensitive in the SQL sentences
+        # for all engines. Added option to not always LOWER since customers do have lower/upper case
+        # databases logged in MC
+        filter_value = condition.value.lower()
+        value = value.lower()
+
+        if force_regexp or condition.comparison_type == FilterType.REGEXP:
+            regexp = f"^{filter_value}$"
+            return re.match(regexp, value) is not None
+        elif condition.comparison_type == FilterType.PREFIX:
+            return value.startswith(filter_value)
+        elif condition.comparison_type == FilterType.SUFFIX:
+            return value.endswith(filter_value)
+        elif condition.comparison_type == FilterType.SUBSTRING:
+            return filter_value in value
+        else:  # filter_type == FilterType.EXACT_MATCH
+            return filter_value == value
+
+
+@dataclass
+class BaseAllowBlockList(Generic[FilterRuleT], DataClassJsonMixin):
+    rules: Optional[List[FilterRuleT]] = field(default_factory=list)
+    default_effect: RuleEffect = RuleEffect.ALLOW
+
+    @property
+    def other_effect(self) -> RuleEffect:
+        return RuleEffect.ALLOW if self.default_effect == RuleEffect.BLOCK else RuleEffect.BLOCK
+
+    def get_default_effect_rules(
+        self, condition: Optional[Callable[[FilterRuleT], bool]] = None
+    ) -> List[FilterRuleT]:
+        return list(
+            filter(
+                lambda f: f.effect == self.default_effect and (condition is None or condition(f)),
+                self.rules or [],
+            )
+        )
+
+    def get_other_effect_rules(
+        self, condition: Optional[Callable[[FilterRuleT], bool]] = None
+    ) -> List[FilterRuleT]:
+        return list(
+            filter(
+                lambda f: f.effect != self.default_effect and (condition is None or condition(f)),
+                self.rules or [],
+            )
+        )

pycarlo/features/metadata/metadata_allow_block_list.py

@@ -0,0 +1,94 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from dataclasses_json import config, dataclass_json
+
+from pycarlo.common import get_logger
+from pycarlo.features.metadata.base_allow_block_list import (
+    BaseAllowBlockList,
+    ComparisonType,
+    FilterRule,
+    FilterType,
+    RuleCondition,
+)
+
+logger = get_logger(__name__)
+
+# For documentation and samples check the link below:
+# https://www.notion.so/montecarlodata/Catalog-Schema-Filtering-59edd6eff7f74c94ab6bfca75d2e3ff1
+
+
+@dataclass_json
+@dataclass
+class MetadataFilter(FilterRule):
+    type: FilterType = FilterType.EXACT_MATCH
+
+    # we're using exclude=_exclude_none_values to prevent these properties to be serialized to json
+    # when None, to keep the json doc simpler
+    project: Optional[str] = field(metadata=config(exclude=lambda x: x is None), default=None)
+    dataset: Optional[str] = field(metadata=config(exclude=lambda x: x is None), default=None)
+    table_type: Optional[str] = field(metadata=config(exclude=lambda x: x is None), default=None)
+    table_name: Optional[str] = field(metadata=config(exclude=lambda x: x is None), default=None)
+
+    def __post_init__(self):
+        # For backwards compatibility, we now create a set of conditions based on the
+        # metadata-specific fields.
+        self.conditions = self.conditions or []
+        if self.table_name is not None:
+            is_target_field = self.filter_type_target_field() == "table_name"
+            condition = RuleCondition(
+                comparison_type=self.type if is_target_field else ComparisonType.EXACT_MATCH,
+                attribute_name="table_name",
+                value=self.table_name,
+            )
+            self.conditions.append(condition)
+
+        if self.dataset is not None:
+            is_target_field = self.filter_type_target_field() == "dataset"
+            condition = RuleCondition(
+                comparison_type=self.type if is_target_field else ComparisonType.EXACT_MATCH,
+                attribute_name="dataset",
+                value=self.dataset,
+            )
+            self.conditions.append(condition)
+
+        if self.project is not None:
+            is_target_field = self.filter_type_target_field() == "project"
+            condition = RuleCondition(
+                comparison_type=self.type if is_target_field else ComparisonType.EXACT_MATCH,
+                attribute_name="project",
+                value=self.project,
+            )
+            self.conditions.append(condition)
+
+        if self.table_type is not None:
+            condition = RuleCondition(
+                comparison_type=ComparisonType.EXACT_MATCH,
+                attribute_name="table_type",
+                value=self.table_type,
+            )
+            self.conditions.append(condition)
+
+    def filter_type_target_field(self) -> str:
+        """
+        The field that is evaluated using filter type. Other fields should be
+        compared using exact match.
+        """
+        if self.table_name is not None:
+            return "table_name"
+        if self.dataset is not None:
+            return "dataset"
+        if self.project is not None:
+            return "project"
+
+        logger.exception("Invalid filter, missing target values")
+        return ""
+
+
+@dataclass_json
+@dataclass
+class MetadataAllowBlockList(BaseAllowBlockList[MetadataFilter]):
+    filters: List[MetadataFilter] = field(default_factory=list)
+
+    def __post_init__(self):
+        self.rules = self.filters