Merge branch 'main' into add-ingestion-agent-reference-to-ingestion-p…

…ipeline
open-metadata · Feb 11, 2025 · 6903bce · 6903bce
2 parents dc76fd7 + e1b3e08
commit 6903bce
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py b/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py
@@ -54,7 +54,6 @@ def __init__(
         sample_query: Optional[str] = None,
         storage_config: DataStorageConfig = None,
         sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
-        table_type: TableType = None,
         **kwargs,
     ):
         super().__init__(
@@ -68,7 +67,7 @@ def __init__(
             sample_data_count=sample_data_count,
             **kwargs,
         )
-        self.raw_dataset_type: TableType = table_type
+        self.raw_dataset_type: Optional[TableType] = entity.tableType
 
     def set_tablesample(self, selectable: SqaTable):
         """Set the TABLESAMPLE clause for BigQuery

diff --git a/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_bigquery_sampling.py b/ingestion/tests/unit/profiler/sqlalchemy/bigquery/test_bigquery_sampling.py
@@ -127,14 +127,25 @@ def test_sampling_for_views(self, sampler_mock):
         """
         Test view sampling
         """
+        view_entity = Table(
+            id=uuid4(),
+            name="user",
+            columns=[
+                EntityColumn(
+                    name=ColumnName("id"),
+                    dataType=DataType.INT,
+                ),
+            ],
+            tableType=TableType.View,
+        )
+
         sampler = BigQuerySampler(
             service_connection_config=self.bq_conn,
             ometa_client=None,
-            entity=self.table_entity,
+            entity=view_entity,
             sample_config=SampleConfig(
                 profileSampleType=ProfileSampleType.PERCENTAGE, profileSample=50.0
             ),
-            table_type=TableType.View,
         )
         query: CTE = sampler.get_sample_query()
         expected_query = (
@@ -151,10 +162,22 @@ def test_sampling_view_with_partition(self, sampler_mock):
         """
         Test view sampling with partition
         """
+        view_entity = Table(
+            id=uuid4(),
+            name="user",
+            columns=[
+                EntityColumn(
+                    name=ColumnName("id"),
+                    dataType=DataType.INT,
+                ),
+            ],
+            tableType=TableType.View,
+        )
+
         sampler = BigQuerySampler(
             service_connection_config=self.bq_conn,
             ometa_client=None,
-            entity=self.table_entity,
+            entity=view_entity,
             sample_config=SampleConfig(
                 profileSampleType=ProfileSampleType.PERCENTAGE, profileSample=50.0
             ),