openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

metadata/ingestion/api/topology_runner.py

@@ -13,14 +13,12 @@ Mixin to be used by service sources to dynamically
 generate the _run based on their topology.
 """
 import traceback
+from collections import defaultdict
 from functools import singledispatchmethod
-from typing import Any, Dict, Generic, Iterable, List, TypeVar, Union
+from typing import Any, Generic, Iterable, List, Type, TypeVar
 
 from pydantic import BaseModel
 
-from metadata.generated.schema.api.data.createStoredProcedure import (
-    CreateStoredProcedureRequest,
-)
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
@@ -34,25 +32,18 @@ from metadata.ingestion.models.topology import (
     ServiceTopology,
     TopologyContext,
     TopologyNode,
-    get_ctx_default,
     get_topology_node,
     get_topology_root,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.utils import model_str
-from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
-from metadata.utils.source_hash_utils import (
-    SOURCE_HASH_EXCLUDE_FIELDS,
-    generate_source_hash,
-)
+from metadata.utils.source_hash import generate_source_hash
 
 logger = ingestion_logger()
 
 C = TypeVar("C", bound=BaseModel)
 
-CACHED_ENTITIES = "cached_entities"
-
 
 class MissingExpectedEntityAckException(Exception):
     """
@@ -71,6 +62,9 @@ class TopologyRunnerMixin(Generic[C]):
     context: TopologyContext
     metadata: OpenMetadata
 
+    # The cache will have the shape {`child_stage.type_`: {`name`: `hash`}}
+    cache = defaultdict(dict)
+
     def process_nodes(self, nodes: List[TopologyNode]) -> Iterable[Entity]:
         """
         Given a list of nodes, either roots or children,
@@ -78,78 +72,116 @@ class TopologyRunnerMixin(Generic[C]):
 
         The execution tree is created in a depth-first fashion.
 
+        Note that this is used to handle the metadata ingestion for all our services. Therefore,
+        the first Node is always expected to be a Service. This is important because:
+        - Services (root Nodes) are flagged with `overwrite=False` -> In `yield_and_update_context` if
+          the stage is flagged as `overwrite=False`, we won't send any PUT/PATCH request to the API if
+          the service already exists.
+        - Then, when we iterate over Services' children (databases, pipelines, dashboards,...), we will
+          initialize a cache listing its children. This is used to to compare the fingerprint of the
+          stored entity vs. incoming entity and see if we need to:
+          1. Create a new entity (PUT) - if no fingerprint is found for the given name
+          2. Update some fields from the entity (PATCH) - if there's a fingerprint mismatch
+          3. Do nothing - if the fingerprints are the same.
+
+        The fingerprint is stored in the db in the field `sourceHash` in each entity.
+
         :param nodes: Topology Nodes to process
         :return: recursively build the execution tree
         """
         for node in nodes:
             logger.debug(f"Processing node {node}")
             node_producer = getattr(self, node.producer)
-            child_nodes = (
-                [get_topology_node(child, self.topology) for child in node.children]
-                if node.children
-                else []
-            )
+            child_nodes = self._get_child_nodes(node)
 
-            for element in node_producer() or []:
+            # Each node producer will give us a list of entities that we need
+            # to process. Each of the internal stages will sink result to OM API.
+            # E.g., in the DB topology, at the Table TopologyNode, the node_entity
+            # will be each `table`
+            for node_entity in node_producer() or []:
                 for stage in node.stages:
-                    logger.debug(f"Processing stage: {stage}")
-
-                    stage_fn = getattr(self, stage.processor)
-                    for entity_request in stage_fn(element) or []:
-                        try:
-                            # yield and make sure the data is updated
-                            yield from self.sink_request(
-                                stage=stage, entity_request=entity_request
-                            )
-                        except ValueError as err:
-                            logger.debug(traceback.format_exc())
-                            logger.warning(
-                                f"Unexpected value error when processing stage: [{stage}]: {err}"
-                            )
-
-                    # init the cache dict
-                    if stage.cache_entities:
-                        self._init_cache_dict(stage=stage, child_nodes=child_nodes)
-
-                # processing for all stages completed now cleaning the cache if applicable
+                    yield from self._process_stage(
+                        stage=stage, node_entity=node_entity, child_nodes=child_nodes
+                    )
+
+                # Once we are done processing all the stages,
                 for stage in node.stages:
