openmetadata-ingestion 1.5.12.0__py3-none-any.whl → 1.5.13.0__py3-none-any.whl

metadata/generated/schema/type/entityRelationship.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityRelationship.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-11-26T06:50:54+00:00
+#   timestamp: 2024-11-28T18:44:30+00:00
 
 from __future__ import annotations
 

metadata/ingestion/lineage/parser.py

@@ -15,12 +15,12 @@ import traceback
 from collections import defaultdict
 from copy import deepcopy
 from logging.config import DictConfigurator
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import sqlparse
 from cached_property import cached_property
 from collate_sqllineage import SQLPARSE_DIALECT
-from collate_sqllineage.core.models import Column, Table
+from collate_sqllineage.core.models import Column, DataFunction, Table
 from collate_sqllineage.exceptions import SQLLineageException
 from collate_sqllineage.runner import LineageRunner
 from sqlparse.sql import Comparison, Identifier, Parenthesis, Statement
@@ -110,7 +110,7 @@ class LineageParser:
         return []
 
     @cached_property
-    def source_tables(self) -> List[Table]:
+    def source_tables(self) -> List[Union[Table, DataFunction]]:
         """
         Get a list of source tables
         """
@@ -374,7 +374,9 @@ class LineageParser:
         if not self._clean_query:
             return []
         return [
-            self.clean_table_name(table) for table in tables if isinstance(table, Table)
+            self.clean_table_name(table)
+            for table in tables
+            if isinstance(table, Table) or isinstance(table, DataFunction)
         ]
 
     @classmethod

metadata/ingestion/lineage/sql_lineage.py

@@ -13,9 +13,17 @@ Helper functions to handle SQL lineage operations
 """
 import itertools
 import traceback
-from typing import Any, Iterable, List, Optional, Tuple
+from typing import Any, Iterable, List, Optional, Tuple, Union
+
+from collate_sqllineage.core.models import Column, DataFunction
+from collate_sqllineage.core.models import Table as LineageTable
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.storedProcedure import (
+    Language,
+    StoredProcedure,
+    StoredProcedureType,
+)
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -163,6 +171,116 @@ def get_table_fqn_from_query_name(
     return database_query, schema_query, table
 
 
+def handle_udf_column_lineage(
+    column_lineage_original: dict,
+    column_lineage_generated: List[Tuple[Column, Column]],
+):
+    """
+    Handle UDF column lineage
+    """
+    try:
+        result = {}
+        intermediate_column_lineage = populate_column_lineage_map(
+            column_lineage_generated
+        )
+        # Iterate through the original dictionary
+        for source_table, mappings in column_lineage_original.items():
+            result[source_table] = {}
+            for intermediate_table, column_pairs in mappings.items():
+                # Iterate through each column mapping in the original dictionary
+                for source_column, intermediate_column in column_pairs:
+                    if intermediate_table in intermediate_column_lineage:
+                        # Check intermediate dictionary for mappings
+                        for (
+                            target_table,
+                            target_mappings,
+                        ) in intermediate_column_lineage[intermediate_table].items():
+                            for inter_col, target_col in target_mappings:
+                                if intermediate_column == inter_col:
+                                    # Append to the result dictionary
+                                    if target_table not in result[source_table]:
+                                        result[source_table][target_table] = []
+                                    result[source_table][target_table].append(
+                                        (source_column, target_col)
+                                    )
+        column_lineage_original.update(result)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.error(f"Error handling UDF column lineage: {exc}")
+
+
+def get_source_table_names(
+    metadata: OpenMetadata,
+    dialect: Dialect,
+    source_table: Union[DataFunction, LineageTable],
+    database_name: Optional[str],
+    schema_name: Optional[str],
+    service_name: Optional[str],
+    timeout_seconds: int,
+    column_lineage: dict,
+    procedure: Optional[StoredProcedure] = None,
+) -> Iterable[Tuple[Optional[EntityReference], str]]:
+    """
+    Get source table names from DataFunction
+    """
+    try:
+        if not isinstance(source_table, DataFunction):
+            yield EntityReference(
+                id=procedure.id.root, type="storedProcedure"
+            ) if procedure else None, str(source_table)
+        else:
+            database_query, schema_query, table = get_table_fqn_from_query_name(
+                str(source_table)
+            )
+            function_fqn_string = build_es_fqn_search_string(
+                database_query or database_name,
+                schema_query or schema_name,
+                service_name,
+                table,
+            )
+            es_result_entities: Optional[
+                List[StoredProcedure]
+            ] = metadata.es_search_from_fqn(
+                entity_type=StoredProcedure,
+                fqn_search_string=function_fqn_string,
+            )
+            if es_result_entities:
+                for entity in es_result_entities:
+                    if (
+                        entity.storedProcedureType == StoredProcedureType.UDF
+                        and entity.storedProcedureCode
+                        and entity.storedProcedureCode.language == Language.SQL
+                    ):
+                        expected_table_name = str(source_table).replace(
+                            f"{DEFAULT_SCHEMA_NAME}.", ""
+                        )
+                        lineage_parser = LineageParser(
+                            f"create table {str(expected_table_name)} as {entity.storedProcedureCode.code}",
+                            dialect=dialect,
+                            timeout_seconds=timeout_seconds,
+                        )
+                        handle_udf_column_lineage(
+                            column_lineage, lineage_parser.column_lineage
+                        )
+                        for source in lineage_parser.source_tables or []:
+                            yield from get_source_table_names(
+                                metadata,
+                                dialect,
+                                source,
+                                database_name,
+                                schema_name,
+                                service_name,
+                                timeout_seconds,
+                                column_lineage,
+                                procedure or entity,
+                            )
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.error(
+            f"Error getting source table names for table [{source_table}]: {exc}"
+        )
+
+
 def get_table_entities_from_query(
     metadata: OpenMetadata,
     service_name: str,
@@ -252,6 +370,7 @@ def _build_table_lineage(
     masked_query: str,
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    procedure: Optional[EntityReference] = None,
 ) -> Either[AddLineageRequest]:
     """
     Prepare the lineage request generator
