diff --git a/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py b/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py index 1712ccd6cef8..cd82565506b6 100644 --- 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 index 4f8de1611375..3279f5007670 100644 --- 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 ),