-                    if stage.clear_cache:
-                        self.clear_context(stage=stage)
+                    if stage.clear_context:
+                        self.context.clear_stage(stage=stage)
 
                 # process all children from the node being run
                 yield from self.process_nodes(child_nodes)
 
-            if node.post_process:
-                logger.debug(f"Post processing node {node}")
-                for process in node.post_process:
-                    try:
-                        yield from self.check_context_and_handle(process)
-                    except Exception as exc:
-                        logger.debug(traceback.format_exc())
-                        logger.warning(
-                            f"Could not run Post Process `{process}` from Topology Runner -- {exc}"
-                        )
+            yield from self._run_node_post_process(node=node)
+
+    def _get_child_nodes(self, node: TopologyNode) -> List[TopologyNode]:
+        """Compute children nodes if any"""
+        return (
+            [get_topology_node(child, self.topology) for child in node.children]
+            if node.children
+            else []
+        )
 
-    def _init_cache_dict(self, stage: NodeStage, child_nodes: List[TopologyNode]):
+    def _process_stage(
+        self, stage: NodeStage, node_entity: Any, child_nodes: List[TopologyNode]
+    ) -> Iterable[Entity]:
         """
-        Method to call the API to fill the entities cache
+        For each entity produced in the Node Producer, iterate over all the Node's Stages and
+        yield the assets to pass down the workflow.
+
+        For each node_entity processed, we will cache - if needed - its children.
+        E.g., when processing DB Schemas, we will store its tables to compare the fingerprint
+        and decide if we need to PUT or PATCH at the sink.
         """
+        logger.debug(f"Processing stage: {stage}")
+
+        stage_fn = getattr(self, stage.processor)
+        for entity_request in stage_fn(node_entity) or []:
+            try:
+                # yield and make sure the data is updated
+                yield from self.sink_request(stage=stage, entity_request=entity_request)
+            except ValueError as err:
+                logger.debug(traceback.format_exc())
+                logger.warning(
+                    f"Unexpected value error when processing stage: [{stage}]: {err}"
+                )
+
+        if stage.cache_entities:
+            self._init_cache_dict(stage=stage, child_nodes=child_nodes)
 
-        if not self.context.__dict__.get(CACHED_ENTITIES):
-            self.context.__dict__[CACHED_ENTITIES] = {}
+    def _run_node_post_process(self, node: TopologyNode) -> Iterable[Entity]:
+        """
+        If the node has post_process steps, iterate over them and yield the result
+        """
+        if node.post_process:
+            logger.debug(f"Post processing node {node}")
+            for process in node.post_process:
+                try:
+                    node_post_process = getattr(self, process)
+                    for entity_request in node_post_process():
+                        yield entity_request
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Could not run Post Process `{process}` due to [{exc}]"
+                    )
+
+    def _init_cache_dict(
+        self, stage: NodeStage, child_nodes: List[TopologyNode]
+    ) -> None:
+        """
+        Method to call the API to fill the entities cache.
+
+        The cache will be part of the context
+        """
         for child_node in child_nodes or []:
             for child_stage in child_node.stages or []:
                 if child_stage.use_cache:
-                    entity_fqn = self.fqn_from_context(
+                    entity_fqn = self.context.fqn_from_stage(
                         stage=stage,
                         entity_name=self.context.__dict__[stage.context],
                     )
 
-                    if not self.context.__dict__[CACHED_ENTITIES].get(
-                        child_stage.type_
-                    ):
-                        self.context.__dict__[CACHED_ENTITIES][child_stage.type_] = {}
-
                     self.get_fqn_source_hash_dict(
                         parent_type=stage.type_,
                         child_type=child_stage.type_,
@@ -157,12 +189,11 @@ class TopologyRunnerMixin(Generic[C]):
                     )
 
     def get_fqn_source_hash_dict(
-        self, parent_type: Entity, child_type: Entity, entity_fqn: str
-    ) -> Dict:
+        self, parent_type: Type[Entity], child_type: Type[Entity], entity_fqn: str
+    ) -> None:
         """
         Get all the entities and store them as fqn:sourceHash in a dict
         """
