fix: Fix partitioning by DATE column (googleapis#1074)

bnaul · chalmerlowe · web-flow · commit ad69c630833b · 2024-05-13T16:04:43.000-04:00
Fixes googleapis#1056. Co-authored-by: Chalmer Lowe <chalmerlowe@google.com>
diff --git 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
@@ -561,7 +561,8 @@ 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 +582,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
@@ -193,6 +193,24 @@ 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):
