acryl-datahub 0.15.0.6rc3__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -0,0 +1,374 @@
+from __future__ import annotations
+
+import abc
+from typing import (
+    Any,
+    List,
+    Sequence,
+    TypedDict,
+    Union,
+)
+
+import pydantic
+
+from datahub.configuration.common import ConfigModel
+from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
+from datahub.ingestion.graph.client import entity_type_to_graphql
+from datahub.ingestion.graph.filters import SearchFilterRule
+from datahub.metadata.schema_classes import EntityTypeName
+from datahub.metadata.urns import DataPlatformUrn, DomainUrn
+
+_AndSearchFilterRule = TypedDict(
+    "_AndSearchFilterRule", {"and": List[SearchFilterRule]}
+)
+_OrFilters = List[_AndSearchFilterRule]
+
+
+class _BaseFilter(ConfigModel):
+    class Config:
+        # We can't wrap this in a TYPE_CHECKING block because the pydantic plugin
+        # doesn't recognize it properly. So unfortunately we'll need to live
+        # with the deprecation warning w/ pydantic v2.
+        allow_population_by_field_name = True
+        if PYDANTIC_VERSION_2:
+            populate_by_name = True
+
+    @abc.abstractmethod
+    def compile(self) -> _OrFilters:
+        pass
+
+
+def _flexible_entity_type_to_graphql(entity_type: str) -> str:
+    if entity_type.upper() == entity_type:
+        # Assume that we were passed a graphql EntityType enum value,
+        # so no conversion is needed.
+        return entity_type
+    return entity_type_to_graphql(entity_type)
+
+
+class _EntityTypeFilter(_BaseFilter):
+    entity_type: List[str] = pydantic.Field(
+        description="The entity type to filter on. Can be 'dataset', 'chart', 'dashboard', 'corpuser', etc.",
+    )
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="_entityType",
+            condition="EQUAL",
+            values=[_flexible_entity_type_to_graphql(t) for t in self.entity_type],
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
+class _EntitySubtypeFilter(_BaseFilter):
+    entity_type: str
+    entity_subtype: str = pydantic.Field(
+        description="The entity subtype to filter on. Can be 'Table', 'View', 'Source', etc. depending on the native platform's concepts.",
+    )
+
+    def compile(self) -> _OrFilters:
+        rules = [
+            SearchFilterRule(
+                field="_entityType",
+                condition="EQUAL",
+                values=[_flexible_entity_type_to_graphql(self.entity_type)],
+            ),
+            SearchFilterRule(
+                field="typeNames",
+                condition="EQUAL",
+                values=[self.entity_subtype],
+            ),
+        ]
+        return [{"and": rules}]
+
+
+class _PlatformFilter(_BaseFilter):
+    platform: List[str]
+    # TODO: Add validator to convert string -> list of strings
+
+    @pydantic.validator("platform", each_item=True)
+    def validate_platform(cls, v: str) -> str:
+        # Subtle - we use the constructor instead of the from_string method
+        # because coercion is acceptable here.
+        return str(DataPlatformUrn(v))
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="platform.keyword",
+            condition="EQUAL",
+            values=self.platform,
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
+class _DomainFilter(_BaseFilter):
+    domain: List[str]
+
+    @pydantic.validator("domain", each_item=True)
+    def validate_domain(cls, v: str) -> str:
+        return str(DomainUrn.from_string(v))
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="domains",
+            condition="EQUAL",
+            values=self.domain,
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
+class _EnvFilter(_BaseFilter):
+    # Note that not all entity types have an env (e.g. dashboards / charts).
+    # If the env filter is specified, these will be excluded.
+    env: List[str]
+
+    def compile(self) -> _OrFilters:
+        return [
+            # For most entity types, we look at the origin field.
+            {
+                "and": [
+                    SearchFilterRule(
+                        field="origin",
+                        condition="EQUAL",
+                        values=self.env,
+                    ),
+                ]
+            },
+            # For containers, we now have an "env" property as of
+            # https://github.com/datahub-project/datahub/pull/11214
+            # Prior to this, we put "env" in the customProperties. But we're
+            # not bothering with that here.
+            {
+                "and": [
+                    SearchFilterRule(
+                        field="env",
+                        condition="EQUAL",
+                        values=self.env,
+                    ),
+                ]
+            },
+        ]
+
+
+class _CustomCondition(_BaseFilter):
+    """Represents a single field condition"""
+
+    field: str
+    condition: str
+    values: List[str]
+
+    def compile(self) -> _OrFilters:
+        rule = SearchFilterRule(
+            field=self.field,
+            condition=self.condition,
+            values=self.values,
+        )
+        return [{"and": [rule]}]
+
+
+class _And(_BaseFilter):
+    """Represents an AND conjunction of filters"""
+
+    and_: Sequence["Filter"] = pydantic.Field(alias="and")
+    # TODO: Add validator to ensure that the "and" field is not empty
+
+    def compile(self) -> _OrFilters:
+        # The "and" operator must be implemented by doing a Cartesian product
+        # of the OR clauses.
+        # Example 1:
+        # (A or B) and (C or D) ->
+        # (A and C) or (A and D) or (B and C) or (B and D)
+        # Example 2:
+        # (A or B) and (C or D) and (E or F) ->
+        # (A and C and E) or (A and C and F) or (A and D and E) or (A and D and F) or
+        # (B and C and E) or (B and C and F) or (B and D and E) or (B and D and F)
+
+        # Start with the first filter's OR clauses
+        result = self.and_[0].compile()
+
+        # For each subsequent filter
+        for filter in self.and_[1:]:
+            new_result = []
+            # Get its OR clauses
+            other_clauses = filter.compile()
+
+            # Create Cartesian product
+            for existing_clause in result:
+                for other_clause in other_clauses:
+                    # Merge the AND conditions from both clauses
+                    new_result.append(self._merge_ands(existing_clause, other_clause))
+
+            result = new_result
+
+        return result
+
+    @classmethod
+    def _merge_ands(
+        cls, a: _AndSearchFilterRule, b: _AndSearchFilterRule
+    ) -> _AndSearchFilterRule:
+        return {
+            "and": [
+                *a["and"],
+                *b["and"],
+            ]
+        }
+
+
+class _Or(_BaseFilter):
+    """Represents an OR conjunction of filters"""
+
+    or_: Sequence["Filter"] = pydantic.Field(alias="or")
+    # TODO: Add validator to ensure that the "or" field is not empty
+
+    def compile(self) -> _OrFilters:
+        merged_filter = []
+        for filter in self.or_:
+            merged_filter.extend(filter.compile())
+        return merged_filter
+
+
+class _Not(_BaseFilter):
+    """Represents a NOT filter"""
+
+    not_: "Filter" = pydantic.Field(alias="not")
+
+    @pydantic.validator("not_", pre=False)
+    def validate_not(cls, v: "Filter") -> "Filter":
+        inner_filter = v.compile()
+        if len(inner_filter) != 1:
+            raise ValueError(
+                "Cannot negate a filter with multiple OR clauses [not yet supported]"
+            )
+        return v
+
+    def compile(self) -> _OrFilters:
+        # TODO: Eventually we'll want to implement a full DNF normalizer.
+        # https://en.wikipedia.org/wiki/Disjunctive_normal_form#Conversion_to_DNF
+
+        inner_filter = self.not_.compile()
+        assert len(inner_filter) == 1  # validated above
+
+        # ¬(A and B) -> (¬A) OR (¬B)
+        and_filters = inner_filter[0]["and"]
+        final_filters: _OrFilters = []
+        for rule in and_filters:
+            final_filters.append({"and": [rule.negate()]})
+
+        return final_filters
+
+
+# TODO: With pydantic 2, we can use a RootModel with a
+# discriminated union to make the error messages more informative.
+Filter = Union[
+    _And,
+    _Or,
+    _Not,
+    _EntityTypeFilter,
+    _EntitySubtypeFilter,
+    _PlatformFilter,
+    _DomainFilter,
+    _EnvFilter,
+    _CustomCondition,
+]
+
+
+# Required to resolve forward references to "Filter"
+if PYDANTIC_VERSION_2:
+    _And.model_rebuild()  # type: ignore
+    _Or.model_rebuild()  # type: ignore
+    _Not.model_rebuild()  # type: ignore
+else:
+    _And.update_forward_refs()
+    _Or.update_forward_refs()
+    _Not.update_forward_refs()
+
+
+def load_filters(obj: Any) -> Filter:
+    if PYDANTIC_VERSION_2:
+        return pydantic.TypeAdapter(Filter).validate_python(obj)  # type: ignore
+    else:
+        return pydantic.parse_obj_as(Filter, obj)  # type: ignore
+
+
+# We need FilterDsl for two reasons:
+# 1. To provide wrapper methods around lots of filters while avoid bloating the
+#    yaml spec.
+# 2. Pydantic models in general don't support positional arguments, making the
+#    calls feel repetitive (e.g. Platform(platform=...)).
+#    See https://github.com/pydantic/pydantic/issues/6792
+#    We also considered using dataclasses / pydantic dataclasses, but
+#    ultimately decided that they didn't quite suit our requirements,
+#    particularly with regards to the field aliases for and/or/not.
+class FilterDsl:
+    @staticmethod
+    def and_(*args: "Filter") -> _And:
+        return _And(and_=list(args))
+
+    @staticmethod
+    def or_(*args: "Filter") -> _Or:
+        return _Or(or_=list(args))
+
+    @staticmethod
+    def not_(arg: "Filter") -> _Not:
+        return _Not(not_=arg)
+
+    @staticmethod
+    def entity_type(
+        entity_type: Union[EntityTypeName, Sequence[EntityTypeName]],
+    ) -> _EntityTypeFilter:
+        return _EntityTypeFilter(
+            entity_type=(
+                [entity_type] if isinstance(entity_type, str) else list(entity_type)
+            )
+        )
+
+    @staticmethod
+    def entity_subtype(entity_type: str, subtype: str) -> _EntitySubtypeFilter:
+        return _EntitySubtypeFilter(
+            entity_type=entity_type,
+            entity_subtype=subtype,
+        )
+
+    @staticmethod
+    def platform(platform: Union[str, List[str]], /) -> _PlatformFilter:
+        return _PlatformFilter(
+            platform=[platform] if isinstance(platform, str) else platform
+        )
+
+    # TODO: Add a platform_instance filter
+
+    @staticmethod
+    def domain(domain: Union[str, List[str]], /) -> _DomainFilter:
+        return _DomainFilter(domain=[domain] if isinstance(domain, str) else domain)
+
+    @staticmethod
+    def env(env: Union[str, List[str]], /) -> _EnvFilter:
+        return _EnvFilter(env=[env] if isinstance(env, str) else env)
+
+    @staticmethod
+    def has_custom_property(key: str, value: str) -> _CustomCondition:
+        return _CustomCondition(
+            field="customProperties",
+            condition="EQUAL",
+            values=[f"{key}={value}"],
+        )
+
+    # TODO: Add a soft-deletion status filter
+    # TODO: add a container / browse path filter
+    # TODO add shortcut for custom filters
+
+    @staticmethod
+    def custom_filter(
+        field: str, condition: str, values: List[str]
+    ) -> _CustomCondition:
+        return _CustomCondition(
+            field=field,
+            condition=condition,
+            values=values,
+        )