-        params = {}
         if parent_type in (Database, DatabaseSchema):
             if child_type == StoredProcedure:
                 params = {"databaseSchema": entity_fqn}
@@ -177,24 +208,10 @@ class TopologyRunnerMixin(Generic[C]):
         )
         for entity in entities_list:
             if entity.sourceHash:
-                self.context.__dict__[CACHED_ENTITIES][child_type][
+                self.cache[child_type][
                     model_str(entity.fullyQualifiedName)
                 ] = entity.sourceHash
 
-    def check_context_and_handle(self, post_process: str):
-        """Based on the post_process step, check context and
-        evaluate if we can run it based on available class attributes
-
-        Args:
-            post_process: the name of the post_process step
-        """
-        if post_process == "mark_tables_as_deleted" and not self.context.database:
-            raise ValueError("No Database found in  `self.context`")
-
-        node_post_process = getattr(self, post_process)
-        for entity_request in node_post_process():
-            yield entity_request
-
     def _iter(self) -> Iterable[Either]:
         """
         This is the implementation for the entrypoint of our Source classes, which
@@ -207,56 +224,6 @@ class TopologyRunnerMixin(Generic[C]):
         """
         yield from self.process_nodes(get_topology_root(self.topology))
 
-    def _replace_context(self, key: str, value: Any) -> None:
-        """
-        Update the key of the context with the given value
-        :param key: element to update from the source context
-        :param value: value to use for the update
-        """
-        self.context.__dict__[key] = value
-
-    def _append_context(self, key: str, value: Any) -> None:
-        """
-        Update the key of the context with the given value
-        :param key: element to update from the source context
-        :param value: value to use for the update
-        """
-        self.context.__dict__[key].append(value)
-
-    def clear_context(self, stage: NodeStage) -> None:
-        """
-        Clear the available context
-        :param key: element to update from the source context
-        """
-        self.context.__dict__[stage.context] = get_ctx_default(stage)
-
-    def fqn_from_context(self, stage: NodeStage, entity_name: str) -> str:
-        """
-        Read the context
-        :param stage: Topology node being processed
-        :param entity_request: Request sent to the sink
-        :return: Entity FQN derived from context
-        """
-        context_names = [
-            self.context.__dict__[dependency]
-            for dependency in stage.consumer or []  # root nodes do not have consumers
-        ]
-        return fqn._build(  # pylint: disable=protected-access
-            *context_names, entity_name
-        )
-
-    def update_context(
-        self, stage: NodeStage, context: Union[str, OMetaTagAndClassification]
-    ):
-        """Append or update context"""
-        # We'll store the entity_name in the topology context instead of the entity_fqn
-        # and build the fqn on the fly wherever required.
-        # This is mainly because we need the context in other places
-        if stage.context and not stage.cache_all:
-            self._replace_context(key=stage.context, value=context)
-        if stage.context and stage.cache_all:
-            self._append_context(key=stage.context, value=context)
-
     def create_patch_request(
         self, original_entity: Entity, create_request: C
     ) -> PatchRequest:
@@ -284,30 +251,30 @@ class TopologyRunnerMixin(Generic[C]):
         """
         entity = None
         entity_name = model_str(right.name)
-        entity_fqn = self.fqn_from_context(stage=stage, entity_name=entity_name)
+        entity_fqn = self.context.fqn_from_stage(stage=stage, entity_name=entity_name)
 
-        # we get entity from OM if we do not want to overwrite existing data in OM
+        # If we don't want to write data in OM, we'll return what we fetch from the API.
         # This will be applicable for service entities since we do not want to overwrite the data
+        same_fingerprint = False
         if not stage.overwrite and not self._is_force_overwrite_enabled():
             entity = self.metadata.get_by_name(
                 entity=stage.type_,
                 fqn=entity_fqn,
-                fields=["*"],  # Get all the available data from the Entity
+                fields=["*"],
             )
+            if entity:
+                same_fingerprint = True
+
         create_entity_request_hash = generate_source_hash(
             create_request=entity_request.right,
-            exclude_fields=SOURCE_HASH_EXCLUDE_FIELDS,
         )
 
         if hasattr(entity_request.right, "sourceHash"):
             entity_request.right.sourceHash = create_entity_request_hash
 
-        skip_processing_entity = False
         if entity is None and stage.use_cache:
             # check if we find the entity in the entities list
-            entity_source_hash = self.context.__dict__[CACHED_ENTITIES][
-                stage.type_
-            ].get(entity_fqn)
+            entity_source_hash = self.cache[stage.type_].get(entity_fqn)
             if entity_source_hash:
                 # if the source hash is present, compare it with new hash
                 if entity_source_hash != create_entity_request_hash:
@@ -315,7 +282,7 @@ class TopologyRunnerMixin(Generic[C]):
                     entity = self.metadata.get_by_name(
                         entity=stage.type_,
                         fqn=entity_fqn,
-                        fields=["*"],  # Get all the available data from the Entity
+                        fields=["*"],
                     )
 
                     # we return the entity for a patch update
@@ -329,9 +296,9 @@ class TopologyRunnerMixin(Generic[C]):
                     logger.debug(
                         f"No changes detected for {str(stage.type_.__name__)} '{entity_fqn}'"
                     )
-                    skip_processing_entity = True
+                    same_fingerprint = True
 
-        if not skip_processing_entity:
+        if not same_fingerprint:
             # We store the generated source hash and yield the request
 
             yield entity_request
@@ -357,7 +324,7 @@ class TopologyRunnerMixin(Generic[C]):
                     "for the service connection."
                 )
 
-        self.update_context(stage=stage, context=entity_name)
+        self.context.update_context_name(stage=stage, right=right)
 
     @yield_and_update_context.register
     def _(
@@ -373,7 +340,7 @@ class TopologyRunnerMixin(Generic[C]):
         lineage has been properly drawn. We'll skip the process for now.
         """
         yield entity_request
-        self.update_context(stage=stage, context=right.edge.fromEntity.name.__root__)
+        self.context.update_context_name(stage=stage, right=right.edge.fromEntity)
 
     @yield_and_update_context.register
     def _(
@@ -382,11 +349,16 @@ class TopologyRunnerMixin(Generic[C]):
         stage: NodeStage,
         entity_request: Either[C],
     ) -> Iterable[Either[Entity]]:
-        """Tag implementation for the context information"""
+        """
+        Tag implementation for the context information.
+
+        We need the full OMetaTagAndClassification in the context
+        to build the TagLabels during the ingestion. We need to bundle
+        both CreateClassificationRequest and CreateTagRequest.
+        """
         yield entity_request
 
-        # We'll keep the tag fqn in the context and use if required
-        self.update_context(stage=stage, context=right)
+        self.context.update_context_value(stage=stage, value=right)
 
     @yield_and_update_context.register
     def _(
@@ -398,30 +370,7 @@ class TopologyRunnerMixin(Generic[C]):
         """Custom Property implementation for the context information"""
         yield entity_request
 
-        # We'll keep the tag fqn in the context and use if required
-        self.update_context(stage=stage, context=right)
-
-    @yield_and_update_context.register
-    def _(
-        self,
-        right: CreateStoredProcedureRequest,
-        stage: NodeStage,
-        entity_request: Either[C],
-    ) -> Iterable[Either[Entity]]:
-        """Tag implementation for the context information"""
-        yield entity_request
-
-        procedure_fqn = fqn.build(
-            metadata=self.metadata,
-            entity_type=StoredProcedure,
-            service_name=self.context.database_service,
-            database_name=self.context.database,
-            schema_name=self.context.database_schema,
-            procedure_name=right.name.__root__,
-        )
-
-        # We'll keep the tag fqn in the context and use if required
-        self.update_context(stage=stage, context=procedure_fqn)
+        self.context.update_context_value(stage=stage, value=right)
 
     def sink_request(
         self, stage: NodeStage, entity_request: Either[C]

metadata/ingestion/bulksink/metadata_usage.py

@@ -35,11 +35,13 @@ from metadata.generated.schema.entity.data.table import (
     Table,
     TableJoins,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle
 from metadata.generated.schema.type.tableUsageCount import TableColumn, TableUsageCount
 from metadata.generated.schema.type.usageRequest import UsageRequest
-from metadata.ingestion.api.models import StackTraceError
 from metadata.ingestion.api.steps import BulkSink
 from metadata.ingestion.lineage.sql_lineage import (
     get_column_fqn,
@@ -86,6 +88,10 @@ class MetadataUsageBulkSink(BulkSink):
         self.table_usage_map = {}
         self.today = datetime.today().strftime("%Y-%m-%d")
 
+    @property
+    def name(self) -> str:
+        return "OpenMetadata"
+
     @classmethod
     def create(cls, config_dict: dict, metadata: OpenMetadata):
         config = MetadataUsageSinkConfig.parse_obj(config_dict)
@@ -147,7 +153,7 @@ class MetadataUsageBulkSink(BulkSink):
                     StackTraceError(
                         name=value_dict["table_entity"].fullyQualifiedName.__root__,
                         error=f"Failed to update usage for {name} :{exc}",
-                        stack_trace=traceback.format_exc(),
+                        stackTrace=traceback.format_exc(),
                     )
                 )
 
@@ -240,7 +246,7 @@ class MetadataUsageBulkSink(BulkSink):
                         StackTraceError(
                             name=table_usage.table,
                             error=error,
-                            stack_trace=traceback.format_exc(),
+                            stackTrace=traceback.format_exc(),
                         )
                     )
                 except Exception as exc:
@@ -252,7 +258,7 @@ class MetadataUsageBulkSink(BulkSink):
                     logger.warning(error)
                     self.status.failed(
                         StackTraceError(
-                            name=name, error=error, stack_trace=traceback.format_exc()
+                            name=name, error=error, stackTrace=traceback.format_exc()
                         )
                     )
             else:
@@ -379,7 +385,7 @@ class MetadataUsageBulkSink(BulkSink):
                 StackTraceError(
                     name=table_usage.table,
                     error=error,
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

metadata/ingestion/lineage/models.py

@@ -17,6 +17,9 @@ from typing import Dict
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaType,
 )
+from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
+    AzureSQLType,
+)
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigqueryType,
 )
@@ -105,6 +108,7 @@ MAP_CONNECTION_TYPE_DIALECT: Dict[str, Dialect] = {
     str(DeltaLakeType.DeltaLake.value): Dialect.SPARKSQL,
     str(SQLiteType.SQLite.value): Dialect.SQLITE,
     str(MssqlType.Mssql.value): Dialect.TSQL,
+    str(AzureSQLType.AzureSQL.value): Dialect.TSQL,
 }
 
 

metadata/ingestion/lineage/parser.py

@@ -19,10 +19,10 @@ from typing import Any, Dict, List, Optional, Tuple
 
 import sqlparse
 from cached_property import cached_property
-from sqllineage import SQLPARSE_DIALECT
-from sqllineage.core.models import Column, Table
-from sqllineage.exceptions import SQLLineageException
-from sqllineage.runner import LineageRunner
+from collate_sqllineage import SQLPARSE_DIALECT
+from collate_sqllineage.core.models import Column, Table
+from collate_sqllineage.exceptions import SQLLineageException
+from collate_sqllineage.runner import LineageRunner
 from sqlparse.sql import Comparison, Identifier, Parenthesis, Statement
 
 from metadata.generated.schema.type.tableUsageCount import TableColumn, TableColumnJoin
@@ -98,24 +98,30 @@ class LineageParser:
         """
         Get a list of intermediate tables
         """
-        # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
-        return self.retrieve_tables(self.parser.intermediate_tables)
+        if self.parser:
+            # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
+            return self.retrieve_tables(self.parser.intermediate_tables)
+        return []
 
     @cached_property
     def source_tables(self) -> List[Table]:
         """
         Get a list of source tables
         """
-        # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
-        return self.retrieve_tables(self.parser.source_tables)
+        if self.parser:
+            # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
+            return self.retrieve_tables(self.parser.source_tables)
+        return []
 
     @cached_property
     def target_tables(self) -> List[Table]:
         """
         Get a list of target tables
         """
-        # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
-        return self.retrieve_tables(self.parser.target_tables)
+        if self.parser:
+            # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
+            return self.retrieve_tables(self.parser.target_tables)
+        return []
 
     # pylint: disable=protected-access
     @cached_property
@@ -124,6 +130,8 @@ class LineageParser:
         Get a list of tuples of column lineage
         """
         column_lineage = []
+        if self.parser is None:
+            return []
         try:
             if self.parser._dialect == SQLPARSE_DIALECT:
                 return self.parser.get_column_lineage()
@@ -311,7 +319,10 @@ class LineageParser:
             )
 
             if not table_left or not table_right:
-                logger.warning(f"Cannot find ingredients from {comparison}")
+                logger.warning(
+                    f"Can't extract table names when parsing JOIN information from {comparison}"
+                )
+                logger.debug(f"Query: {sql_statement}")
                 continue
 
             left_table_column = TableColumn(table=table_left, column=column_left)
@@ -331,6 +342,8 @@ class LineageParser:
         :return: for each table name, list all joins against other tables
         """
         join_data = defaultdict(list)
+        if self.parser is None:
+            return join_data
         # These are @lazy_property, not properly being picked up by IDEs. Ignore the warning
         for statement in self.parser.statements():
             self.stateful_add_joins_from_statement(join_data, sql_statement=statement)
@@ -378,7 +391,10 @@ class LineageParser:
     @staticmethod
     def _evaluate_best_parser(
         query: str, dialect: Dialect, timeout_seconds: int
-    ) -> LineageRunner:
+    ) -> Optional[LineageRunner]:
+        if query is None:
+            return None
+
         @timeout(seconds=timeout_seconds)
         def get_sqlfluff_lineage_runner(qry: str, dlct: str) -> LineageRunner:
             lr_dialect = LineageRunner(qry, dialect=dlct)

