openmetadata-ingestion 1.7.0.0rc2__py3-none-any.whl → 1.7.0.1__py3-none-any.whl

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

@@ -56,6 +56,7 @@ from metadata.utils import fqn
 from metadata.utils.sqlalchemy_utils import (
     get_display_datatype,
     get_table_comment_wrapper,
+    get_view_definition_wrapper,
 )
 
 dialect = SnowflakeDialect()
@@ -271,25 +272,16 @@ def get_stream_names(self, connection, schema, **kw):
 
 
 @reflection.cache
-def get_view_definition(  # pylint: disable=unused-argument
-    self, connection, view_name, schema=None, **kw
-):
-    """
-    Gets the view definition
-    """
-    schema = schema or self.default_schema_name
-    view_name = f"{schema}.{view_name}" if schema else view_name
-    cursor = connection.execute(
-        SNOWFLAKE_GET_VIEW_DEFINITION.format(view_name=view_name)
+def get_view_definition(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    return get_view_definition_wrapper(
+        self,
+        connection,
+        table_name=table_name,
+        schema=schema,
+        query=SNOWFLAKE_GET_VIEW_DEFINITION,
     )
-    n2i = self.__class__._map_name_to_idx(cursor)  # pylint: disable=protected-access
-    try:
-        ret = cursor.fetchone()
-        if ret:
-            return ret[n2i["text"]]
-    except Exception:
-        pass
-    return None
 
 
 @reflection.cache

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

@@ -72,7 +72,6 @@ from metadata.ingestion.source.database.unitycatalog.models import (
     ForeignConstrains,
     Type,
 )
-from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
 from metadata.utils.helpers import retry_with_docker_host
@@ -97,7 +96,6 @@ class UnitycatalogSource(
         self.source_config: DatabaseServiceMetadataPipeline = (
             self.config.sourceConfig.config
         )
-        self.context.get_global().table_views = []
         self.metadata = metadata
         self.service_connection: UnityCatalogConnection = (
             self.config.serviceConnection.root.config
@@ -346,19 +344,6 @@ class UnitycatalogSource(
             )
             yield Either(right=table_request)
 
-            if table_type == TableType.View or table.view_definition:
-                self.context.get_global().table_views.append(
-                    TableView(
-                        table_name=table_name,
-                        schema_name=schema_name,
-                        db_name=db_name,
-                        view_definition=(
-                            f'CREATE VIEW "{db_name}"."{schema_name}"'
-                            f'."{table_name}" AS {table.view_definition}'
-                        ),
-                    )
-                )
-
             self.register_record(table_request=table_request)
         except Exception as exc:
             yield Either(

metadata/ingestion/source/messaging/kinesis/metadata.py

@@ -185,6 +185,9 @@ class KinesisSource(MessagingServiceSource):
         try:
             while has_more_partitions:
                 partitions = self.kinesis.list_shards(**args.dict())
+                # Handle the case when NextToken is not present
+                if "NextToken" not in partitions:
+                    partitions["NextToken"] = None
                 kinesis_partitions_model = KinesisPartitions(**partitions)
                 all_partitions.extend(
                     [

metadata/profiler/source/database/bigquery/profiler_source.py

@@ -13,8 +13,6 @@
 Bigquery Profiler source
 """
 
-from copy import deepcopy
-
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
@@ -22,12 +20,8 @@ from metadata.generated.schema.entity.services.connections.database.bigQueryConn
 from metadata.generated.schema.metadataIngestion.workflow import (
     OpenMetadataWorkflowConfig,
 )
-from metadata.generated.schema.security.credentials.gcpValues import (
-    GcpCredentialsValues,
-    MultipleProjectId,
-    SingleProjectId,
-)
 from metadata.profiler.source.database.base.profiler_source import ProfilerSource
+from metadata.utils.bigquery_utils import copy_service_config
 
 
 class BigQueryProfilerSource(ProfilerSource):
@@ -46,16 +40,4 @@ class BigQueryProfilerSource(ProfilerSource):
         Returns:
             DatabaseConnection
         """
-        config_copy: BigQueryConnection = deepcopy(
-            config.source.serviceConnection.root.config  # type: ignore
-        )
-
-        if isinstance(config_copy.credentials.gcpConfig, GcpCredentialsValues):
-            if isinstance(
-                config_copy.credentials.gcpConfig.projectId, MultipleProjectId
-            ):
-                config_copy.credentials.gcpConfig.projectId = SingleProjectId(
-                    database.name.root
-                )
-
-        return config_copy
+        return copy_service_config(config, database.name.root)

metadata/sampler/processor.py

@@ -17,6 +17,9 @@ from typing import Optional, cast
 
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
+    BigQueryConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -38,6 +41,7 @@ from metadata.profiler.source.metadata import ProfilerSourceAndEntity
 from metadata.sampler.config import get_config_for_table
 from metadata.sampler.models import SampleConfig, SampleData, SamplerResponse
 from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.bigquery_utils import copy_service_config
 from metadata.utils.profiler_utils import get_context_entities
 from metadata.utils.service_spec.service_spec import import_sampler_class
 
@@ -133,6 +137,9 @@ class SamplerProcessor(Processor):
         Returns:
             DatabaseService.__config__
         """
+        if isinstance(config.source.serviceConnection.root.config, BigQueryConnection):
+            return copy_service_config(config, database.name.root)
+
         config_copy = deepcopy(
             config.source.serviceConnection.root.config  # type: ignore
         )

metadata/sampler/sqlalchemy/bigquery/sampler.py

@@ -12,16 +12,18 @@
 Helper module to handle data sampling
 for the profiler
 """
-from typing import Dict, Optional, Union
+import traceback
+from typing import Dict, List, Optional, Union
 
 from sqlalchemy import Column
 from sqlalchemy import Table as SqaTable
-from sqlalchemy import text
+from sqlalchemy import inspect, text
 from sqlalchemy.orm import Query
 
 from metadata.generated.schema.entity.data.table import (
     ProfileSampleType,
     Table,
+    TableData,
     TableType,
 )
 from metadata.generated.schema.entity.services.connections.connectionBasicType import (
@@ -35,6 +37,11 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.sampler.models import SampleConfig
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
+from metadata.utils.logger import profiler_interface_registry_logger
+
+logger = profiler_interface_registry_logger()
+
+RANDOM_LABEL = "random"
 
 
 class BigQuerySampler(SQASampler):
@@ -123,3 +130,62 @@ class BigQuerySampler(SQASampler):
             )
 
         return super().get_sample_query(column=column)
+
+    def fetch_sample_data(self, columns: Optional[List[Column]] = None) -> TableData:
+        """
+        Use the sampler to retrieve sample data rows as per limit given by user
+
+        Args:
+            columns (Optional[List]): List of columns to fetch
+        Returns:
+            TableData to be added to the Table Entity
+        """
+        if self.sample_query:
+            return self._fetch_sample_data_from_user_query()
+
+        # Add new RandomNumFn column
+        ds = self.get_dataset()
+        if not columns:
+            sqa_columns = [col for col in inspect(ds).c if col.name != RANDOM_LABEL]
+        else:
+            # we can't directly use columns as it is bound to self.raw_dataset and not the rnd table.
+            # If we use it, it will result in a cross join between self.raw_dataset and rnd table
+            names = [col.name for col in columns]
+            sqa_columns = [
+                col
+                for col in inspect(ds).c
+                if col.name != RANDOM_LABEL and col.name in names
+            ]
+        with self.get_client() as client:
+            try:
+                query = (
+                    str(client.query(*sqa_columns)).replace(
+                        f"`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
+                        f"`{self.entity.database.name}`.`{ds.__table_args__['schema']}`.`{ds.__tablename__}`",
+                    )
+                    + f" LIMIT {self.sample_limit}"
+                )
+                sqa_sample = client.execute(text(query))
+            except Exception:
+                logger.debug(
+                    "Cannot fetch sample data with random sampling. Falling back to 100 rows."
+                )
+                logger.debug(traceback.format_exc())
+                ds_inspect = inspect(self.raw_dataset)
+                sqa_columns = list(ds_inspect.c)
+
+                schema = ds_inspect.class_.__table_args__["schema"]
+                table = ds_inspect.class_.__tablename__
+
+                query = (
+                    str(client.query(*sqa_columns)).replace(
+                        f"`{schema}`.`{table}`",
+                        f"`{self.entity.database.name}`.`{schema}`.`{table}`",
+                    )
+                    + f" LIMIT {self.sample_limit}"
+                )
+                sqa_sample = client.execute(text(query))
+            return TableData(
+                columns=[column.name for column in sqa_columns],
+                rows=[list(row) for row in sqa_sample],
+            )

metadata/utils/bigquery_utils.py

@@ -13,10 +13,22 @@
 Utils module of BigQuery
 """
 
+from copy import deepcopy
 from typing import List, Optional
 
 from google.cloud import bigquery
 
+from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
+    BigQueryConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+    MultipleProjectId,
+    SingleProjectId,
+)
 from metadata.utils.credentials import (
     get_gcp_default_credentials,
     get_gcp_impersonate_credentials,
@@ -56,3 +68,26 @@ def get_bigquery_client(
     return bigquery.Client(
         credentials=credentials, project=project_id, location=location
     )
+
+
+def copy_service_config(
+    config: OpenMetadataWorkflowConfig, database_name: str
+) -> BigQueryConnection:
+    """Handles multiple project id in the service config and replace it with the database name
+
+    Args:
+        config (OpenMetadataWorkflowConfig): openmetadata workflow config
+        database_name (str): database name
+
+    Returns:
+        BigQueryConnection: bigquery connection
+    """
+    config_copy: BigQueryConnection = deepcopy(
+        config.source.serviceConnection.root.config  # type: ignore
+    )
+
+    if isinstance(config_copy.credentials.gcpConfig, GcpCredentialsValues):
+        if isinstance(config_copy.credentials.gcpConfig.projectId, MultipleProjectId):
+            config_copy.credentials.gcpConfig.projectId = SingleProjectId(database_name)
+
+    return config_copy

metadata/utils/sqlalchemy_utils.py

@@ -71,10 +71,13 @@ def get_all_view_definitions(self, connection, query):
     Method to fetch view definition of all available views
     """
     self.all_view_definitions: Dict[Tuple[str, str], str] = {}
-    self.current_db: str = connection.engine.url.database
+    self.current_db: str = connection.engine.url.database  # type: ignore
     result = connection.execute(query)
     for view in result:
-        self.all_view_definitions[(view.view_name, view.schema)] = view.view_def
+        if hasattr(view, "view_def") and hasattr(view, "schema"):
+            self.all_view_definitions[(view.view_name, view.schema)] = view.view_def
+        elif hasattr(view, "VIEW_DEF") and hasattr(view, "SCHEMA"):
+            self.all_view_definitions[(view.VIEW_NAME, view.SCHEMA)] = view.VIEW_DEF
 
 
 def get_view_definition_wrapper(self, connection, query, table_name, schema=None):