acryl-datahub 0.14.1.13rc9__py3-none-any.whl → 0.15.0__py3-none-any.whl

datahub/metadata/schemas/FormInfo.avsc

@@ -97,6 +97,10 @@
                   "namespace": "com.linkedin.pegasus2avro.form",
                   "fields": [
                     {
+                      "Searchable": {
+                        "fieldName": "structuredPropertyPromptUrns",
+                        "fieldType": "URN"
+                      },
                       "java": {
                         "class": "com.linkedin.pegasus2avro.common.urn.Urn"
                       },

datahub/metadata/schemas/StructuredPropertyDefinition.avsc

@@ -359,7 +359,7 @@
       ],
       "name": "lastModified",
       "default": null,
-      "doc": "Created Audit stamp"
+      "doc": "Last Modified Audit stamp"
     }
   ]
 }

datahub/metadata/schemas/StructuredPropertyKey.avsc

@@ -6,6 +6,7 @@
     "entityCategory": "core",
     "entityAspects": [
       "propertyDefinition",
+      "structuredPropertySettings",
       "institutionalMemory",
       "status"
     ],

datahub/metadata/schemas/StructuredPropertySettings.avsc

@@ -0,0 +1,114 @@
+{
+  "type": "record",
+  "Aspect": {
+    "name": "structuredPropertySettings"
+  },
+  "name": "StructuredPropertySettings",
+  "namespace": "com.linkedin.pegasus2avro.structured",
+  "fields": [
+    {
+      "Searchable": {
+        "fieldType": "BOOLEAN"
+      },
+      "type": "boolean",
+      "name": "isHidden",
+      "default": false,
+      "doc": "Whether or not this asset should be hidden in the main application"
+    },
+    {
+      "Searchable": {
+        "fieldType": "BOOLEAN"
+      },
+      "type": "boolean",
+      "name": "showInSearchFilters",
+      "default": false,
+      "doc": "Whether or not this asset should be displayed as a search filter"
+    },
+    {
+      "Searchable": {
+        "fieldType": "BOOLEAN"
+      },
+      "type": "boolean",
+      "name": "showInAssetSummary",
+      "default": false,
+      "doc": "Whether or not this asset should be displayed in the asset sidebar"
+    },
+    {
+      "Searchable": {
+        "fieldType": "BOOLEAN"
+      },
+      "type": "boolean",
+      "name": "showAsAssetBadge",
+      "default": false,
+      "doc": "Whether or not this asset should be displayed as an asset badge on other\nasset's headers"
+    },
+    {
+      "Searchable": {
+        "fieldType": "BOOLEAN"
+      },
+      "type": "boolean",
+      "name": "showInColumnsTable",
+      "default": false,
+      "doc": "Whether or not this asset should be displayed as a column in the schema field table\nin a Dataset's \"Columns\" tab."
+    },
+    {
+      "Searchable": {
+        "/time": {
+          "fieldName": "lastModifiedSettings",
+          "fieldType": "DATETIME"
+        }
+      },
+      "type": [
+        "null",
+        {
+          "type": "record",
+          "name": "AuditStamp",
+          "namespace": "com.linkedin.pegasus2avro.common",
+          "fields": [
+            {
+              "type": "long",
+              "name": "time",
+              "doc": "When did the resource/association/sub-resource move into the specific lifecycle stage represented by this AuditEvent."
+            },
+            {
+              "java": {
+                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+              },
+              "type": "string",
+              "name": "actor",
+              "doc": "The entity (e.g. a member URN) which will be credited for moving the resource/association/sub-resource into the specific lifecycle stage. It is also the one used to authorize the change.",
+              "Urn": "Urn"
+            },
+            {
+              "java": {
+                "class": "com.linkedin.pegasus2avro.common.urn.Urn"
+              },
+              "type": [
+                "null",
+                "string"
+              ],
+              "name": "impersonator",
+              "default": null,
+              "doc": "The entity (e.g. a service URN) which performs the change on behalf of the Actor and must be authorized to act as the Actor.",
+              "Urn": "Urn"
+            },
+            {
+              "type": [
+                "null",
+                "string"
+              ],
+              "name": "message",
+              "default": null,
+              "doc": "Additional context around how DataHub was informed of the particular change. For example: was the change created by an automated process, or manually."
+            }
+          ],
+          "doc": "Data captured on a resource/association/sub-resource level giving insight into when that resource/association/sub-resource moved into a particular lifecycle stage, and who acted to move it into that specific lifecycle stage."
+        }
+      ],
+      "name": "lastModified",
+      "default": null,
+      "doc": "Last Modified Audit stamp"
+    }
+  ],
+  "doc": "Settings specific to a structured property entity"
+}

