Skip to content

Commit f2e1ba8

Browse files
committed
add integration test for BigQueryEngineAdapter.get_bq_schema
1 parent 90fa431 commit f2e1ba8

File tree

1 file changed

+33
-0
lines changed

1 file changed

+33
-0
lines changed

tests/core/engine_adapter/integration/test_integration_bigquery.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,39 @@ def test_compare_nested_values_in_table_diff(ctx: TestContext):
341341
ctx.engine_adapter.drop_table(target_table)
342342

343343

344+
def test_get_bq_schema(ctx: TestContext, engine_adapter: BigQueryEngineAdapter):
345+
from google.cloud.bigquery import SchemaField
346+
347+
table = ctx.table("test")
348+
349+
engine_adapter.execute(f"""
350+
CREATE TABLE {table.sql(dialect=ctx.dialect)} (
351+
id STRING NOT NULL,
352+
user_data STRUCT<id STRING NOT NULL, name STRING NOT NULL, address STRING>,
353+
tags ARRAY<STRING>,
354+
score NUMERIC,
355+
created_at DATETIME
356+
)
357+
""")
358+
359+
bg_schema = engine_adapter.get_bq_schema(table)
360+
assert len(bg_schema) == 5
361+
assert bg_schema[0] == SchemaField(name="id", field_type="STRING", mode="REQUIRED")
362+
assert bg_schema[1] == SchemaField(
363+
name="user_data",
364+
field_type="RECORD",
365+
mode="NULLABLE",
366+
fields=[
367+
SchemaField(name="id", field_type="STRING", mode="REQUIRED"),
368+
SchemaField(name="name", field_type="STRING", mode="REQUIRED"),
369+
SchemaField(name="address", field_type="STRING", mode="NULLABLE"),
370+
],
371+
)
372+
assert bg_schema[2] == SchemaField(name="tags", field_type="STRING", mode="REPEATED")
373+
assert bg_schema[3] == SchemaField(name="score", field_type="NUMERIC", mode="NULLABLE")
374+
assert bg_schema[4] == SchemaField(name="created_at", field_type="DATETIME", mode="NULLABLE")
375+
376+
344377
def test_column_types(ctx: TestContext):
345378
model_name = ctx.table("test")
346379
sqlmesh = ctx.create_context()

0 commit comments

Comments
 (0)