openmetadata-ingestion 1.5.4.0__py3-none-any.whl → 1.5.4.1__py3-none-any.whl

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+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-09-13T12:48:15+00:00
+#   timestamp: 2024-09-19T13:48:25+00:00
 
 from __future__ import annotations
 

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -331,19 +331,9 @@ class ESMixin(Generic[T]):
                 else:
                     break
 
-            # Get the data
-            for hit in response.hits.hits:
-                try:
-                    yield self.get_by_name(
-                        entity=entity,
-                        fqn=hit.source["fullyQualifiedName"],
-                        fields=fields,
-                        nullable=False,  # Raise an error if we don't find the Entity
-                    )
-                except Exception as exc:
-                    logger.warning(
-                        f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
-                    )
+            yield from self._yield_hits_from_api(
+                response=response, entity=entity, fields=fields
+            )
 
             # Get next page
             last_hit = response.hits.hits[-1] if response.hits.hits else None
@@ -362,3 +352,20 @@ class ESMixin(Generic[T]):
             logger.debug(traceback.format_exc())
             logger.warning(f"Error while getting ES response: {exc}")
         return None
+
+    def _yield_hits_from_api(
+        self, response: ESResponse, entity: Type[T], fields: Optional[List[str]]
+    ) -> Iterator[T]:
+        """Get the data from the API based on ES responses"""
+        for hit in response.hits.hits:
+            try:
+                yield self.get_by_name(
+                    entity=entity,
+                    fqn=hit.source["fullyQualifiedName"],
+                    fields=fields,
+                    nullable=False,  # Raise an error if we don't find the Entity
+                )
+            except Exception as exc:
+                logger.warning(
+                    f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
+                )

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

@@ -17,6 +17,7 @@ from datetime import datetime
 from functools import partial
 from typing import Optional
 
+from google.api_core.exceptions import NotFound
 from google.cloud.datacatalog_v1 import PolicyTagManagerClient
 from sqlalchemy.engine import Engine
 
