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