From 43ac9e82cf2f0ba96388f9a9a0422d1d0b3829a4 Mon Sep 17 00:00:00 2001 From: Brett Naul Date: Fri, 5 Apr 2024 12:45:04 -0400 Subject: [PATCH 1/2] fix: Fix partitioning by DATE column Fixes #1056. --- sqlalchemy_bigquery/base.py | 5 +++++ tests/system/test_sqlalchemy_bigquery.py | 11 +++++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index e80f2891..636ef9c0 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -831,6 +831,11 @@ def _process_time_partitioning( if time_partitioning.field is not None: field = time_partitioning.field if isinstance( + table.columns[time_partitioning.field].type, + sqlalchemy.sql.sqltypes.DATE, + ): + return f"PARTITION BY {field}" + elif isinstance( table.columns[time_partitioning.field].type, sqlalchemy.sql.sqltypes.TIMESTAMP, ): diff --git a/tests/system/test_sqlalchemy_bigquery.py b/tests/system/test_sqlalchemy_bigquery.py index 457a8ea8..423e1ad5 100644 --- a/tests/system/test_sqlalchemy_bigquery.py +++ b/tests/system/test_sqlalchemy_bigquery.py @@ -561,7 +561,14 @@ def test_dml(engine, session, table_dml): assert len(result) == 0 -def test_create_table(engine, bigquery_dataset): +@pytest.mark.parametrize( + "time_partitioning_field", + [ + "timestamp_c", + "date_c" + ] +) +def test_create_table(engine, bigquery_dataset, time_partitioning_field): meta = MetaData() Table( f"{bigquery_dataset}.test_table_create", @@ -581,7 +588,7 @@ def test_create_table(engine, bigquery_dataset): bigquery_friendly_name="test table name", bigquery_expiration_timestamp=datetime.datetime(2183, 3, 26, 8, 30, 0), bigquery_time_partitioning=TimePartitioning( - field="timestamp_c", + field=time_partitioning_field, expiration_ms=1000 * 60 * 60 * 24 * 30, # 30 days ), bigquery_require_partition_filter=True, From 709c0d0fc3bee585654f115bc05026db00358db9 Mon Sep 17 00:00:00 2001 From: chalmer lowe Date: Tue, 7 May 2024 20:39:47 +0000 Subject: [PATCH 2/2] update with a new unit test --- tests/unit/test_table_options.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/unit/test_table_options.py b/tests/unit/test_table_options.py index 2147fb1d..13291a9a 100644 --- a/tests/unit/test_table_options.py +++ b/tests/unit/test_table_options.py @@ -193,6 +193,23 @@ def test_table_time_partitioning_with_timestamp_dialect_option(faux_conn): ) +def test_table_time_partitioning_with_date_dialect_option(faux_conn): + # expect table creation to fail as SQLite does not support partitioned tables + with pytest.raises(sqlite3.OperationalError): + setup_table( + faux_conn, + "some_table_2", + sqlalchemy.Column("id", sqlalchemy.Integer), + sqlalchemy.Column("createdAt", sqlalchemy.DATE), + bigquery_time_partitioning=TimePartitioning(field="createdAt"), + ) + # confirm that the following code creates the correct SQL string + assert " ".join(faux_conn.test_data["execute"][-1][0].strip().split()) == ( + "CREATE TABLE `some_table_2` ( `id` INT64, `createdAt` DATE )" + " PARTITION BY createdAt" + ) + + def test_table_time_partitioning_dialect_option_partition_expiration_days(faux_conn): # expect table creation to fail as SQLite does not support partitioned tables with pytest.raises(sqlite3.OperationalError):