datahub/specific/chart.py

@@ -1,10 +1,8 @@
-import time
 from typing import Dict, List, Optional, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
 from datahub.metadata.schema_classes import (
     AccessLevelClass,
-    AuditStampClass,
     ChangeAuditStampsClass,
     ChartInfoClass as ChartInfo,
     ChartTypeClass,
@@ -47,43 +45,6 @@ class ChartPatchBuilder(MetadataPatchProposal):
         )
         self.ownership_patch_helper = OwnershipPatchHelper(self)
 
-    def _mint_auditstamp(self, message: Optional[str] = None) -> AuditStampClass:
-        """
-        Creates an AuditStampClass instance with the current timestamp and other default values.
-
-        Args:
-            message: The message associated with the audit stamp (optional).
-
-        Returns:
-            An instance of AuditStampClass.
-        """
-        return AuditStampClass(
-            time=int(time.time() * 1000.0),
-            actor="urn:li:corpuser:datahub",
-            message=message,
-        )
-
-    def _ensure_urn_type(
-        self, entity_type: str, edges: List[Edge], context: str
-    ) -> None:
-        """
-        Ensures that the destination URNs in the given edges have the specified entity type.
-
-        Args:
-            entity_type: The entity type to check against.
-            edges: A list of Edge objects.
-            context: The context or description of the operation.
-
-        Raises:
-            ValueError: If any of the destination URNs is not of the specified entity type.
-        """
-        for e in edges:
-            urn = Urn.create_from_string(e.destinationUrn)
-            if not urn.get_type() == entity_type:
-                raise ValueError(
-                    f"{context}: {e.destinationUrn} is not of type {entity_type}"
-                )
-
     def add_owner(self, owner: Owner) -> "ChartPatchBuilder":
         """
         Adds an owner to the ChartPatchBuilder.

datahub/specific/dashboard.py

@@ -1,10 +1,8 @@
-import time
 from typing import Dict, List, Optional, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
 from datahub.metadata.schema_classes import (
     AccessLevelClass,
-    AuditStampClass,
     ChangeAuditStampsClass,
     DashboardInfoClass as DashboardInfo,
     EdgeClass as Edge,
@@ -46,43 +44,6 @@ class DashboardPatchBuilder(MetadataPatchProposal):
         )
         self.ownership_patch_helper = OwnershipPatchHelper(self)
 
-    def _mint_auditstamp(self, message: Optional[str] = None) -> AuditStampClass:
-        """
-        Creates an AuditStampClass instance with the current timestamp and other default values.
-
-        Args:
-            message: The message associated with the audit stamp (optional).
-
-        Returns:
-            An instance of AuditStampClass.
-        """
-        return AuditStampClass(
-            time=int(time.time() * 1000.0),
-            actor="urn:li:corpuser:datahub",
-            message=message,
-        )
-
-    def _ensure_urn_type(
-        self, entity_type: str, edges: List[Edge], context: str
-    ) -> None:
-        """
-        Ensures that the destination URNs in the given edges have the specified entity type.
-
-        Args:
-            entity_type: The entity type to check against.
-            edges: A list of Edge objects.
-            context: The context or description of the operation.
-
-        Raises:
-            ValueError: If any of the destination URNs is not of the specified entity type.
-        """
-        for e in edges:
-            urn = Urn.create_from_string(e.destinationUrn)
-            if not urn.get_type() == entity_type:
-                raise ValueError(
-                    f"{context}: {e.destinationUrn} is not of type {entity_type}"
-                )
-
     def add_owner(self, owner: Owner) -> "DashboardPatchBuilder":
         """
         Adds an owner to the DashboardPatchBuilder.

