diff --git a/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py b/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py index 1712ccd6cef..cd82565506b 100644 --- a/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py +++ b/ingestion/src/metadata/sampler/sqlalchemy/bigquery/sampler.py @@ -54,7 +54,6 @@ class BigQuerySampler(SQASampler): 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 @@ class BigQuerySampler(SQASampler): 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 4f8de161137..3279f500767 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 @@ class SampleTest(TestCase): """ 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 @@ class SampleTest(TestCase): """ 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 ),