openmetadata-ingestion 1.6.0.0rc2__py3-none-any.whl → 1.6.1.0__py3-none-any.whl

metadata/ingestion/ometa/mixins/ingestion_pipeline_mixin.py

@@ -22,6 +22,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
 )
 from metadata.ingestion.api.parser import parse_ingestion_pipeline_config_gracefully
 from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.utils import quote
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -46,7 +47,7 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status: Pipeline Status data to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=pipeline_status.model_dump_json(),
         )
         logger.debug(
@@ -64,7 +65,8 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status_run_id: Pipeline Status run id
         """
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus/{pipeline_status_run_id}"
+            f"{self.get_suffix(IngestionPipeline)}/"
+            f"{quote(ingestion_pipeline_fqn)}/pipelineStatus/{pipeline_status_run_id}"
         )
         if resp:
             return PipelineStatus(**resp)
@@ -99,7 +101,7 @@ class OMetaIngestionPipelineMixin:
         params = {"startTs": start_ts, "endTs": end_ts}
 
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=params,
         )
 

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -19,7 +19,7 @@ from metadata.generated.schema.entity.feed.suggestion import Suggestion, Suggest
 from metadata.generated.schema.type import basic
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.ingestion.ometa.client import REST
-from metadata.ingestion.ometa.utils import model_str
+from metadata.ingestion.ometa.utils import model_str, quote
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -65,7 +65,7 @@ class OMetaSuggestionsMixin:
         self.client.put(
             f"{self.get_suffix(Suggestion)}/accept-all?"
             f"userId={model_str(user_id)}&"
-            f"entityFQN={model_str(fqn)}&"
+            f"entityFQN={quote(fqn)}&"
             f"suggestionType={suggestion_type.value}",
         )
 
@@ -79,6 +79,6 @@ class OMetaSuggestionsMixin:
         self.client.put(
             f"{self.get_suffix(Suggestion)}/reject-all?"
             f"userId={model_str(user_id)}&"
-            f"entityFQN={model_str(fqn)}&"
+            f"entityFQN={quote(fqn)}&"
             f"suggestionType={suggestion_type.value}",
         )

metadata/ingestion/ometa/mixins/table_mixin.py

@@ -257,7 +257,7 @@ class OMetaTableMixin:
         profile_type_url = profile_type.__name__[0].lower() + profile_type.__name__[1:]
 
         resp = self.client.get(
-            f"{self.get_suffix(Table)}/{fqn}/{profile_type_url}?limit={limit}{url_after}",
+            f"{self.get_suffix(Table)}/{quote(fqn)}/{profile_type_url}?limit={limit}{url_after}",
             data={"startTs": start_ts, "endTs": end_ts},
         )
 

metadata/ingestion/source/dashboard/dashboard_service.py

@@ -479,6 +479,7 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
         to_entity: Union[Dashboard, DashboardDataModel],
         from_entity: Union[Table, DashboardDataModel, Dashboard],
         column_lineage: List[ColumnLineage] = None,
+        sql: Optional[str] = None,
     ) -> Optional[Either[AddLineageRequest]]:
         if from_entity and to_entity:
             return Either(
@@ -494,6 +495,7 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
                         ),
                         lineageDetails=LineageDetails(
                             source=LineageSource.DashboardLineage,
+                            sqlQuery=sql,
                             columnsLineage=column_lineage,
                         ),
                     )

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

@@ -37,12 +37,13 @@ from metadata.utils.logger import ingestion_logger
 logger = ingestion_logger()
 
 API_VERSION = "MicroStrategyLibrary/api"
+HEADERS = {"Content-Type": "application/json"}
 APPLICATION_TYPE = 35
 
 
 class MicroStrategyClient:
     """
-    Client Handling API communication with Metabase
+    Client Handling API communication with MicroStrategy
     """
 
     def _get_base_url(self, path=None):
@@ -56,18 +57,43 @@ class MicroStrategyClient:
     ):
         self.config = config
 
-        self.auth_params: AuthHeaderCookie = self._get_auth_header_and_cookies()
+        self.auth_params = self._get_auth_header_and_cookies()
 
         client_config = ClientConfig(
             base_url=clean_uri(str(self.config.hostPort)),
             api_version=API_VERSION,
-            extra_headers=self.auth_params.auth_header,
+            extra_headers=self.auth_params.auth_header if self.auth_params else None,
             allow_redirects=True,
-            cookies=self.auth_params.auth_cookies,
+            cookies=self.auth_params.auth_cookies if self.auth_params else None,
         )
 
         self.client = REST(client_config)
-        self._set_api_session()
+
+    def get_auth_params(self) -> AuthHeaderCookie:
+        """
+        Test whether we can fetch auth_token from the api
+        """
+        data = {
+            "username": self.config.username,
+            "password": self.config.password.get_secret_value(),
+            "loginMode": int(self.config.loginMode),
+            "applicationType": APPLICATION_TYPE,
+        }
+        response = requests.post(
+            url=self._get_base_url("auth/login"), json=data, headers=HEADERS, timeout=60
+        )
+        response.raise_for_status()
+        if (
+            not response.ok
+            or response.status_code != 204
+            or "X-MSTR-AuthToken" not in response.headers
+        ):
+            raise SourceConnectionException(
+                f"Failed to Fetch Token, please validate your credentials and login_mode : {response.text}"
+            )
+        return AuthHeaderCookie(
+            auth_header=response.headers, auth_cookies=response.cookies
+        )
 
     def _get_auth_header_and_cookies(self) -> Optional[AuthHeaderCookie]:
         """
@@ -77,35 +103,25 @@ class MicroStrategyClient:
         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": self.config.loginMode,
-                "applicationType": APPLICATION_TYPE,
-            }
-            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
-            )
+            auth_data = self.get_auth_params()
+            if auth_data:
+                self._set_api_session(auth_data)
+                return auth_data
         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"