datahub/specific/datajob.py

@@ -1,9 +1,7 @@
-import time
 from typing import Dict, List, Optional, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
 from datahub.metadata.schema_classes import (
-    AuditStampClass,
     DataJobInfoClass as DataJobInfo,
     DataJobInputOutputClass as DataJobInputOutput,
     EdgeClass as Edge,
@@ -16,10 +14,9 @@ from datahub.metadata.schema_classes import (
     SystemMetadataClass,
     TagAssociationClass as Tag,
 )
+from datahub.metadata.urns import SchemaFieldUrn, TagUrn, Urn
 from datahub.specific.custom_properties import CustomPropertiesPatchHelper
 from datahub.specific.ownership import OwnershipPatchHelper
-from datahub.utilities.urns.tag_urn import TagUrn
-from datahub.utilities.urns.urn import Urn
 
 
 class DataJobPatchBuilder(MetadataPatchProposal):
@@ -45,43 +42,6 @@ class DataJobPatchBuilder(MetadataPatchProposal):
         )
         self.ownership_patch_helper = OwnershipPatchHelper(self)
 
-    def _mint_auditstamp(self, message: Optional[str] = None) -> AuditStampClass:
-        """
-        Creates an AuditStampClass instance with the current timestamp and other default values.
-
-        Args:
-            message: The message associated with the audit stamp (optional).
-
-        Returns:
-            An instance of AuditStampClass.
-        """
-        return AuditStampClass(
-            time=int(time.time() * 1000.0),
-            actor="urn:li:corpuser:datahub",
-            message=message,
-        )
-
-    def _ensure_urn_type(
-        self, entity_type: str, edges: List[Edge], context: str
-    ) -> None:
-        """
-        Ensures that the destination URNs in the given edges have the specified entity type.
-
-        Args:
-            entity_type: The entity type to check against.
-            edges: A list of Edge objects.
-            context: The context or description of the operation.
-
-        Raises:
-            ValueError: If any of the destination URNs is not of the specified entity type.
-        """
-        for e in edges:
-            urn = Urn.create_from_string(e.destinationUrn)
-            if not urn.get_type() == entity_type:
-                raise ValueError(
-                    f"{context}: {e.destinationUrn} is not of type {entity_type}"
-                )
-
     def add_owner(self, owner: Owner) -> "DataJobPatchBuilder":
         """
         Adds an owner to the DataJobPatchBuilder.
@@ -142,7 +102,7 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
         Notes:
             If `input` is an Edge object, it is used directly. If `input` is a Urn object or string,
-            it is converted to an Edge object and added with default audit stamps.
+            it is converted to an Edge object and added without any audit stamps.
         """
         if isinstance(input, Edge):
             input_urn: str = input.destinationUrn
@@ -154,8 +114,6 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
             input_edge = Edge(
                 destinationUrn=input_urn,
-                created=self._mint_auditstamp(),
-                lastModified=self._mint_auditstamp(),
             )
 
         self._ensure_urn_type("dataJob", [input_edge], "add_input_datajob")
@@ -225,7 +183,7 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
         Notes:
             If `input` is an Edge object, it is used directly. If `input` is a Urn object or string,
-            it is converted to an Edge object and added with default audit stamps.
+            it is converted to an Edge object and added without any audit stamps.
         """
         if isinstance(input, Edge):
             input_urn: str = input.destinationUrn
@@ -237,8 +195,6 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
             input_edge = Edge(
                 destinationUrn=input_urn,
-                created=self._mint_auditstamp(),
-                lastModified=self._mint_auditstamp(),
             )
 
         self._ensure_urn_type("dataset", [input_edge], "add_input_dataset")
@@ -310,7 +266,7 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
         Notes:
             If `output` is an Edge object, it is used directly. If `output` is a Urn object or string,
-            it is converted to an Edge object and added with default audit stamps.
+            it is converted to an Edge object and added without any audit stamps.
         """
         if isinstance(output, Edge):
             output_urn: str = output.destinationUrn
