MINOR BQ sampler type missing (#19696)

* fix: missing entity type in bq sampler

* fix: failing tests
This commit is contained in:
Teddy 2025-02-11 10:46:34 -08:00 committed by GitHub
parent 92c9f0754b
commit e1b3e08317
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 27 additions and 5 deletions

View File

@ -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

View File

@ -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
),