openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

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

@@ -13,6 +13,7 @@
 DBT source methods.
 """
 import traceback
+from copy import deepcopy
 from datetime import datetime
 from typing import Any, Iterable, List, Optional, Union
 
@@ -324,7 +325,41 @@ class DbtSource(DbtServiceSource):
             None,
         )
 
-    # pylint: disable=too-many-locals, too-many-branches
+    def _add_dbt_freshness_test_from_sources(
+        self, key: str, manifest_node, manifest_entities, dbt_objects: DbtObjects
+    ):
+        # in dbt manifest sources node name is table/view name (not test name like with test nodes)
+        # so in order for the test creation to be named precisely I am amending manifest node name within it's deepcopy
+        manifest_node_new = deepcopy(manifest_node)
+        manifest_node_new.name = manifest_node_new.name + "_freshness"
+
+        freshness_test_result = next(
+            (item for item in dbt_objects.dbt_sources.results if item.unique_id == key),
+            None,
+        )
+
+        if freshness_test_result:
+            self.context.get().dbt_tests[key + "_freshness"] = {
+                DbtCommonEnum.MANIFEST_NODE.value: manifest_node_new
+            }
+            self.context.get().dbt_tests[key + "_freshness"][
+                DbtCommonEnum.UPSTREAM.value
+            ] = self.parse_upstream_nodes(manifest_entities, manifest_node)
+            self.context.get().dbt_tests[key + "_freshness"][
+                DbtCommonEnum.RESULTS.value
+            ] = freshness_test_result
+
+    def add_dbt_sources(
+        self, key: str, manifest_node, manifest_entities, dbt_objects: DbtObjects
+    ) -> None:
+        """
+        Method to append dbt test cases based on sources file for later processing
+        """
+        self._add_dbt_freshness_test_from_sources(
+            key, manifest_node, manifest_entities, dbt_objects
+        )
+
+    # pylint: disable=too-many-locals, too-many-branches, too-many-statements
     def yield_data_models(
         self, dbt_objects: DbtObjects
     ) -> Iterable[Either[DataModelLink]]:
@@ -376,6 +411,17 @@ class DbtSource(DbtServiceSource):
                         )
                         continue
 
+                    if (
+                        dbt_objects.dbt_sources
+                        and resource_type == SkipResourceTypeEnum.SOURCE.value
+                    ):
+                        self.add_dbt_sources(
+                            key,
+                            manifest_node=manifest_node,
+                            manifest_entities=manifest_entities,
+                            dbt_objects=dbt_objects,
+                        )
+
                     # Skip the ephemeral nodes since it is not materialized
                     if check_ephemeral_node(manifest_node):
                         logger.debug(f"Skipping ephemeral DBT node: {key}.")
@@ -549,6 +595,29 @@ class DbtSource(DbtServiceSource):
                         f"Failed to parse the DBT node {node} to get upstream nodes: {exc}"
                     )
                     continue
+
+        if dbt_node.resource_type == SkipResourceTypeEnum.SOURCE.value:
+            parent_fqn = fqn.build(
+                self.metadata,
+                entity_type=Table,
+                service_name="*",
+                database_name=get_corrected_name(dbt_node.database),
+                schema_name=get_corrected_name(dbt_node.schema_),
+                table_name=dbt_node.name,
+            )
+
+            # check if the parent table exists in OM before adding it to the upstream list
+            parent_table_entity: Optional[
+                Union[Table, List[Table]]
+            ] = get_entity_from_es_result(
+                entity_list=self.metadata.es_search_from_fqn(
+                    entity_type=Table, fqn_search_string=parent_fqn
+                ),
+                fetch_multiple_entities=False,
+            )
+            if parent_table_entity:
+                upstream_nodes.append(parent_fqn)
+
         return upstream_nodes
 
     def parse_data_model_columns(

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

@@ -20,12 +20,14 @@ from pydantic import BaseModel
 class DbtFiles(BaseModel):
     dbt_catalog: Optional[dict] = None
     dbt_manifest: dict
+    dbt_sources: Optional[dict] = None
     dbt_run_results: Optional[List[dict]] = None
 
 
 class DbtObjects(BaseModel):
     dbt_catalog: Optional[Any] = None
     dbt_manifest: Any
+    dbt_sources: Optional[Any] = None
     dbt_run_results: Optional[List[Any]] = None
 
 

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

@@ -1,7 +1,10 @@
 from metadata.ingestion.source.database.dynamodb.metadata import DynamodbSource
 from metadata.profiler.interface.nosql.profiler_interface import NoSQLProfilerInterface
+from metadata.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=DynamodbSource, profiler_class=NoSQLProfilerInterface
+    metadata_source_class=DynamodbSource,
+    profiler_class=NoSQLProfilerInterface,
+    sampler_class=NoSQLSampler,
 )

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

@@ -1,7 +1,10 @@
 from metadata.ingestion.source.database.mongodb.metadata import MongodbSource
 from metadata.profiler.interface.nosql.profiler_interface import NoSQLProfilerInterface
+from metadata.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
-    metadata_source_class=MongodbSource, profiler_class=NoSQLProfilerInterface
+    metadata_source_class=MongodbSource,
+    profiler_class=NoSQLProfilerInterface,
+    sampler_class=NoSQLSampler,
 )

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

@@ -38,6 +38,7 @@ from metadata.ingestion.connections.test_connections import (
     test_connection_db_schema_sources,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.mysql.queries import MYSQL_TEST_GET_QUERIES
 from metadata.utils.constants import THREE_MIN
 
 
@@ -73,10 +74,14 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
+    queries = {
+        "GetQueries": MYSQL_TEST_GET_QUERIES,
+    }
     return test_connection_db_schema_sources(
         metadata=metadata,
         engine=engine,
         service_connection=service_connection,
         automation_workflow=automation_workflow,
         timeout_seconds=timeout_seconds,
+        queries=queries,
     )

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

@@ -9,38 +9,20 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """
-Mysql lineage module
+MYSQL lineage module
 """
-from typing import Optional
-
-from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
-)
-from metadata.generated.schema.metadataIngestion.workflow import (
-    Source as WorkflowSource,
-)
-from metadata.ingestion.api.steps import InvalidSourceException
-from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.lineage_source import LineageSource
-from metadata.utils.logger import ingestion_logger
-
-logger = ingestion_logger()
+from metadata.ingestion.source.database.mysql.queries import MYSQL_SQL_STATEMENT
+from metadata.ingestion.source.database.mysql.query_parser import MysqlQueryParserSource
 
 
-class MysqlLineageSource(LineageSource):
+class MysqlLineageSource(MysqlQueryParserSource, LineageSource):
+    sql_stmt = MYSQL_SQL_STATEMENT
+    filters = """
+        AND (
+            lower(argument) LIKE '%%create%%table%%select%%'
+            OR lower(argument) LIKE '%%insert%%into%%select%%'
+            OR lower(argument) LIKE '%%update%%'
+            OR lower(argument) LIKE '%%merge%%'
+        )
     """
-    Mysql lineage source implements view lineage
-    """
-
-    @classmethod
-    def create(
-        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
-    ):
-        """Create class instance"""
-        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
-        connection: MysqlConnection = config.serviceConnection.root.config
-        if not isinstance(connection, MysqlConnection):
-            raise InvalidSourceException(
-                f"Expected MysqlConnection, but got {connection}"
-            )
-        return cls(config, metadata)