@@ -322,15 +278,13 @@ class DataJobPatchBuilder(MetadataPatchProposal):
 
             output_edge = Edge(
                 destinationUrn=output_urn,
-                created=self._mint_auditstamp(),
-                lastModified=self._mint_auditstamp(),
             )
 
         self._ensure_urn_type("dataset", [output_edge], "add_output_dataset")
         self._add_patch(
             DataJobInputOutput.ASPECT_NAME,
             "add",
-            path=f"/outputDatasetEdges/{self.quote(str(output))}",
+            path=f"/outputDatasetEdges/{self.quote(output_urn)}",
             value=output_edge,
         )
         return self
@@ -392,9 +346,7 @@ class DataJobPatchBuilder(MetadataPatchProposal):
             ValueError: If the input is not a Schema Field urn.
         """
         input_urn = str(input)
-        urn = Urn.create_from_string(input_urn)
-        if not urn.get_type() == "schemaField":
-            raise ValueError(f"Input {input} is not a Schema Field urn")
+        assert SchemaFieldUrn.from_string(input_urn)
 
         self._add_patch(
             DataJobInputOutput.ASPECT_NAME,
@@ -466,9 +418,7 @@ class DataJobPatchBuilder(MetadataPatchProposal):
             ValueError: If the output is not a Schema Field urn.
         """
         output_urn = str(output)
-        urn = Urn.create_from_string(output_urn)
-        if not urn.get_type() == "schemaField":
-            raise ValueError(f"Input {output} is not a Schema Field urn")
+        assert SchemaFieldUrn.from_string(output_urn)
 
         self._add_patch(
             DataJobInputOutput.ASPECT_NAME,

datahub/sql_parsing/schema_resolver.py

@@ -123,6 +123,13 @@ class SchemaResolver(Closeable, SchemaResolverInterface):
         )
         return urn
 
+    def resolve_urn(self, urn: str) -> Tuple[str, Optional[SchemaInfo]]:
+        schema_info = self._resolve_schema_info(urn)
+        if schema_info:
+            return urn, schema_info
+
+        return urn, None
+
     def resolve_table(self, table: _TableName) -> Tuple[str, Optional[SchemaInfo]]:
         urn = self.get_urn_for_table(table)
 
