From 9da6421082ad929375a37124d435df3325165372 Mon Sep 17 00:00:00 2001
From: chalmer lowe <chalmerlowe@google.com>
Date: Mon, 13 May 2024 15:44:52 +0000
Subject: [PATCH] recreates edits from bnaul PR to add Date partition
 capability

Co-authored by: Brett Naul brettnaul@gmail.com
---
 sqlalchemy_bigquery/base.py              |  5 +++++
 tests/system/test_sqlalchemy_bigquery.py |  5 +++--
 tests/unit/test_table_options.py         | 18 ++++++++++++++++++
 3 files changed, 26 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..7ea4ccc6 100644
--- 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
index 2147fb1d..2b757e04 100644
--- 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):