acryl-datahub 0.15.0rc15__py3-none-any.whl → 0.15.0rc17__py3-none-any.whl

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/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/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(

datahub/testing/doctest.py

@@ -0,0 +1,12 @@
+import doctest
+from types import ModuleType
+
+
+def assert_doctest(module: ModuleType) -> None:
+    result = doctest.testmod(
+        module,
+        raise_on_error=True,
+        verbose=True,
+    )
+    if result.attempted == 0:
+        raise ValueError(f"No doctests found in {module.__name__}")

datahub/utilities/partition_executor.py

@@ -268,7 +268,7 @@ class BatchPartitionExecutor(Closeable):
         self.process_batch = process_batch
         self.min_process_interval = min_process_interval
         self.read_from_pending_interval = read_from_pending_interval
-        assert self.max_workers > 1
+        assert self.max_workers >= 1
 
         self._state_lock = threading.Lock()
         self._executor = ThreadPoolExecutor(

datahub/utilities/sql_lineage_parser_impl.py

@@ -1,160 +0,0 @@
-import contextlib
-import logging
-import re
-import unittest
-import unittest.mock
-from typing import Dict, List, Optional, Set
-
-from sqllineage.core.holders import Column, SQLLineageHolder
-from sqllineage.exceptions import SQLLineageException
-
-from datahub.utilities.sql_parser_base import SQLParser, SqlParserException
-
-with contextlib.suppress(ImportError):
-    import sqlparse
-    from networkx import DiGraph
-    from sqllineage.core import LineageAnalyzer
-
-    import datahub.utilities.sqllineage_patch
-logger = logging.getLogger(__name__)
-
-
-class SqlLineageSQLParserImpl(SQLParser):
-    _DATE_SWAP_TOKEN = "__d_a_t_e"
-    _HOUR_SWAP_TOKEN = "__h_o_u_r"
-    _TIMESTAMP_SWAP_TOKEN = "__t_i_m_e_s_t_a_m_p"
-    _DATA_SWAP_TOKEN = "__d_a_t_a"
-    _ADMIN_SWAP_TOKEN = "__a_d_m_i_n"
-    _MYVIEW_SQL_TABLE_NAME_TOKEN = "__my_view__.__sql_table_name__"
-    _MYVIEW_LOOKER_TOKEN = "my_view.SQL_TABLE_NAME"
-
-    def __init__(self, sql_query: str, use_raw_names: bool = False) -> None:
-        super().__init__(sql_query)
-        original_sql_query = sql_query
-        self._use_raw_names = use_raw_names
-
-        # SqlLineageParser makes mistakes on lateral flatten queries, use the prefix
-        if "lateral flatten" in sql_query:
-            sql_query = sql_query[: sql_query.find("lateral flatten")]
-
-        # Replace reserved words that break SqlLineageParser
-        self.token_to_original: Dict[str, str] = {
-            self._DATE_SWAP_TOKEN: "date",
-            self._HOUR_SWAP_TOKEN: "hour",
-            self._TIMESTAMP_SWAP_TOKEN: "timestamp",
-            self._DATA_SWAP_TOKEN: "data",
-            self._ADMIN_SWAP_TOKEN: "admin",
-        }
-        for replacement, original in self.token_to_original.items():
-            # Replace original tokens with replacement. Since table and column name can contain a hyphen('-'),
-            # also prevent original tokens appearing as part of these names with a hyphen from getting substituted.
-            sql_query = re.sub(
-                rf"((?<!-)\b{original}\b)(?!-)",
-                rf"{replacement}",
-                sql_query,
-                flags=re.IGNORECASE,
-            )
-
-        # SqlLineageParser lowercarese tablenames and we need to replace Looker specific token which should be uppercased
-        sql_query = re.sub(
-            rf"(\${{{self._MYVIEW_LOOKER_TOKEN}}})",
-            rf"{self._MYVIEW_SQL_TABLE_NAME_TOKEN}",
-            sql_query,
-        )
-
-        # SqlLineageParser does not handle "encode" directives well. Remove them
-        sql_query = re.sub(r"\sencode [a-zA-Z]*", "", sql_query, flags=re.IGNORECASE)
-
-        # Replace lookml templates with the variable otherwise sqlparse can't parse ${
-        sql_query = re.sub(r"(\${)(.+)(})", r"\2", sql_query)
-        if sql_query != original_sql_query:
-            logger.debug(f"Rewrote original query {original_sql_query} as {sql_query}")
-
-        self._sql = sql_query
-        self._stmt_holders: Optional[List[LineageAnalyzer]] = None
-        self._sql_holder: Optional[SQLLineageHolder] = None
-        try:
-            self._stmt = [
-                s
-                for s in sqlparse.parse(
-                    # first apply sqlparser formatting just to get rid of comments, which cause
-                    # inconsistencies in parsing output
-                    sqlparse.format(
-                        self._sql.strip(),
-                        strip_comments=True,
-                        use_space_around_operators=True,
-                    ),
-                )
-                if s.token_first(skip_cm=True)
-            ]
-
-            with unittest.mock.patch(
-                "sqllineage.core.handlers.source.SourceHandler.end_of_query_cleanup",
-                datahub.utilities.sqllineage_patch.end_of_query_cleanup_patch,
-            ):
-                with unittest.mock.patch(
-                    "sqllineage.core.holders.SubQueryLineageHolder.add_column_lineage",
-                    datahub.utilities.sqllineage_patch.add_column_lineage_patch,
-                ):
-                    self._stmt_holders = [
-                        LineageAnalyzer().analyze(stmt) for stmt in self._stmt
-                    ]
-                    self._sql_holder = SQLLineageHolder.of(*self._stmt_holders)
-        except SQLLineageException as e:
-            raise SqlParserException(
-                f"SQL lineage analyzer error '{e}' for query: '{self._sql}"
-            ) from e
-
-    def get_tables(self) -> List[str]:
-        result: List[str] = []
-        if self._sql_holder is None:
-            logger.error("sql holder not present so cannot get tables")
-            return result
-        for table in self._sql_holder.source_tables:
-            table_normalized = re.sub(
-                r"^<default>.",
-                "",
-                (
-                    str(table)
-                    if not self._use_raw_names
-                    else f"{table.schema.raw_name}.{table.raw_name}"
-                ),
-            )
-            result.append(str(table_normalized))
-
-        # We need to revert TOKEN replacements
-        for token, replacement in self.token_to_original.items():
-            result = [replacement if c == token else c for c in result]
-        result = [
-            self._MYVIEW_LOOKER_TOKEN if c == self._MYVIEW_SQL_TABLE_NAME_TOKEN else c
-            for c in result
-        ]
-
-        # Sort tables to make the list deterministic
-        result.sort()
-
-        return result
-
-    def get_columns(self) -> List[str]:
-        if self._sql_holder is None:
-            raise SqlParserException("sql holder not present so cannot get columns")
-        graph: DiGraph = self._sql_holder.graph  # For mypy attribute checking
-        column_nodes = [n for n in graph.nodes if isinstance(n, Column)]
-        column_graph = graph.subgraph(column_nodes)
-
-        target_columns = {column for column, deg in column_graph.out_degree if deg == 0}
-
-        result: Set[str] = set()
-        for column in target_columns:
-            # Let's drop all the count(*) and similard columns which are expression actually if it does not have an alias
-            if not any(ele in column.raw_name for ele in ["*", "(", ")"]):
-                result.add(str(column.raw_name))
-
-        # Reverting back all the previously renamed words which confuses the parser
-        result = {"date" if c == self._DATE_SWAP_TOKEN else c for c in result}
-        result = {
-            "timestamp" if c == self._TIMESTAMP_SWAP_TOKEN else c for c in list(result)
-        }
-
-        # swap back renamed date column
-        return list(result)

datahub/utilities/sql_parser.py

@@ -1,94 +0,0 @@
-import logging
-import multiprocessing
-import traceback
-from multiprocessing import Process, Queue
-from typing import Any, List, Optional, Tuple
-
-from datahub.utilities.sql_lineage_parser_impl import SqlLineageSQLParserImpl
-from datahub.utilities.sql_parser_base import SQLParser
-
-logger = logging.getLogger(__name__)
-
-
-def sql_lineage_parser_impl_func_wrapper(
-    queue: Optional[multiprocessing.Queue], sql_query: str, use_raw_names: bool = False
-) -> Optional[Tuple[List[str], List[str], Any]]:
-    """
-    The wrapper function that computes the tables and columns using the SqlLineageSQLParserImpl
-    and puts the results on the shared IPC queue. This is used to isolate SqlLineageSQLParserImpl
-    functionality in a separate process, and hence protect our sources from memory leaks originating in
-    the sqllineage module.
-    :param queue: The shared IPC queue on to which the results will be put.
-    :param sql_query: The SQL query to extract the tables & columns from.
-    :param use_raw_names: Parameter used to ignore sqllineage's default lowercasing.
-    :return: None.
-    """
-    exception_details: Optional[Tuple[BaseException, str]] = None
-    tables: List[str] = []
-    columns: List[str] = []
-    try:
-        parser = SqlLineageSQLParserImpl(sql_query, use_raw_names)
-        tables = parser.get_tables()
-        columns = parser.get_columns()
-    except BaseException as e:
-        exc_msg = traceback.format_exc()
-        exception_details = (e, exc_msg)
-        logger.debug(exc_msg)
-
-    if queue is not None:
-        queue.put((tables, columns, exception_details))
-        return None
-    else:
-        return (tables, columns, exception_details)
-
-
-class SqlLineageSQLParser(SQLParser):
-    def __init__(
-        self,
-        sql_query: str,
-        use_external_process: bool = False,
-        use_raw_names: bool = False,
-    ) -> None:
-        super().__init__(sql_query, use_external_process)
-        if use_external_process:
-            self.tables, self.columns = self._get_tables_columns_process_wrapped(
-                sql_query, use_raw_names
-            )
-        else:
-            return_tuple = sql_lineage_parser_impl_func_wrapper(
-                None, sql_query, use_raw_names
-            )
-            if return_tuple is not None:
-                (
-                    self.tables,
-                    self.columns,
-                    some_exception,
-                ) = return_tuple
-
-    @staticmethod
-    def _get_tables_columns_process_wrapped(
-        sql_query: str, use_raw_names: bool = False
-    ) -> Tuple[List[str], List[str]]:
-        # Invoke sql_lineage_parser_impl_func_wrapper in a separate process to avoid
-        # memory leaks from sqllineage module used by SqlLineageSQLParserImpl. This will help
-        # shield our sources like lookml & redash, that need to parse a large number of SQL statements,
-        # from causing significant memory leaks in the datahub cli during ingestion.
-        queue: multiprocessing.Queue = Queue()
-        process: multiprocessing.Process = Process(
-            target=sql_lineage_parser_impl_func_wrapper,
-            args=(queue, sql_query, use_raw_names),
-        )
-        process.start()
-        tables, columns, exception_details = queue.get(block=True)
-        if exception_details is not None:
-            raise exception_details[0](f"Sub-process exception: {exception_details[1]}")
-        return tables, columns
-
-    def get_tables(self) -> List[str]:
-        return self.tables
-
-    def get_columns(self) -> List[str]:
-        return self.columns
-
-
-DefaultSQLParser = SqlLineageSQLParser

datahub/utilities/sql_parser_base.py

@@ -1,21 +0,0 @@
-from abc import ABCMeta, abstractmethod
-from typing import List
-
-
-class SqlParserException(Exception):
-    """Raised when sql parser fails"""
-
-    pass
-
-
-class SQLParser(metaclass=ABCMeta):
-    def __init__(self, sql_query: str, use_external_process: bool = True) -> None:
-        self._sql_query = sql_query
-
-    @abstractmethod
-    def get_tables(self) -> List[str]:
-        pass
-
-    @abstractmethod
-    def get_columns(self) -> List[str]:
-        pass