metadata/ingestion/lineage/sql_lineage.py

@@ -16,6 +16,9 @@ from typing import Any, Iterable, List, Optional, Tuple
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.type.entityLineage import (
     ColumnLineage,
     EntitiesEdge,
@@ -23,7 +26,7 @@ from metadata.generated.schema.type.entityLineage import (
 )
 from metadata.generated.schema.type.entityLineage import Source as LineageSource
 from metadata.generated.schema.type.entityReference import EntityReference
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.models import Dialect
 from metadata.ingestion.lineage.parser import LINEAGE_PARSING_TIMEOUT, LineageParser
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -183,19 +186,8 @@ def get_table_entities_from_query(
     table_entities = search_table_entities(
         metadata=metadata,
         service_name=service_name,
-        database=database_query,
-        database_schema=schema_query,
-        table=table,
-    )
-
-    if table_entities:
-        return table_entities
-
-    table_entities = search_table_entities(
-        metadata=metadata,
-        service_name=service_name,
-        database=database_name,
-        database_schema=database_schema,
+        database=database_query if database_query else database_name,
+        database_schema=schema_query if schema_query else database_schema,
         table=table,
     )
 
@@ -337,7 +329,7 @@ def _create_lineage_by_table_name(
             left=StackTraceError(
                 name="Lineage",
                 error=f"Error creating lineage for service [{service_name}] from table [{from_table}]: {exc}",
-                stack_trace=traceback.format_exc(),
+                stackTrace=traceback.format_exc(),
             )
         )
 
@@ -440,7 +432,7 @@ def get_lineage_by_query(
             left=StackTraceError(
                 name="Lineage",
                 error=f"Ingesting lineage failed for service [{service_name}]: {exc}",
-                stack_trace=traceback.format_exc(),
+                stackTrace=traceback.format_exc(),
             )
         )
 
@@ -481,6 +473,6 @@ def get_lineage_via_table_entity(
             left=StackTraceError(
                 name="Lineage",
                 error=f"Failed to create view lineage for database [{database_name}] and table [{table_entity}]: {exc}",
-                stack_trace=traceback.format_exc(),
+                stackTrace=traceback.format_exc(),
             )
         )