openmetadata-ingestion 1.6.2.0__py3-none-any.whl → 1.10.14.4__py3-none-any.whl

metadata/ingestion/connections/test_connections.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -38,7 +38,7 @@ from metadata.generated.schema.entity.services.connections.testConnectionResult
 )
 from metadata.generated.schema.type.basic import Timestamp
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import kill_active_connections
+from metadata.ingestion.source.connections_utils import kill_active_connections
 from metadata.profiler.orm.functions.conn_test import ConnTestFn
 from metadata.utils.constants import THREE_MIN
 from metadata.utils.logger import cli_logger

metadata/ingestion/lineage/masker.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -13,17 +13,23 @@ Query masking utilities
 """
 
 import traceback
+from typing import Optional
 
-import sqlparse
-from sqlfluff.core import Linter
+from cachetools import LRUCache
+from collate_sqllineage.runner import SQLPARSE_DIALECT, LineageRunner
 from sqlparse.sql import Comparison
 from sqlparse.tokens import Literal, Number, String
 
 from metadata.ingestion.lineage.models import Dialect
+from metadata.utils.execution_time_tracker import calculate_execution_time
 
 MASK_TOKEN = "?"
 
+# Cache size is 128 to avoid memory issues
+masked_query_cache = LRUCache(maxsize=128)
 
+
+# pylint: disable=protected-access
 def get_logger():
     # pylint: disable=import-outside-toplevel
     from metadata.utils.logger import utils_logger
@@ -31,18 +37,14 @@ def get_logger():
     return utils_logger()
 
 
-def mask_literals_with_sqlparse(query: str):
+def mask_literals_with_sqlparse(query: str, parser: LineageRunner):
     """
     Mask literals in a query using sqlparse.
     """
     logger = get_logger()
 
     try:
-        parsed = sqlparse.parse(query)  # Parse the query
-
-        if not parsed:
-            return query
-        parsed = parsed[0]
+        parsed = parser._parsed_result
 
         def mask_token(token):
             # Mask all literals: strings, numbers, or other literal values
@@ -53,7 +55,6 @@ def mask_literals_with_sqlparse(query: str):
                 Literal.Number.Integer,
                 Literal.Number.Float,
                 Literal.String.Single,
-                Literal.String.Symbol,
             ):
                 token.value = MASK_TOKEN
             elif token.is_group:
@@ -79,17 +80,16 @@ def mask_literals_with_sqlparse(query: str):
     return query
 
 
-def mask_literals_with_sqlfluff(query: str, dialect: str = Dialect.ANSI.value) -> str:
+def mask_literals_with_sqlfluff(query: str, parser: LineageRunner) -> str:
     """
     Mask literals in a query using SQLFluff.
     """
     logger = get_logger()
     try:
-        # Initialize SQLFluff linter
-        linter = Linter(dialect=dialect)
+        if not parser._evaluated:
+            parser._eval()
 
-        # Parse the query
-        parsed = linter.parse_string(query)
+        parsed = parser._parsed_result
 
         def replace_literals(segment):
             """Recursively replace literals with placeholders."""
@@ -114,18 +114,36 @@ def mask_literals_with_sqlfluff(query: str, dialect: str = Dialect.ANSI.value) -
     return query
 
 
-def mask_query(query: str, dialect: str = Dialect.ANSI.value) -> str:
+@calculate_execution_time(context="MaskQuery")
+def mask_query(
+    query: str,
+    dialect: str = Dialect.ANSI.value,
+    parser: Optional[LineageRunner] = None,
+    parser_required: bool = False,
+) -> str:
+    """
+    Mask a query using sqlparse or sqlfluff.
+    """
     logger = get_logger()
     try:
-        sqlfluff_masked_query = mask_literals_with_sqlfluff(query, dialect)
-        sqlparse_masked_query = mask_literals_with_sqlparse(query)
-        # compare both masked queries and return the one with more masked tokens
-        if sqlfluff_masked_query.count(MASK_TOKEN) >= sqlparse_masked_query.count(
-            MASK_TOKEN
-        ):
-            return sqlfluff_masked_query
-        return sqlparse_masked_query
+        if masked_query_cache.get((query, dialect)):
+            return masked_query_cache.get((query, dialect))
+        if parser_required and not parser:
+            return None
+        if not parser:
+            try:
+                parser = LineageRunner(query, dialect=dialect)
+                len(parser.source_tables)
+            except Exception:
+                parser = LineageRunner(query)
+                len(parser.source_tables)
+        if parser._dialect == SQLPARSE_DIALECT:
+            masked_query = mask_literals_with_sqlparse(query, parser)
+        else:
+            masked_query = mask_literals_with_sqlfluff(query, parser)
+        masked_query_cache[(query, dialect)] = masked_query
+        return masked_query
     except Exception as exc:
         logger.debug(f"Failed to mask query with sqlfluff: {exc}")
         logger.debug(traceback.format_exc())
-    return query
+    return None

metadata/ingestion/lineage/models.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -37,12 +37,21 @@ from metadata.generated.schema.entity.services.connections.database.db2Connectio
 from metadata.generated.schema.entity.services.connections.database.deltaLakeConnection import (
     DeltaLakeType,
 )
+from metadata.generated.schema.entity.services.connections.database.exasolConnection import (
+    ExasolType,
+)
+from metadata.generated.schema.entity.services.connections.database.greenplumConnection import (
+    GreenplumType,
+)
 from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
     HiveType,
 )
 from metadata.generated.schema.entity.services.connections.database.impalaConnection import (
     ImpalaType,
 )
+from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
+    MariaDBType,
+)
 from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
     MssqlType,
 )
@@ -58,6 +67,9 @@ from metadata.generated.schema.entity.services.connections.database.postgresConn
 from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
     RedshiftType,
 )
+from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
+    SingleStoreType,
+)
 from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
     SnowflakeType,
 )
@@ -67,9 +79,15 @@ from metadata.generated.schema.entity.services.connections.database.sqliteConnec
 from metadata.generated.schema.entity.services.connections.database.teradataConnection import (
     TeradataType,
 )
+from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
+    TrinoType,
+)
 from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
     DatabricksType as UnityCatalogType,
 )
+from metadata.generated.schema.entity.services.connections.database.verticaConnection import (
+    VerticaType,
+)
 from metadata.utils.singleton import Singleton
 
 
@@ -99,6 +117,9 @@ class Dialect(Enum):
     SQLITE = "sqlite"
     TERADATA = "teradata"
     TSQL = "tsql"
+    MARIADB = "mariadb"
+    TRINO = "trino"
+    VERTICA = "vertica"
 
 
 MAP_CONNECTION_TYPE_DIALECT: Dict[str, Dialect] = {
@@ -120,6 +141,12 @@ MAP_CONNECTION_TYPE_DIALECT: Dict[str, Dialect] = {
     str(MssqlType.Mssql.value): Dialect.TSQL,
     str(AzureSQLType.AzureSQL.value): Dialect.TSQL,
     str(TeradataType.Teradata.value): Dialect.TERADATA,
+    str(MariaDBType.MariaDB.value): Dialect.MARIADB,
+    str(SingleStoreType.SingleStore.value): Dialect.MYSQL,
+    str(ExasolType.Exasol.value): Dialect.EXASOL,
+    str(TrinoType.Trino.value): Dialect.TRINO,
+    str(VerticaType.Vertica.value): Dialect.VERTICA,
+    str(GreenplumType.Greenplum.value): Dialect.POSTGRES,
 }
 
 

metadata/ingestion/lineage/parser.py

@@ -1,8 +1,8 @@
-#  Copyright 2021 Collate
-#  Licensed under the Apache License, Version 2.0 (the "License");
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
 #  you may not use this file except in compliance with the License.
 #  You may obtain a copy of the License at
-#  http://www.apache.org/licenses/LICENSE-2.0
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
 #  Unless required by applicable law or agreed to in writing, software
 #  distributed under the License is distributed on an "AS IS" BASIS,
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
@@ -28,6 +28,7 @@ from sqlparse.sql import Comparison, Identifier, Parenthesis, Statement
 from metadata.generated.schema.type.tableUsageCount import TableColumn, TableColumnJoin
 from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.lineage.models import Dialect
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.helpers import (
     find_in_iter,
     get_formatted_entity_name,
@@ -71,12 +72,16 @@ class LineageParser:
         self.query_parsing_success = True
         self.query_parsing_failure_reason = None
         self.dialect = dialect
-        self._masked_query = mask_query(self.query, dialect.value)
+        self.masked_query = None
         self._clean_query = self.clean_raw_query(query)
-        self._masked_clean_query = mask_query(self._clean_query, dialect.value)
         self.parser = self._evaluate_best_parser(
             self._clean_query, dialect=dialect, timeout_seconds=timeout_seconds
         )
+        if self.masked_query is None:
+            self.masked_query = (
+                mask_query(self._clean_query, parser=self.parser, parser_required=True)
+                or self._clean_query
+            )
 
     @cached_property
     def involved_tables(self) -> Optional[List[Table]]:
@@ -87,16 +92,22 @@ class LineageParser:
         :return: List of involved tables
         """
         try:
+            logger.debug(f"[UsageSink] Source tables: {self.source_tables}")
+            logger.debug(f"[UsageSink] Intermediate tables: {self.intermediate_tables}")
+            logger.debug(f"[UsageSink] Target tables: {self.target_tables}")
+
             return list(
                 set(self.source_tables)
                 .union(set(self.intermediate_tables))
                 .union(set(self.target_tables))
             )
+
         except SQLLineageException as exc:
             logger.debug(traceback.format_exc())
             logger.warning(
-                f"Cannot extract source table information from query [{self._masked_query}]: {exc}"
+                f"Cannot extract source table information from query [{self.masked_query or self.query}]: {exc}"
             )
+
             return None
 
     @cached_property
@@ -171,14 +182,38 @@ class LineageParser:
     def table_aliases(self) -> Dict[str, str]:
         """
         Prepare a dictionary in the shape of {alias: table_name} from
-        the parser tables
+        the parser tables, with detailed logging for debugging.
         :return: alias dict
         """
-        return {
+        # Check if involved_tables is present
+        if not self.involved_tables:
+            logger.debug(
+                "[UsageSink] No involved tables found — alias map will be empty."
+            )
+            return {}
+
+        # Log raw involved tables for inspection
+        logger.debug(
+            f"[UsageSink] Involved tables before alias mapping: {[str(t) for t in self.involved_tables]}"
+        )
+
+        # Log alias/name pairs for each table
+        logger.debug(
+            f"[UsageSink] Table alias-name pairs: {[(t.alias, str(t)) for t in self.involved_tables]}"
+        )
+
+        # Build the alias dictionary
+        alias_map = {
             table.alias: str(table).replace("<default>.", "")
             for table in self.involved_tables
+            if getattr(table, "alias", None)  # Only include if alias is not None
         }
 
+        # Log the final computed alias map
+        logger.debug(f"[UsageSink] Final computed alias map: {alias_map}")
+
+        return alias_map
+
     def get_table_name_from_list(
         self,
         database_name: Optional[str],
@@ -194,6 +229,11 @@ class LineageParser:
         :return: table name from parser info
         """
         tables = self.clean_table_list
+        logger.debug(
+            f"[UsageSink] Searching for table '{table_name}' "
+            f"in schema '{schema_name}', database '{database_name}', "
+            f"within tables list: {tables}"
+        )
         table = find_in_iter(element=table_name, container=tables)
         if table:
             return table
@@ -221,8 +261,15 @@ class LineageParser:
         :param identifier: comparison identifier
         :return: table name and column name from the identifier
         """
+        logger.debug(f"[DEBUG] Raw identifier object: {identifier!r}")
+        logger.debug(f"[DEBUG] Identifier type: {type(identifier)}")
+        logger.debug(f"[DEBUG] Identifier value: {getattr(identifier, 'value', None)}")
+
         aliases = self.table_aliases
+        logger.debug(f"[DEBUG] Current table aliases: {aliases}")
+
         values = identifier.value.split(".")
+        logger.debug(f"[DEBUG] Split identifier values: {values}")
 
         if len(values) > 4:
             logger.debug(f"Invalid comparison element from identifier: {identifier}")
@@ -232,6 +279,12 @@ class LineageParser:
             [None] * (4 - len(values))
         ) + values
 
+        logger.debug(
+            f"[DEBUG] Parsed components => "
+            f"database_name={database_name}, schema_name={schema_name}, "
+            f"table_or_alias={table_or_alias}, column_name={column_name}"
+        )
+
         if not table_or_alias or not column_name:
             logger.debug(
                 f"Cannot obtain comparison elements from identifier {identifier}"
@@ -334,12 +387,10 @@ class LineageParser:
                 )
 
                 if not table_left or not table_right:
-                    logger.warning(
-                        f"Can't extract table names when parsing JOIN information from {comparison}"
-                    )
                     logger.debug(
-                        f"Query: {mask_query(sql_statement, self.dialect.value)}"
+                        f"Can't extract table names when parsing JOIN information from {comparison}"
                     )
+                    logger.debug(f"Query: {self.masked_query or self.query}")
                     continue
 
                 left_table_column = TableColumn(table=table_left, column=column_left)
@@ -410,6 +461,7 @@ class LineageParser:
 
         return clean_query.strip()
 
+    @calculate_execution_time(context="EvaluateBestParser")
     def _evaluate_best_parser(
         self, query: str, dialect: Dialect, timeout_seconds: int
     ) -> Optional[LineageRunner]:
@@ -422,10 +474,9 @@ class LineageParser:
             lr_dialect.get_column_lineage()
             return lr_dialect
 
-        sqlfluff_count = 0
         try:
             lr_sqlfluff = get_sqlfluff_lineage_runner(query, dialect.value)
-            sqlfluff_count = len(lr_sqlfluff.get_column_lineage()) + len(
+            _ = len(lr_sqlfluff.get_column_lineage()) + len(
                 set(lr_sqlfluff.source_tables).union(
                     set(lr_sqlfluff.target_tables).union(
                         set(lr_sqlfluff.intermediate_tables)
@@ -438,46 +489,48 @@ class LineageParser:
                 f"Lineage with SqlFluff failed for the [{dialect.value}]. "
                 f"Parser has been running for more than {timeout_seconds} seconds."
             )
-            logger.debug(
-                f"{self.query_parsing_failure_reason}] query: [{self._masked_clean_query}]"
-            )
             lr_sqlfluff = None
         except Exception:
             self.query_parsing_success = False
             self.query_parsing_failure_reason = (
                 f"Lineage with SqlFluff failed for the [{dialect.value}]"
             )
-            logger.debug(
-                f"{self.query_parsing_failure_reason} query: [{self._masked_clean_query}]"
-            )
             lr_sqlfluff = None
 
-        lr_sqlparser = LineageRunner(query)
+        if lr_sqlfluff:
+            return lr_sqlfluff
+
+        @timeout(seconds=timeout_seconds)
+        def get_sqlparser_lineage_runner(qry: str) -> LineageRunner:
+            lr_sqlparser = LineageRunner(qry)
+            lr_sqlparser.get_column_lineage()
+            return lr_sqlparser
+
+        lr_sqlparser = None
         try:
-            sqlparser_count = len(lr_sqlparser.get_column_lineage()) + len(
+            lr_sqlparser = get_sqlparser_lineage_runner(query)
+            _ = len(lr_sqlparser.get_column_lineage()) + len(
                 set(lr_sqlparser.source_tables).union(
                     set(lr_sqlparser.target_tables).union(
                         set(lr_sqlparser.intermediate_tables)
                     )
                 )
             )
+        except TimeoutError:
+            self.query_parsing_success = False
+            self.query_parsing_failure_reason = (
+                f"Lineage with SqlParser failed for the [{dialect.value}]. "
+                f"Parser has been running for more than {timeout_seconds} seconds."
+            )
+            return None
         except Exception:
             # if both runner have failed we return the usual one
-            return lr_sqlfluff if lr_sqlfluff else lr_sqlparser
+            logger.debug(f"Failed to parse query with sqlparse & sqlfluff: {query}")
+            return lr_sqlfluff if lr_sqlfluff else None
 
-        if lr_sqlfluff:
-            # if sqlparser retrieve more lineage info that sqlfluff
-            if sqlparser_count > sqlfluff_count:
-                self.query_parsing_success = False
-                self.query_parsing_failure_reason = (
-                    "Lineage computed with SqlFluff did not perform as expected "
-                    f"for the [{dialect.value}]"
-                )
-                logger.debug(
-                    f"{self.query_parsing_failure_reason} query: [{self._masked_clean_query}]"
-                )
-                return lr_sqlparser
-            return lr_sqlfluff
+        logger.debug(
+            f"Using sqlparse for lineage parsing for query: {self.masked_query or self.query}"
+        )
         return lr_sqlparser
 
     @staticmethod