datahub/specific/dataset.py

@@ -15,6 +15,7 @@ from datahub.metadata.schema_classes import (
     UpstreamClass as Upstream,
     UpstreamLineageClass as UpstreamLineage,
 )
+from datahub.metadata.urns import DatasetUrn, TagUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.structured_properties import (
@@ -22,8 +23,6 @@ from datahub.specific.aspect_helpers.structured_properties import (
 )
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
 from datahub.specific.aspect_helpers.terms import HasTermsPatch
-from datahub.utilities.urns.tag_urn import TagUrn
-from datahub.utilities.urns.urn import Urn
 
 _Parent = TypeVar("_Parent", bound=MetadataPatchProposal)
 
@@ -104,12 +103,12 @@ class DatasetPatchBuilder(
 ):
     def __init__(
         self,
-        urn: str,
+        urn: Union[str, DatasetUrn],
         system_metadata: Optional[SystemMetadataClass] = None,
         audit_header: Optional[KafkaAuditHeaderClass] = None,
     ) -> None:
         super().__init__(
-            urn, system_metadata=system_metadata, audit_header=audit_header
+            str(urn), system_metadata=system_metadata, audit_header=audit_header
         )
 
     @classmethod

datahub/sql_parsing/_sqlglot_patch.py

@@ -172,17 +172,9 @@ def _patch_lineage() -> None:
          derived_tables = [
              source.expression.parent
              for source in scope.sources.values()
-@@ -254,6 +257,7 @@ def to_node(
-         if dt.comments and dt.comments[0].startswith("source: ")
-     }
-
-+    c: exp.Column
-     for c in source_columns:
-         table = c.table
-         source = scope.sources.get(table)
 @@ -281,8 +285,21 @@ def to_node(
-             # it means this column's lineage is unknown. This can happen if the definition of a source used in a query
-             # is not passed into the `sources` map.
+             # is unknown. This can happen if the definition of a source used in a query is not
+             # passed into the `sources` map.
              source = source or exp.Placeholder()
 +
 +            subfields = []

datahub/sql_parsing/schema_resolver.py

@@ -13,7 +13,7 @@ from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier
 from datahub.metadata.schema_classes import SchemaFieldClass, SchemaMetadataClass
 from datahub.metadata.urns import DataPlatformUrn
-from datahub.sql_parsing._models import _TableName as _TableName  # noqa: I250
+from datahub.sql_parsing._models import _TableName as _TableName
 from datahub.sql_parsing.sql_parsing_common import PLATFORMS_WITH_CASE_SENSITIVE_TABLES
 from datahub.utilities.file_backed_collections import ConnectionWrapper, FileBackedDict
 from datahub.utilities.urns.field_paths import get_simple_field_path_from_v2_field_path

datahub/sql_parsing/split_statements.py

@@ -8,11 +8,11 @@ END_KEYWORD = "END"
 
 CONTROL_FLOW_KEYWORDS = [
     "GO",
-    r"BEGIN\w+TRY",
-    r"BEGIN\w+CATCH",
+    r"BEGIN\s+TRY",
+    r"BEGIN\s+CATCH",
     "BEGIN",
-    r"END\w+TRY",
-    r"END\w+CATCH",
+    r"END\s+TRY",
+    r"END\s+CATCH",
     # This isn't strictly correct, but we assume that IF | (condition) | (block) should all be split up
     # This mainly ensures that IF statements don't get tacked onto the previous statement incorrectly
     "IF",
@@ -73,25 +73,31 @@ class _StatementSplitter:
         # what a given END is closing.
         self.current_case_statements = 0
 
-    def _is_keyword_at_position(self, pos: int, keyword: str) -> bool:
+    def _is_keyword_at_position(self, pos: int, keyword: str) -> Tuple[bool, str]:
         """
         Check if a keyword exists at the given position using regex word boundaries.
         """
         sql = self.sql
 
-        if pos + len(keyword) > len(sql):
-            return False
+        keyword_length = len(keyword.replace(r"\s+", " "))
+
+        if pos + keyword_length > len(sql):
+            return False, ""
 
         # If we're not at a word boundary, we can't generate a keyword.
         if pos > 0 and not (
             bool(re.match(r"\w\W", sql[pos - 1 : pos + 1]))
             or bool(re.match(r"\W\w", sql[pos - 1 : pos + 1]))
         ):
-            return False
+            return False, ""
 
-        pattern = rf"^{re.escape(keyword)}\b"
+        pattern = rf"^{keyword}\b"
         match = re.match(pattern, sql[pos:], re.IGNORECASE)
-        return bool(match)
+        is_match = bool(match)
+        actual_match = (
+            sql[pos:][match.start() : match.end()] if match is not None else ""
+        )
+        return is_match, actual_match
 
     def _look_ahead_for_keywords(self, keywords: List[str]) -> Tuple[bool, str, int]:
         """
@@ -99,7 +105,8 @@ class _StatementSplitter:
         """
 
         for keyword in keywords:
-            if self._is_keyword_at_position(self.i, keyword):
+            is_match, keyword = self._is_keyword_at_position(self.i, keyword)
+            if is_match:
                 return True, keyword, len(keyword)
         return False, "", 0
 
@@ -118,7 +125,7 @@ class _StatementSplitter:
 
     def process(self) -> Iterator[str]:
         if not self.sql or not self.sql.strip():
-            return
+            yield from ()
 
         prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
         while self.i < len(self.sql):
@@ -181,7 +188,7 @@ class _StatementSplitter:
     def _process_normal(self, most_recent_real_char: str) -> Iterator[str]:
         c = self.sql[self.i]
 
-        if self._is_keyword_at_position(self.i, CASE_KEYWORD):
+        if self._is_keyword_at_position(self.i, CASE_KEYWORD)[0]:
             self.current_case_statements += 1
 
         is_control_keyword, keyword, keyword_len = self._look_ahead_for_keywords(

datahub/sql_parsing/sql_parsing_common.py

@@ -24,12 +24,19 @@ DIALECTS_WITH_CASE_INSENSITIVE_COLS = {
     # For SQL server, the default collation rules mean that all identifiers (schema, table, column names)
     # are case preserving but case insensitive.
     "mssql",
+    # Oracle automatically converts unquoted identifiers to uppercase.
+    # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/Database-Object-Names-and-Qualifiers.html#GUID-3C59E44A-5140-4BCA-B9E1-3039C8050C49
+    # In our Oracle connector, we then normalize column names to lowercase. This behavior
+    # actually comes from the underlying Oracle sqlalchemy dialect.
+    # https://github.com/sqlalchemy/sqlalchemy/blob/d9b4d8ff3aae504402d324f3ebf0b8faff78f5dc/lib/sqlalchemy/dialects/oracle/base.py#L2579
+    "oracle",
 }
 DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # In some dialects, column identifiers are effectively case insensitive
     # because they are automatically converted to uppercase. Most other systems
     # automatically lowercase unquoted identifiers.
     "snowflake",
+    "oracle",
 }
 assert DIALECTS_WITH_DEFAULT_UPPERCASE_COLS.issubset(
     DIALECTS_WITH_CASE_INSENSITIVE_COLS

datahub/sql_parsing/sqlglot_lineage.py

@@ -473,7 +473,7 @@ def _create_table_ddl_cll(
     return column_lineage
 
 
-def _select_statement_cll(  # noqa: C901
+def _select_statement_cll(
     statement: _SupportedColumnLineageTypes,
     dialect: sqlglot.Dialect,
     root_scope: sqlglot.optimizer.Scope,

datahub/sql_parsing/sqlglot_utils.py

@@ -56,10 +56,7 @@ def get_dialect(platform: DialectOrStr) -> sqlglot.Dialect:
 def is_dialect_instance(
     dialect: sqlglot.Dialect, platforms: Union[str, Iterable[str]]
 ) -> bool:
-    if isinstance(platforms, str):
-        platforms = [platforms]
-    else:
-        platforms = list(platforms)
+    platforms = [platforms] if isinstance(platforms, str) else list(platforms)
 
     dialects = [get_dialect(platform) for platform in platforms]
 

datahub/testing/check_sql_parser_result.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import pathlib
 from typing import Any, Dict, Optional
 
@@ -8,11 +7,10 @@ import deepdiff
 from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier
 from datahub.sql_parsing.schema_resolver import SchemaInfo, SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult, sqlglot_lineage
+from datahub.testing.pytest_hooks import get_golden_settings
 
 logger = logging.getLogger(__name__)
 
-UPDATE_FILES = os.environ.get("UPDATE_SQLPARSER_FILES", "false").lower() == "true"
-
 
 def assert_sql_result_with_resolver(
     sql: str,
@@ -22,6 +20,8 @@ def assert_sql_result_with_resolver(
     allow_table_error: bool = False,
     **kwargs: Any,
 ) -> None:
+    settings = get_golden_settings()
+
     # HACK: Our BigQuery source overwrites this value and doesn't undo it.
     # As such, we need to handle that here.
     BigqueryTableIdentifier._BQ_SHARDED_TABLE_SUFFIX = "_yyyymmdd"
@@ -47,15 +47,14 @@ def assert_sql_result_with_resolver(
         )
 
     txt = res.json(indent=4)
-    if UPDATE_FILES:
+    if settings.update_golden:
         expected_file.write_text(txt)
         return
 
     if not expected_file.exists():
         expected_file.write_text(txt)
         raise AssertionError(
-            f"Expected file {expected_file} does not exist. "
-            "Created it with the expected output. Please verify it."
+            f"Missing expected golden file; run with --update-golden-files to create it: {expected_file}"
         )
 
     expected = SqlParsingResult.parse_raw(expected_file.read_text())

datahub/testing/compare_metadata_json.py

@@ -16,6 +16,7 @@ from deepdiff import DeepDiff
 from datahub.ingestion.sink.file import write_metadata_file
 from datahub.ingestion.source.file import read_metadata_file
 from datahub.testing.mcp_diff import CannotCompareMCPs, MCPDiff, get_aspects_by_urn
+from datahub.testing.pytest_hooks import get_golden_settings
 
 logger = logging.getLogger(__name__)
 
@@ -40,26 +41,26 @@ def load_json_file(filename: Union[str, os.PathLike]) -> MetadataJson:
 def assert_metadata_files_equal(
     output_path: Union[str, os.PathLike],
     golden_path: Union[str, os.PathLike],
-    update_golden: bool,
-    copy_output: bool,
     ignore_paths: Sequence[str] = (),
     ignore_paths_v2: Sequence[str] = (),
     ignore_order: bool = True,
 ) -> None:
+    settings = get_golden_settings()
+
     golden_exists = os.path.isfile(golden_path)
 
-    if copy_output:
+    if settings.copy_output:
         shutil.copyfile(str(output_path), str(golden_path) + ".output")
         logger.info(f"Copied output file to {golden_path}.output")
 
-    if not update_golden and not golden_exists:
+    if not settings.update_golden and not golden_exists:
         raise FileNotFoundError(
             "Golden file does not exist. Please run with the --update-golden-files option to create."
         )
 
     output = load_json_file(output_path)
 
-    if update_golden and not golden_exists:
+    if settings.update_golden and not golden_exists:
         shutil.copyfile(str(output_path), str(golden_path))
         return
     else:
@@ -87,7 +88,7 @@ def assert_metadata_files_equal(
     ignore_paths = (*ignore_paths, *default_exclude_paths)
 
     diff = diff_metadata_json(output, golden, ignore_paths, ignore_order=ignore_order)
-    if diff and update_golden:
+    if diff and settings.update_golden:
         if isinstance(diff, MCPDiff) and diff.is_delta_valid:
             logger.info(f"Applying delta to golden file {golden_path}")
             diff.apply_delta(golden)