metadata/ingestion/source/database/mysql/queries.py

@@ -0,0 +1,43 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.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
+#  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.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+SQL Queries used during ingestion
+"""
+import textwrap
+
+MYSQL_SQL_STATEMENT = textwrap.dedent(
+    """
+SELECT 
+	NULL `database_name`,
+	argument `query_text`,
+	event_time `start_time`,
+    NULL `end_time`,
+	NULL `duration`,
+	NULL `schema_name`,
+    NULL `query_type`,
+    NULL `user_name`,
+    NULL `aborted`
+FROM mysql.general_log
+WHERE command_type = 'Query' 
+    AND event_time between '{start_time}' and '{end_time}'
+    AND argument NOT LIKE '/* {{"app": "OpenMetadata", %%}} */%%'
+    AND argument NOT LIKE '/* {{"app": "dbt", %%}} */%%'
+    {filters}
+ORDER BY event_time desc
+LIMIT {result_limit};
+"""
+)
+
+MYSQL_TEST_GET_QUERIES = textwrap.dedent(
+    """
+SELECT `argument` from mysql.general_log limit 1;
+"""
+)

metadata/ingestion/source/database/mysql/query_parser.py

@@ -0,0 +1,46 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.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
+#  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.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Mysql query parser module
+"""
+from abc import ABC
+from typing import Optional
+
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.query_parser_source import QueryParserSource
+
+
+class MysqlQueryParserSource(QueryParserSource, ABC):
+    """
+    Mysql base for Usage and Lineage
+    """
+
+    filters: str
+
+    @classmethod
+    def create(
+        cls, config_dict, metadata: OpenMetadata, pipeline_name: Optional[str] = None
+    ):
+        """Create class instance"""
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: MysqlConnection = config.serviceConnection.root.config
+        if not isinstance(connection, MysqlConnection):
+            raise InvalidSourceException(
+                f"Expected MysqlConnection, but got {connection}"
+            )
+        return cls(config, metadata)

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