@@ -293,3 +300,19 @@ def _convert_schema_field_list_to_info(
 
 def _convert_schema_aspect_to_info(schema_metadata: SchemaMetadataClass) -> SchemaInfo:
     return _convert_schema_field_list_to_info(schema_metadata.fields)
+
+
+def match_columns_to_schema(
+    schema_info: SchemaInfo, input_columns: List[str]
+) -> List[str]:
+    column_from_gms: List[str] = list(schema_info.keys())  # list() to silent lint
+
+    gms_column_map: Dict[str, str] = {
+        column.lower(): column for column in column_from_gms
+    }
+
+    output_columns: List[str] = [
+        gms_column_map.get(column.lower(), column) for column in input_columns
+    ]
+
+    return output_columns

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -1383,8 +1383,7 @@ class SqlParsingAggregator(Closeable):
         return QueryUrn(query_id).urn()
 
     @classmethod
-    def _composite_query_id(cls, composed_of_queries: Iterable[QueryId]) -> str:
-        composed_of_queries = list(composed_of_queries)
+    def _composite_query_id(cls, composed_of_queries: List[QueryId]) -> str:
         combined = json.dumps(composed_of_queries)
         return f"composite_{generate_hash(combined)}"
 

datahub/sql_parsing/sqlglot_lineage.py

@@ -1181,6 +1181,45 @@ def sqlglot_lineage(
         )
 
 
+@functools.lru_cache(maxsize=128)
+def create_and_cache_schema_resolver(
+    platform: str,
+    env: str,
+    graph: Optional[DataHubGraph] = None,
+    platform_instance: Optional[str] = None,
+    schema_aware: bool = True,
+) -> SchemaResolver:
+    return create_schema_resolver(
+        platform=platform,
+        env=env,
+        graph=graph,
+        platform_instance=platform_instance,
+        schema_aware=schema_aware,
+    )
+
+
+def create_schema_resolver(
+    platform: str,
+    env: str,
+    graph: Optional[DataHubGraph] = None,
+    platform_instance: Optional[str] = None,
+    schema_aware: bool = True,
+) -> SchemaResolver:
+    if graph and schema_aware:
+        return graph._make_schema_resolver(
+            platform=platform,
+            platform_instance=platform_instance,
+            env=env,
+        )
+
+    return SchemaResolver(
+        platform=platform,
+        platform_instance=platform_instance,
+        env=env,
+        graph=None,
+    )
+
+
 def create_lineage_sql_parsed_result(
     query: str,
     default_db: Optional[str],
@@ -1191,21 +1230,17 @@ def create_lineage_sql_parsed_result(
     graph: Optional[DataHubGraph] = None,
     schema_aware: bool = True,
 ) -> SqlParsingResult:
+    schema_resolver = create_schema_resolver(
+        platform=platform,
+        platform_instance=platform_instance,
+        env=env,
+        schema_aware=schema_aware,
+        graph=graph,
+    )
+
+    needs_close: bool = True
     if graph and schema_aware:
         needs_close = False
-        schema_resolver = graph._make_schema_resolver(
-            platform=platform,
-            platform_instance=platform_instance,
-            env=env,
-        )
-    else:
-        needs_close = True
-        schema_resolver = SchemaResolver(
-            platform=platform,
-            platform_instance=platform_instance,
-            env=env,
-            graph=None,
-        )
 
     try:
         return sqlglot_lineage(
@@ -1243,13 +1278,19 @@ def infer_output_schema(result: SqlParsingResult) -> Optional[List[SchemaFieldCl
 def view_definition_lineage_helper(
     result: SqlParsingResult, view_urn: str
 ) -> SqlParsingResult:
-    if result.query_type is QueryType.SELECT:
+    if result.query_type is QueryType.SELECT or (
+        result.out_tables and result.out_tables != [view_urn]
+    ):
         # Some platforms (e.g. postgres) store only <select statement> from view definition
         # `create view V as <select statement>` . For such view definitions, `result.out_tables` and
         # `result.column_lineage[].downstream` are empty in `sqlglot_lineage` response, whereas upstream
         # details and downstream column details are extracted correctly.
         # Here, we inject view V's urn in `result.out_tables` and `result.column_lineage[].downstream`
         # to get complete lineage result.
+
+        # Some platforms(e.g. mssql) may have slightly different view name in view definition than
+        # actual view name used elsewhere. Therefore we overwrite downstream table for such cases as well.
+
         result.out_tables = [view_urn]
         if result.column_lineage:
             for col_result in result.column_lineage:

datahub/sql_parsing/sqlglot_utils.py

@@ -121,7 +121,7 @@ _BASIC_NORMALIZATION_RULES = {
     # Remove /* */ comments.
     re.compile(r"/\*.*?\*/", re.DOTALL): "",
     # Remove -- comments.
-    re.compile(r"--.*$"): "",
+    re.compile(r"--.*$", re.MULTILINE): "",
     # Replace all runs of whitespace with a single space.
     re.compile(r"\s+"): " ",
     # Remove leading and trailing whitespace and trailing semicolons.
@@ -131,10 +131,16 @@ _BASIC_NORMALIZATION_RULES = {
     # Replace anything that looks like a string with a placeholder.
     re.compile(r"'[^']*'"): "?",
     # Replace sequences of IN/VALUES with a single placeholder.
-    re.compile(r"\b(IN|VALUES)\s*\(\?(?:, \?)*\)", re.IGNORECASE): r"\1 (?)",
+    # The r" ?" makes it more robust to uneven spacing.
+    re.compile(r"\b(IN|VALUES)\s*\( ?\?(?:, ?\?)* ?\)", re.IGNORECASE): r"\1 (?)",
     # Normalize parenthesis spacing.
     re.compile(r"\( "): "(",
     re.compile(r" \)"): ")",
+    # Fix up spaces before commas in column lists.
+    # e.g. "col1 , col2" -> "col1, col2"
+    # e.g. "col1,col2" -> "col1, col2"
+    re.compile(r"\b ,"): ",",
+    re.compile(r"\b,\b"): ", ",
 }
 _TABLE_NAME_NORMALIZATION_RULES = {
     # Replace UUID-like strings with a placeholder (both - and _ variants).

datahub/telemetry/telemetry.py

@@ -7,7 +7,7 @@ import sys
 import uuid
 from functools import wraps
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, TypeVar
 
 from mixpanel import Consumer, Mixpanel
 from typing_extensions import ParamSpec
@@ -16,10 +16,12 @@ import datahub as datahub_package
 from datahub.cli.config_utils import DATAHUB_ROOT_FOLDER
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import ExceptionWithProps
-from datahub.ingestion.graph.client import DataHubGraph
 from datahub.metadata.schema_classes import _custom_package_path
 from datahub.utilities.perf_timer import PerfTimer
 
+if TYPE_CHECKING:
+    from datahub.ingestion.graph.client import DataHubGraph
+
 logger = logging.getLogger(__name__)
 
 DATAHUB_FOLDER = Path(DATAHUB_ROOT_FOLDER)
@@ -117,7 +119,11 @@ class Telemetry:
     tracking_init: bool = False
     sentry_enabled: bool = False
 
+    context_properties: Dict[str, Any] = {}
+
     def __init__(self):
+        self.context_properties = {}
+
         if SENTRY_DSN:
             self.sentry_enabled = True
             try:
@@ -157,6 +163,9 @@ class Telemetry:
             except Exception as e:
                 logger.debug(f"Error connecting to mixpanel: {e}")
 
+        # Initialize the default properties for all events.
+        self.set_context()
+
     def update_config(self) -> bool:
         """
         Update the config file with the current client ID and enabled status.
@@ -238,18 +247,22 @@ class Telemetry:
 
         return False
 
-    def update_capture_exception_context(
+    def set_context(
         self,
-        server: Optional[DataHubGraph] = None,
+        server: Optional["DataHubGraph"] = None,
         properties: Optional[Dict[str, Any]] = None,
     ) -> None:
+        self.context_properties = {
+            **self._server_props(server),
+            **(properties or {}),
+        }
+
         if self.sentry_enabled:
             from sentry_sdk import set_tag
 
             properties = {
                 **_default_telemetry_properties(),
-                **self._server_props(server),
-                **(properties or {}),
+                **self.context_properties,
             }
 
             for key in properties:
@@ -297,7 +310,6 @@ class Telemetry:
         self,
         event_name: str,
         properties: Optional[Dict[str, Any]] = None,
-        server: Optional[DataHubGraph] = None,
     ) -> None:
         """
         Send a single telemetry event.
@@ -323,14 +335,15 @@ class Telemetry:
 
             properties = {
                 **_default_telemetry_properties(),
-                **self._server_props(server),
+                **self.context_properties,
                 **properties,
             }
             self.mp.track(self.client_id, event_name, properties)
         except Exception as e:
             logger.debug(f"Error reporting telemetry: {e}")
 
-    def _server_props(self, server: Optional[DataHubGraph]) -> Dict[str, str]:
+    @classmethod
+    def _server_props(cls, server: Optional["DataHubGraph"]) -> Dict[str, str]:
         if not server:
             return {
                 "server_type": "n/a",
@@ -435,6 +448,7 @@ def with_telemetry(
                             **call_props,
                             "status": "error",
                             **_error_props(e),
+                            "code": e.code,
                         },
                     )
                 telemetry_instance.capture_exception(e)

datahub/testing/compare_metadata_json.py

@@ -117,7 +117,7 @@ def diff_metadata_json(
     ignore_paths: Sequence[str] = (),
     ignore_order: bool = True,
 ) -> Union[DeepDiff, MCPDiff]:
-    ignore_paths = (*ignore_paths, *default_exclude_paths, r"root\[\d+].delta_info")
+    ignore_paths = [*ignore_paths, *default_exclude_paths, r"root\[\d+].delta_info"]
     try:
         if ignore_order:
             golden_map = get_aspects_by_urn(golden)