@@ -149,8 +150,8 @@ def test_connection(
         test_fn = {
             "CheckAccess": partial(test_connection_engine_step, engine),
             "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
-            "GetTables": partial(execute_inspector_func, engine, "get_table_names"),
-            "GetViews": partial(execute_inspector_func, engine, "get_view_names"),
+            "GetTables": partial(get_table_view_names, engine),
+            "GetViews": partial(get_table_view_names, engine),
             "GetTags": test_tags,
             "GetQueries": partial(
                 test_query,
@@ -170,3 +171,28 @@ def test_connection(
         )
 
     test_connection_inner(engine)
+
+
+def get_table_view_names(connection, schema=None):
+    with connection.connect() as conn:
+        current_schema = schema
+        client = conn.connection._client
+        item_types = ["TABLE", "EXTERNAL", "VIEW", "MATERIALIZED_VIEW"]
+        datasets = client.list_datasets()
+        result = []
+        for dataset in datasets:
+            if current_schema is not None and current_schema != dataset.dataset_id:
+                continue
+
+            try:
+                tables = client.list_tables(dataset.reference, page_size=1)
+                for table in tables:
+                    if table.table_type in item_types:
+                        break
+            except NotFound:
+                # It's possible that the dataset was deleted between when we
+                # fetched the list of datasets and when we try to list the
+                # tables from it. See:
+                # https://github.com/googleapis/python-bigquery-sqlalchemy/issues/105
+                pass
+        return result

metadata/ingestion/source/database/common_db_source.py

@@ -11,8 +11,11 @@
 """
 Generic source to build SQL connectors.
 """
+import math
+import time
 import traceback
 from abc import ABC
+from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from typing import Any, Iterable, List, Optional, Tuple, Union, cast
 
@@ -60,6 +63,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
+from metadata.ingestion.models.topology import Queue
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import (
     get_connection,
@@ -73,6 +77,7 @@ from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.db_utils import get_view_lineage
 from metadata.utils.execution_time_tracker import (
+    ExecutionTimeTrackerContextMap,
     calculate_execution_time,
     calculate_execution_time_generator,
 )
@@ -583,9 +588,78 @@ class CommonDbSourceService(
                 )
             )
 
-    @calculate_execution_time_generator()
-    def yield_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
-        logger.info("Processing Lineage for Views")
+    def multithread_process_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
+        """Multithread Processing of a Node"""
+
+        views_list = list(self.context.get().table_views or [])
+        views_length = len(views_list)
+
+        if views_length != 0:
+            chunksize = int(math.ceil(views_length / self.source_config.threads))
+            chunks = [
+                views_list[i : i + chunksize] for i in range(0, views_length, chunksize)
+            ]
+
+            thread_pool = ThreadPoolExecutor(max_workers=self.source_config.threads)
+            queue = Queue()
+
+            futures = [
+                thread_pool.submit(
+                    self._process_view_def_chunk,
+                    chunk,
+                    queue,
+                    self.context.get_current_thread_id(),
+                )
+                for chunk in chunks
+            ]
+
+            while True:
+                if queue.has_tasks():
+                    yield from queue.process()
+
+                else:
+                    if not futures:
+                        break
+
+                    for i, future in enumerate(futures):
+                        if future.done():
+                            future.result()
+                            futures.pop(i)
+
+                time.sleep(0.01)
+
+    def _process_view_def_chunk(
+        self, chunk: List[TableView], queue: Queue, thread_id: int
+    ) -> None:
+        """
+        Process a chunk of view definitions
+        """
+        self.context.copy_from(thread_id)
+        ExecutionTimeTrackerContextMap().copy_from_parent(thread_id)
+        for view in [v for v in chunk if v.view_definition is not None]:
+            for lineage in get_view_lineage(
+                view=view,
+                metadata=self.metadata,
+                service_name=self.context.get().database_service,
+                connection_type=self.service_connection.type.value,
+                timeout_seconds=self.source_config.queryParsingTimeoutLimit,
+            ):
+                if lineage.right is not None:
+                    queue.put(
+                        Either(
+                            right=OMetaLineageRequest(
+                                lineage_request=lineage.right,
+                                override_lineage=self.source_config.overrideViewLineage,
+                            )
+                        )
+                    )
+                else:
+                    queue.put(lineage)
+
+    def _process_view_def_serial(self) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Process view definitions serially
+        """
         for view in [
             v for v in self.context.get().table_views if v.view_definition is not None
         ]:
@@ -606,6 +680,14 @@ class CommonDbSourceService(
                 else:
                     yield lineage
 
+    @calculate_execution_time_generator()
+    def yield_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
+        logger.info("Processing Lineage for Views")
+        if self.source_config.threads > 1:
+            yield from self.multithread_process_view_lineage()
+        else:
+            yield from self._process_view_def_serial()
+
     def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]:
         """
         Search the referred table for foreign constraints

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

@@ -16,6 +16,7 @@ from functools import partial
 from typing import Optional
 
 from sqlalchemy.engine import Engine
+from sqlalchemy.sql import text
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -38,6 +39,7 @@ from metadata.ingestion.connections.test_connections import (
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import kill_active_connections
 from metadata.ingestion.source.database.redshift.queries import (
+    REDSHIFT_GET_ALL_RELATIONS,
     REDSHIFT_GET_DATABASE_NAMES,
     REDSHIFT_TEST_GET_QUERIES,
     REDSHIFT_TEST_PARTITION_DETAILS,
@@ -65,6 +67,11 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
+    table_and_view_query = text(
+        REDSHIFT_GET_ALL_RELATIONS.format(
+            schema_clause="", table_clause="", limit_clause="LIMIT 1"
+        )
+    )
 
     def test_get_queries_permissions(engine_: Engine):
         """Check if we have the right permissions to list queries"""
@@ -78,8 +85,8 @@ def test_connection(
     test_fn = {
         "CheckAccess": partial(test_connection_engine_step, engine),
         "GetSchemas": partial(execute_inspector_func, engine, "get_schema_names"),
-        "GetTables": partial(execute_inspector_func, engine, "get_table_names"),
-        "GetViews": partial(execute_inspector_func, engine, "get_view_names"),
+        "GetTables": partial(test_query, statement=table_and_view_query, engine=engine),
+        "GetViews": partial(test_query, statement=table_and_view_query, engine=engine),
         "GetQueries": partial(test_get_queries_permissions, engine),
         "GetDatabases": partial(
             test_query, statement=REDSHIFT_GET_DATABASE_NAMES, engine=engine

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

@@ -236,7 +236,7 @@ REDSHIFT_TEST_PARTITION_DETAILS = "select * from SVV_TABLE_INFO limit 1"
 # hence we are appending "create view <schema>.<table> as " to select query
 # to generate the column level lineage
 REDSHIFT_GET_ALL_RELATIONS = """
-    SELECT
+    (SELECT
         c.relkind,
         n.oid as "schema_oid",
         n.nspname as "schema",
@@ -255,8 +255,9 @@ REDSHIFT_GET_ALL_RELATIONS = """
             JOIN pg_catalog.pg_user u ON u.usesysid = c.relowner
     WHERE c.relkind IN ('r', 'v', 'm', 'S', 'f')
         AND n.nspname !~ '^pg_' {schema_clause} {table_clause}
+        {limit_clause})
     UNION
-    SELECT
+    (SELECT
         'r' AS "relkind",
         s.esoid AS "schema_oid",
         s.schemaname AS "schema",
@@ -272,7 +273,8 @@ REDSHIFT_GET_ALL_RELATIONS = """
         JOIN svv_external_schemas s ON s.schemaname = t.schemaname
         JOIN pg_catalog.pg_user u ON u.usesysid = s.esowner
     where 1 {schema_clause} {table_clause}
-    ORDER BY "relkind", "schema_oid", "schema";
+    ORDER BY "relkind", "schema_oid", "schema"
+    {limit_clause});
     """
 
 

metadata/ingestion/source/database/redshift/utils.py

@@ -386,7 +386,7 @@ def _get_all_relation_info(self, connection, **kw):  # pylint: disable=unused-ar
     result = connection.execute(
         sa.text(
             REDSHIFT_GET_ALL_RELATIONS.format(
-                schema_clause=schema_clause, table_clause=table_clause
+                schema_clause=schema_clause, table_clause=table_clause, limit_clause=""
             )
         )
     )

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

@@ -45,6 +45,7 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_TEST_FETCH_TAG,
     SNOWFLAKE_TEST_GET_QUERIES,
     SNOWFLAKE_TEST_GET_TABLES,
+    SNOWFLAKE_TEST_GET_VIEWS,
 )
 from metadata.utils.logger import ingestion_logger
 
@@ -176,7 +177,11 @@ def test_connection(
             statement=SNOWFLAKE_TEST_GET_TABLES,
             engine_wrapper=engine_wrapper,
         ),
-        "GetViews": partial(execute_inspector_func, engine_wrapper, "get_view_names"),
+        "GetViews": partial(
+            test_table_query,
+            statement=SNOWFLAKE_TEST_GET_VIEWS,
+            engine_wrapper=engine_wrapper,
+        ),
         "GetQueries": partial(
             test_query, statement=SNOWFLAKE_TEST_GET_QUERIES, engine=engine
         ),

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

@@ -65,9 +65,6 @@ from metadata.ingestion.source.database.external_table_lineage_mixin import (
 from metadata.ingestion.source.database.incremental_metadata_extraction import (
     IncrementalConfig,
 )
-from metadata.ingestion.source.database.life_cycle_query_mixin import (
-    LifeCycleQueryMixin,
-)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.snowflake.models import (
     STORED_PROC_LANGUAGE_MAP,
@@ -148,7 +145,6 @@ SnowflakeDialect._get_schema_foreign_keys = get_schema_foreign_keys
 
 
 class SnowflakeSource(
-    LifeCycleQueryMixin,
     StoredProcedureMixin,
     ExternalTableLineageMixin,
     CommonDbSourceService,

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

@@ -245,6 +245,10 @@ SNOWFLAKE_TEST_GET_TABLES = """
 SELECT TABLE_NAME FROM "{database_name}".information_schema.tables LIMIT 1
 """
 
+SNOWFLAKE_TEST_GET_VIEWS = """
+SELECT TABLE_NAME FROM "{database_name}".information_schema.views LIMIT 1
+"""
+
 SNOWFLAKE_GET_DATABASES = "SHOW DATABASES"