@@ -1,4 +1,10 @@
+from metadata.ingestion.source.database.mysql.lineage import MysqlLineageSource
 from metadata.ingestion.source.database.mysql.metadata import MysqlSource
+from metadata.ingestion.source.database.mysql.usage import MysqlUsageSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=MysqlSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=MysqlSource,
+    lineage_source_class=MysqlLineageSource,
+    usage_source_class=MysqlUsageSource,
+)

metadata/{data_quality/builders/pandas_validator_builder.py → ingestion/source/database/mysql/usage.py}

@@ -8,22 +8,17 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
-
 """
-Builder defining the structure of builders for validators for Pandas sources
+MYSQL usage module
 """
-
-from typing import TYPE_CHECKING
-
-from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
-
-if TYPE_CHECKING:
-    pass
+from metadata.ingestion.source.database.mysql.queries import MYSQL_SQL_STATEMENT
+from metadata.ingestion.source.database.mysql.query_parser import MysqlQueryParserSource
+from metadata.ingestion.source.database.usage_source import UsageSource
 
 
-class PandasValidatorBuilder(IValidatorBuilder):
-    """Builder for Pandas validators"""
+class MysqlUsageSource(MysqlQueryParserSource, UsageSource):
+    sql_stmt = MYSQL_SQL_STATEMENT
+    filters = ""
 
-    def _get_source_type(self) -> str:
-        """Return the test case"""
-        return "pandas"
+    def format_query(self, query: bytes) -> str:
+        return query.decode(errors="ignore").replace("\\n", "\n")

metadata/ingestion/source/database/sample_data.py

@@ -136,7 +136,7 @@ from metadata.parsers.schema_parsers import (
     InvalidSchemaTypeException,
     schema_parser_config_registry,
 )
-from metadata.profiler.api.models import ProfilerResponse, SampleData
+from metadata.sampler.models import SampleData, SamplerResponse
 from metadata.utils import entity_link, fqn
 from metadata.utils.constants import UTF_8
 from metadata.utils.fqn import FQN_SEPARATOR