+                f"Failed to fetch the auth header and cookies due to : [{exc}]"
             )
         return None
 
-    def _set_api_session(self) -> bool:
+    def _set_api_session(self, auth_data: AuthHeaderCookie) -> bool:
         """
         Set the user api session to active this will keep the connection alive
         """
         api_session = requests.put(
             url=self._get_base_url("sessions"),
-            headers=self.auth_params.auth_header,
-            cookies=self.auth_params.auth_cookies,
+            headers=auth_data.auth_header,
+            cookies=auth_data.auth_cookies,
             timeout=60,
         )
         if api_session.ok:

metadata/ingestion/source/dashboard/microstrategy/connection.py

@@ -46,7 +46,10 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"GetProjects": client.get_projects_list}
+    test_fn = {
+        "GetToken": client.get_auth_params,
+        "GetProjects": client.get_projects_list,
+    }
 
     return test_connection_steps(
         metadata=metadata,

metadata/ingestion/source/dashboard/tableau/metadata.py

@@ -69,6 +69,7 @@ from metadata.ingestion.source.dashboard.tableau.models import (
     ChartUrl,
     DataSource,
     DatasourceField,
+    TableAndQuery,
     TableauDashboard,
     TableauTag,
     UpstreamTable,
@@ -149,7 +150,7 @@ class TableauSource(DashboardServiceSource):
         return None
 
     @staticmethod
-    def _get_data_models_tags(dataModels: [DataSource]) -> Set[str]:
+    def _get_data_models_tags(dataModels: List[DataSource]) -> Set[str]:
         """
         Get the tags from the data model in the upstreamDatasources
         """
@@ -382,9 +383,12 @@ class TableauSource(DashboardServiceSource):
                         column=column.id,
                     )
                     for to_column in to_columns:
-                        column_lineage.append(
-                            ColumnLineage(fromColumns=[from_column], toColumn=to_column)
-                        )
+                        if from_column and to_column:
+                            column_lineage.append(
+                                ColumnLineage(
+                                    fromColumns=[from_column], toColumn=to_column
+                                )
+                            )
             return column_lineage
         except Exception as exc:
             logger.debug(f"Error to get column lineage: {exc}")
@@ -451,17 +455,22 @@ class TableauSource(DashboardServiceSource):
                 column.id
                 for field in upstream_data_model.fields
                 for column in field.upstreamColumns
+                if column is not None
             }
             for table in datamodel.upstreamTables or []:
                 om_tables = self._get_database_tables(db_service_entity, table)