@@ -276,7 +395,9 @@ def _build_table_lineage(
             from_table_raw_name=str(from_table_raw_name),
             column_lineage_map=column_lineage_map,
         )
-        lineage_details = LineageDetails(sqlQuery=masked_query, source=lineage_source)
+        lineage_details = LineageDetails(
+            sqlQuery=masked_query, source=lineage_source, pipeline=procedure
+        )
         if col_lineage:
             lineage_details.columnsLineage = col_lineage
         lineage = AddLineageRequest(
@@ -315,6 +436,7 @@ def _create_lineage_by_table_name(
     masked_query: str,
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    procedure: Optional[EntityReference] = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
     This method is to create a lineage between two tables
@@ -358,6 +480,7 @@ def _create_lineage_by_table_name(
                     masked_query=masked_query,
                     column_lineage_map=column_lineage_map,
                     lineage_source=lineage_source,
+                    procedure=procedure,
                 )
 
     except Exception as exc:
@@ -429,17 +552,28 @@ def get_lineage_by_query(
 
         for intermediate_table in lineage_parser.intermediate_tables:
             for source_table in lineage_parser.source_tables:
-                yield from _create_lineage_by_table_name(
-                    metadata,
-                    from_table=str(source_table),
-                    to_table=str(intermediate_table),
-                    service_name=service_name,
+                for procedure, from_table_name in get_source_table_names(
+                    metadata=metadata,
+                    dialect=dialect,
+                    source_table=source_table,
                     database_name=database_name,
                     schema_name=schema_name,
-                    masked_query=masked_query,
-                    column_lineage_map=column_lineage,
-                    lineage_source=lineage_source,
-                )
+                    service_name=service_name,
+                    timeout_seconds=timeout_seconds,
+                    column_lineage=column_lineage,
+                ):
+                    yield from _create_lineage_by_table_name(
+                        metadata,
+                        from_table=str(from_table_name),
+                        to_table=str(intermediate_table),
+                        service_name=service_name,
+                        database_name=database_name,
+                        schema_name=schema_name,
+                        masked_query=masked_query,
+                        column_lineage_map=column_lineage,
+                        lineage_source=lineage_source,
+                        procedure=procedure,
+                    )
             for target_table in lineage_parser.target_tables:
                 yield from _create_lineage_by_table_name(
                     metadata,
@@ -455,17 +589,28 @@ def get_lineage_by_query(
         if not lineage_parser.intermediate_tables:
             for target_table in lineage_parser.target_tables:
                 for source_table in lineage_parser.source_tables:
-                    yield from _create_lineage_by_table_name(
-                        metadata,
-                        from_table=str(source_table),
-                        to_table=str(target_table),
-                        service_name=service_name,
+                    for procedure, from_table_name in get_source_table_names(
+                        metadata=metadata,
+                        dialect=dialect,
+                        source_table=source_table,
                         database_name=database_name,
                         schema_name=schema_name,
-                        masked_query=masked_query,
-                        column_lineage_map=column_lineage,
-                        lineage_source=lineage_source,
-                    )
+                        service_name=service_name,
+                        timeout_seconds=timeout_seconds,
+                        column_lineage=column_lineage,
+                    ):
+                        yield from _create_lineage_by_table_name(
+                            metadata,
+                            from_table=str(from_table_name),
+                            to_table=str(target_table),
+                            service_name=service_name,
+                            database_name=database_name,
+                            schema_name=schema_name,
+                            masked_query=masked_query,
+                            column_lineage_map=column_lineage,
+                            lineage_source=lineage_source,
+                            procedure=procedure,
+                        )
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(
@@ -505,17 +650,28 @@ def get_lineage_via_table_entity(
         to_table_name = table_entity.name.root
 
         for from_table_name in lineage_parser.source_tables:
-            yield from _create_lineage_by_table_name(
-                metadata,
-                from_table=str(from_table_name),
-                to_table=f"{schema_name}.{to_table_name}",
-                service_name=service_name,
+            for procedure, source_table in get_source_table_names(
+                metadata=metadata,
+                dialect=dialect,
+                source_table=from_table_name,
                 database_name=database_name,
                 schema_name=schema_name,
-                masked_query=masked_query,
-                column_lineage_map=column_lineage,
-                lineage_source=lineage_source,
-            ) or []
+                service_name=service_name,
+                timeout_seconds=timeout_seconds,
+                column_lineage=column_lineage,
+            ):
+                yield from _create_lineage_by_table_name(
+                    metadata,
+                    from_table=str(source_table),
+                    to_table=f"{schema_name}.{to_table_name}",
+                    service_name=service_name,
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    masked_query=masked_query,
+                    column_lineage_map=column_lineage,
+                    lineage_source=lineage_source,
+                    procedure=procedure,
+                ) or []
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(

metadata/ingestion/models/patch_request.py

@@ -94,11 +94,13 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "tableConstraints": True,
     "tablePartition": True,
     "location": True,
+    "locationPath": True,
     "schemaDefinition": True,
     "sampleData": True,
     "fileFormat": True,
     # Stored Procedure Fields
     "storedProcedureCode": True,
+    "storedProcedureType": True,
     "code": True,
     # Dashboard Entity Fields
     "chartType": True,

metadata/ingestion/source/dashboard/mstr/client.py

@@ -37,7 +37,7 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 API_VERSION = "MicroStrategyLibrary/api"
-LOGIN_MODE_GUEST = 8
+LOGIN_MODE_GUEST = 1
 APPLICATION_TYPE = 35
 
 
@@ -77,27 +77,31 @@ class MSTRClient:
         To know about the data params below please visit
         https://demo.microstrategy.com/MicroStrategyLibrary/api-docs/index.html#/Authentication/postLogin
         """
-        try:
-            data = {
-                "username": self.config.username,
-                "password": self.config.password.get_secret_value(),
-                "loginMode": LOGIN_MODE_GUEST,
-                "applicationType": APPLICATION_TYPE,
-            }
+        data = {
+            "username": self.config.username,
+            "password": self.config.password.get_secret_value(),
+            "loginMode": LOGIN_MODE_GUEST,
+            "applicationType": APPLICATION_TYPE,
+        }
+        response = requests.post(
+            url=self._get_base_url("auth/login"), data=data, timeout=60
+        )
+        if not response.status_code == 204:
+            logger.error(
+                f"Auth Failed with login mode: {LOGIN_MODE_GUEST}, Error: {response.text}"
+            )
+            logger.error("Trying with login mode: 8")
+            data["loginMode"] = 8
             response = requests.post(
                 url=self._get_base_url("auth/login"), data=data, timeout=60
             )
-            if not response:
-                raise SourceConnectionException()
-            return AuthHeaderCookie(
-                auth_header=response.headers, auth_cookies=response.cookies
-            )
-        except Exception as exc:
-            logger.debug(traceback.format_exc())
-            logger.error(
-                f"Failed to fetch the auth header and cookies due to [{exc}], please validate credentials"
-            )
-        return None
+            if not response.status_code == 204:
+                raise SourceConnectionException(
+                    f"Auth Failed with login mode: 8 and {LOGIN_MODE_GUEST}, Error: {response.text}"
+                )
+        return AuthHeaderCookie(
+            auth_header=response.headers, auth_cookies=response.cookies
+        )
 
     def _set_api_session(self) -> bool:
         """

metadata/ingestion/source/database/databricks/metadata.py

@@ -125,7 +125,7 @@ def _get_table_columns(self, connection, table_name, schema, db_name):
         query = DATABRICKS_GET_TABLE_COMMENTS.format(
             database_name=db_name, schema_name=schema, table_name=table_name
         )
-        cursor = get_table_comment_result(
+        rows = get_table_comment_result(
             self,
             connection=connection,
             query=query,
@@ -134,8 +134,6 @@ def _get_table_columns(self, connection, table_name, schema, db_name):
             schema=schema,
         )
 
-        rows = cursor.fetchall()
-
     except exc.OperationalError as e:
         # Does the table exist?
         regex_fmt = r"TExecuteStatementResp.*SemanticException.*Table not found {}"

metadata/ingestion/source/database/snowflake/connection.py

@@ -220,7 +220,7 @@ def execute_inspector_func(engine_wrapper: SnowflakeEngineWrapper, func_name: st
     the function with name `func_name` and executes it
     """
     _init_database(engine_wrapper)
-    engine_wrapper.engine.execute(f"USE DATABASE {engine_wrapper.database_name}")
+    engine_wrapper.engine.execute(f'USE DATABASE "{engine_wrapper.database_name}"')
     inspector = inspect(engine_wrapper.engine)
     inspector_fn = getattr(inspector, func_name)
     inspector_fn()