@@ -899,22 +899,14 @@ class SampleDataSource(
 
                 self.metadata.ingest_table_sample_data(
                     table_entity,
-                    ProfilerResponse(
+                    sample_data=SamplerResponse(
                         table=table_entity,
-                        profile=CreateTableProfileRequest(
-                            tableProfile=TableProfile(
-                                timestamp=Timestamp(
-                                    int(datetime.now().timestamp() * 1000)
-                                ),
-                                columnCount=1.0,
-                                rowCount=3.0,
-                            )
-                        ),
                         sample_data=SampleData(
                             data=TableData(
                                 rows=table["sampleData"]["rows"],
                                 columns=table["sampleData"]["columns"],
-                            )
+                            ),
+                            store=True,
                         ),
                     ).sample_data.data,
                 )

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

@@ -225,7 +225,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()

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

@@ -1,9 +1,13 @@
+from metadata.data_quality.interface.sqlalchemy.snowflake.test_suite_interface import (
+    SnowflakeTestSuiteInterface,
+)
 from metadata.ingestion.source.database.snowflake.lineage import SnowflakeLineageSource
 from metadata.ingestion.source.database.snowflake.metadata import SnowflakeSource
 from metadata.ingestion.source.database.snowflake.profiler.profiler import (
     SnowflakeProfiler,
 )
 from metadata.ingestion.source.database.snowflake.usage import SnowflakeUsageSource
+from metadata.sampler.sqlalchemy.snowflake.sampler import SnowflakeSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
@@ -11,4 +15,6 @@ ServiceSpec = DefaultDatabaseSpec(
     lineage_source_class=SnowflakeLineageSource,
     usage_source_class=SnowflakeUsageSource,
     profiler_class=SnowflakeProfiler,
+    test_suite_class=SnowflakeTestSuiteInterface,
+    sampler_class=SnowflakeSampler,
 )

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

@@ -4,6 +4,7 @@ from metadata.ingestion.source.database.trino.usage import TrinoUsageSource
 from metadata.profiler.interface.sqlalchemy.trino.profiler_interface import (
     TrinoProfilerInterface,
 )
+from metadata.sampler.sqlalchemy.trino.sampler import TrinoSampler
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
 ServiceSpec = DefaultDatabaseSpec(
@@ -11,4 +12,5 @@ ServiceSpec = DefaultDatabaseSpec(
     lineage_source_class=TrinoLineageSource,
     usage_source_class=TrinoUsageSource,
     profiler_class=TrinoProfilerInterface,
+    sampler_class=TrinoSampler,
 )

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

@@ -1,3 +1,6 @@
+from metadata.data_quality.interface.sqlalchemy.unity_catalog.test_suite_interface import (
+    UnityCatalogTestSuiteInterface,
+)
 from metadata.ingestion.source.database.unitycatalog.lineage import (
     UnitycatalogLineageSource,
 )
@@ -15,4 +18,5 @@ ServiceSpec = DefaultDatabaseSpec(
     lineage_source_class=UnitycatalogLineageSource,
     usage_source_class=UnitycatalogUsageSource,
     profiler_class=UnityCatalogProfilerInterface,
+    test_suite_class=UnityCatalogTestSuiteInterface,
 )

metadata/ingestion/source/pipeline/airflow/connection.py

@@ -16,7 +16,9 @@ from functools import partial, singledispatch
 from typing import Optional
 
 from airflow import settings
+from airflow.models.serialized_dag import SerializedDagModel
 from sqlalchemy.engine import Engine
+from sqlalchemy.orm import sessionmaker
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -102,6 +104,18 @@ def get_connection(connection: AirflowConnection) -> Engine:
         raise SourceConnectionException(msg) from exc
 
 
+class AirflowPipelineDetailsAccessError(Exception):
+    """
+    Raise when Pipeline information is not retrieved
+    """
+
+
+class AirflowTaskDetailsAccessError(Exception):
+    """
+    Raise when Task detail information is not retrieved
+    """
+
+
 def test_connection(
     metadata: OpenMetadata,
     engine: Engine,
@@ -114,8 +128,37 @@ def test_connection(
     of a metadata workflow or during an Automation Workflow
     """
 
-    test_fn = {"CheckAccess": partial(test_connection_engine_step, engine)}
-
+    session_maker = sessionmaker(bind=engine)
+    session = session_maker()
+
+    def test_pipeline_details_access(session):
+        try:
+            result = session.query(SerializedDagModel).first()
+            return result
+        except Exception as e:
+            raise AirflowPipelineDetailsAccessError(
+                f"Pipeline details access error: {e}"
+            )
+
+    def test_task_detail_access(session):
+        try:
+            json_data_column = (
+                SerializedDagModel._data  # For 2.3.0 onwards # pylint: disable=protected-access
+                if hasattr(SerializedDagModel, "_data")
+                else SerializedDagModel.data  # For 2.2.5 and 2.1.4
+            )
+            result = session.query(json_data_column).first()
+
+            retrieved_tasks = result[0]["dag"]["tasks"]
+            return retrieved_tasks
+        except Exception as e:
+            raise AirflowTaskDetailsAccessError(f"Task details access error : {e}")
+
+    test_fn = {
+        "CheckAccess": partial(test_connection_engine_step, engine),
+        "PipelineDetailsAccess": partial(test_pipeline_details_access, session),
+        "TaskDetailAccess": partial(test_task_detail_access, session),
+    }
     return test_connection_steps(
         metadata=metadata,
         test_fn=test_fn,

metadata/ingestion/source/pipeline/airflow/metadata.py

@@ -357,32 +357,42 @@ class AirflowSource(PipelineServiceSource):
             ).filter(
                 DagModel.is_paused == False  # pylint: disable=singleton-comparison
             )
-        for serialized_dag in session_query.yield_per(100):
-            try:
-                data = serialized_dag[1]["dag"]
-                dag = AirflowDagDetails(
-                    dag_id=serialized_dag[0],
-                    fileloc=serialized_dag[2],
-                    data=AirflowDag.model_validate(serialized_dag[1]),
-                    max_active_runs=data.get("max_active_runs", None),
-                    description=data.get("_description", None),
-                    start_date=data.get("start_date", None),
-                    tasks=list(
-                        map(self._extract_serialized_task, data.get("tasks", []))
-                    ),
-                    schedule_interval=get_schedule_interval(data),
-                    owner=self.fetch_dag_owners(data),
-                )
+        limit = 100  # Number of records per batch
+        offset = 0  # Start
+
+        while True:
+            paginated_query = session_query.limit(limit).offset(offset)
+            results = paginated_query.all()
+            if not results:
+                break
+            for serialized_dag in results:
+                try:
+                    data = serialized_dag[1]["dag"]
+                    dag = AirflowDagDetails(
+                        dag_id=serialized_dag[0],
+                        fileloc=serialized_dag[2],
+                        data=AirflowDag.model_validate(serialized_dag[1]),
+                        max_active_runs=data.get("max_active_runs", None),
+                        description=data.get("_description", None),
+                        start_date=data.get("start_date", None),
+                        tasks=list(
+                            map(self._extract_serialized_task, data.get("tasks", []))
+                        ),
+                        schedule_interval=get_schedule_interval(data),
+                        owner=self.fetch_dag_owners(data),
+                    )
 
-                yield dag
-            except ValidationError as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(
-                    f"Error building pydantic model for {serialized_dag} - {err}"
-                )
-            except Exception as err:
-                logger.debug(traceback.format_exc())
-                logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+                    yield dag
+                except ValidationError as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(
+                        f"Error building pydantic model for {serialized_dag} - {err}"
+                    )
+                except Exception as err:
+                    logger.debug(traceback.format_exc())
+                    logger.warning(f"Wild error yielding dag {serialized_dag} - {err}")
+
+            offset += limit
 
     def fetch_dag_owners(self, data) -> Optional[str]:
         """