-                for om_table in om_tables or []:
+                for om_table_and_query in om_tables or []:
                     column_lineage = self._get_column_lineage(
-                        table, om_table, upstream_data_model_entity, upstream_col_set
+                        table,
+                        om_table_and_query.table,
+                        upstream_data_model_entity,
+                        upstream_col_set,
                     )
                     yield self._get_add_lineage_request(
                         to_entity=upstream_data_model_entity,
-                        from_entity=om_table,
+                        from_entity=om_table_and_query.table,
                         column_lineage=column_lineage,
+                        sql=om_table_and_query.query,
                     )
         except Exception as err:
             yield Either(
@@ -698,7 +707,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_table_entities_from_api(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -734,7 +743,7 @@ class TableauSource(DashboardServiceSource):
                     fqn=table_fqn,
                 )
                 if table_entity:
-                    return [table_entity]
+                    return [TableAndQuery(table=table_entity)]
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Error to get tables for lineage using GraphQL Apis: {exc}")
@@ -742,7 +751,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_table_entities_from_query(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         In case we get the table details from the Graphql APIs we process them
         """
@@ -778,7 +787,12 @@ class TableauSource(DashboardServiceSource):
                         database_schema=schema_name,
                         table=table_name,
                     )
-                    tables_list.extend(from_entities)
+                    tables_list.extend(
+                        [
+                            TableAndQuery(table=table, query=custom_sql_table.query)
+                            for table in from_entities
+                        ]
+                    )
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
@@ -787,7 +801,7 @@ class TableauSource(DashboardServiceSource):
 
     def _get_database_tables(
         self, db_service_entity: DatabaseService, table: UpstreamTable
-    ) -> Optional[List[Table]]:
+    ) -> Optional[List[TableAndQuery]]:
         """
         Get the table entities for lineage
         """
@@ -879,6 +893,9 @@ class TableauSource(DashboardServiceSource):
         try:
             return dashboard_details.project.name
         except Exception as exc:
+            logger.info(
+                f"Cannot parse project name for dashboard:{dashboard_details.id} from Tableau server"
+            )
             logger.debug(traceback.format_exc())
             logger.warning(
                 f"Error fetching project name for {dashboard_details.id}: {exc}"

metadata/ingestion/source/dashboard/tableau/models.py

@@ -18,6 +18,7 @@ from typing import Any, Dict, List, Optional, Union
 from pydantic import BaseModel, ConfigDict, Field, validator
 
 from metadata.generated.schema.entity.data.chart import ChartType
+from metadata.generated.schema.entity.data.table import Table
 
 
 class TableauBaseModel(BaseModel):
@@ -172,3 +173,12 @@ class TableauDashboard(TableauBaseModel):
     webpageUrl: Optional[str] = None
     charts: Optional[List[TableauChart]] = None
     dataModels: List[DataSource] = []
+
+
+class TableAndQuery(BaseModel):
+    """
+    Wrapper class for Table entity and associated Query for lineage
+    """
+
+    table: Table
+    query: Optional[str] = None

metadata/ingestion/source/database/bigquery/profiler/system.py

@@ -13,6 +13,7 @@ from metadata.profiler.metrics.system.system import (
     SQASessionProvider,
     SystemMetricsComputer,
 )
+from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.logger import profiler_logger
 from metadata.utils.time_utils import datetime_to_timestamp
 
@@ -25,11 +26,11 @@ class BigQuerySystemMetricsSource(
     """BigQuery system metrics source class"""
 
     def get_kwargs(self, **kwargs):
-        table = kwargs.get("table")
+        runner: QueryRunner = kwargs.get("runner")
         return {
-            "table": table.__table__.name,
-            "dataset_id": table.__table_args__["schema"],
-            "project_id": super().get_session().get_bind().url.host,
+            "table": runner.table_name,
+            "database": runner.session.get_bind().url.database,
+            "schema": runner.schema_name,
             "usage_location": kwargs.get("usage_location"),
         }
 

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

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from pathlib import Path
 from typing import Optional
 
 from sqlalchemy.engine import Engine
@@ -32,13 +33,25 @@ from metadata.ingestion.connections.builders import (
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.utils.constants import THREE_MIN
+from metadata.utils.constants import THREE_MIN, UTF_8
 
 
 def get_connection(connection: Db2Connection) -> Engine:
     """
     Create connection
     """
+    # prepare license
+    # pylint: disable=import-outside-toplevel
+    if connection.license and connection.licenseFileName:
+        import clidriver
+
+        with open(
+            Path(clidriver.__path__[0], "license", connection.licenseFileName),
+            "w",
+            encoding=UTF_8,
+        ) as file:
+            file.write(connection.license)
+
     return create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url_common,

metadata/ingestion/source/database/dbt/constants.py

@@ -22,6 +22,14 @@ REQUIRED_MANIFEST_KEYS = ["name", "schema", "resource_type"]
 # Based on https://schemas.getdbt.com/dbt/catalog/v1.json
 REQUIRED_CATALOG_KEYS = ["name", "type", "index"]
 
+REQUIRED_CONSTRAINT_KEYS = [
+    "type",
+    "name",
+    "expression",
+    "warn_unenforced",
+    "warn_unsupported",
+]
+
 REQUIRED_RESULTS_KEYS = {
     "status",
     "timing",

metadata/ingestion/source/database/dbt/dbt_config.py

@@ -176,7 +176,11 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
         logger.debug(
             "Requesting [dbt_catalog], [dbt_manifest] and [dbt_run_results] data"
         )
-        params_data = {"order_by": "-finished_at", "limit": "1"}
+        params_data = {
+            "order_by": "-finished_at",
+            "limit": "1",
+            "status__in": "[10,20]",
+        }
         if project_id:
             params_data["project_id"] = project_id
 
@@ -194,7 +198,7 @@ def _(config: DbtCloudConfig):  # pylint: disable=too-many-locals
             last_run = runs_data[0]
             run_id = last_run["id"]
             logger.info(
-                f"Retrieved last successful run [{str(run_id)}]: "
+                f"Retrieved last completed run [{str(run_id)}]: "
                 f"Finished {str(last_run['finished_at_humanized'])} (duration: {str(last_run['duration_humanized'])})"
             )
             try:

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -43,6 +43,7 @@ from metadata.ingestion.models.topology import (
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.ingestion.source.database.dbt.constants import (
+    REQUIRED_CONSTRAINT_KEYS,
     REQUIRED_NODE_KEYS,
     REQUIRED_RESULTS_KEYS,
 )
@@ -187,6 +188,20 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
                 ]
                 for key in keys_to_delete:
                     del value[key]
+                if value.get("columns"):
+                    for col_name, value in value[
+                        "columns"
+                    ].items():  # pylint: disable=unused-variable
+                        if value.get("constraints"):
+                            keys_to_delete = [
+                                key
+                                for key in value
+                                if key.lower() not in REQUIRED_CONSTRAINT_KEYS
+                            ]
+                            for key in keys_to_delete:
+                                del value[key]
+                        else:
+                            value["constraints"] = None
 
     def remove_run_result_non_required_keys(self, run_results: List[dict]):
         """
@@ -220,7 +235,9 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             if self.context.get().dbt_file.dbt_catalog
             else None,
             dbt_manifest=parse_manifest(self.context.get().dbt_file.dbt_manifest),
-            dbt_sources=parse_sources(self.context.get().dbt_file.dbt_sources),
+            dbt_sources=parse_sources(self.context.get().dbt_file.dbt_sources)
+            if self.context.get().dbt_file.dbt_sources
+            else None,
             dbt_run_results=[
                 parse_run_results(run_result_file)
                 for run_result_file in self.context.get().dbt_file.dbt_run_results

metadata/ingestion/source/database/mssql/service_spec.py

@@ -1,10 +1,12 @@
 from metadata.ingestion.source.database.mssql.lineage import MssqlLineageSource
 from metadata.ingestion.source.database.mssql.metadata import MssqlSource
 from metadata.ingestion.source.database.mssql.usage import MssqlUsageSource
+from metadata.sampler.sqlalchemy.mssql.sampler import MssqlSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
     metadata_source_class=MssqlSource,
     lineage_source_class=MssqlLineageSource,
     usage_source_class=MssqlUsageSource,
+    sampler_class=MssqlSampler,
 )

metadata/ingestion/source/database/redshift/profiler/system.py

@@ -18,6 +18,7 @@ from metadata.profiler.metrics.system.system import (
     SQASessionProvider,
     SystemMetricsComputer,
 )
+from metadata.profiler.processor.runner import QueryRunner
 from metadata.utils.logger import profiler_logger
 from metadata.utils.profiler_utils import QueryResult
 from metadata.utils.time_utils import datetime_to_timestamp
@@ -45,11 +46,11 @@ class RedshiftSystemMetricsSource(
         return get_metric_result(queries, table)
 
     def get_kwargs(self, **kwargs):
-        table = kwargs.get("table")
+        runner: QueryRunner = kwargs.get("runner")
         return {
-            "table": table.__table__.name,
-            "database": self.get_session().get_bind().url.database,
-            "schema": table.__table__.schema,
+            "table": runner.table_name,
+            "database": runner.session.get_bind().url.database,
+            "schema": runner.schema_name,
         }
 
     def get_deletes(self, **kwargs) -> List[SystemProfile]:

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

@@ -121,7 +121,7 @@ def get_connection(connection: SnowflakeConnection) -> Engine:
             )
         p_key = serialization.load_pem_private_key(
             bytes(connection.privateKey.get_secret_value(), "utf-8"),
-            password=snowflake_private_key_passphrase.encode(),
+            password=snowflake_private_key_passphrase.encode() or None,
             backend=default_backend(),
         )
         pkb = p_key.private_bytes(
@@ -137,11 +137,16 @@ def get_connection(connection: SnowflakeConnection) -> Engine:
             "client_session_keep_alive"
         ] = connection.clientSessionKeepAlive
 
-    return create_generic_db_connection(
+    engine = create_generic_db_connection(
         connection=connection,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )
+    if connection.connectionArguments.root and connection.connectionArguments.root.get(
+        "private_key"
+    ):
+        del connection.connectionArguments.root["private_key"]
+    return engine
 
 
 def test_connection(

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

@@ -12,8 +12,10 @@
 Snowflake lineage module
 """
 
-from typing import Dict, List
+import traceback
+from typing import Dict, Iterator, List
 
+from metadata.generated.schema.type.tableQuery import TableQuery
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_STORED_PROCEDURE_QUERIES,
@@ -27,6 +29,9 @@ from metadata.ingestion.source.database.stored_procedures_mixin import (
     StoredProcedureLineageMixin,
 )
 from metadata.utils.helpers import get_start_and_end
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class SnowflakeLineageSource(
@@ -61,3 +66,32 @@ class SnowflakeLineageSource(
         )
 
         return queries_dict
+
+    def yield_table_query(self) -> Iterator[TableQuery]:
+        """
+        Given an engine, iterate over the query results to
+        yield a TableQuery with query parsing info
+        """
+        for engine in self.get_engine():
+            rows = []
+            with engine.connect() as conn:
+                rows = conn.execution_options(
+                    stream_results=True, max_row_buffer=100
+                ).execute(
+                    self.get_sql_statement(start_time=self.start, end_time=self.end)
+                )
+            # exit from active connection after fetching rows & during
+            # further process of `yield_query_lineage`
+            for row in rows:
+                query_dict = dict(row)
+                try:
+                    yield TableQuery(
+                        dialect=self.dialect.value,
+                        query=query_dict["query_text"],
+                        databaseName=self.get_database_name(query_dict),
+                        serviceName=self.config.serviceName,
+                        databaseSchema=self.get_schema_name(query_dict),
+                    )
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(f"Error processing query_dict {query_dict}: {exc}")