From 5f772c8d313ce8ac10eceb7a247d5f3f6ac46168 Mon Sep 17 00:00:00 2001 From: Brett Naul Date: Thu, 19 Sep 2024 16:02:56 -0400 Subject: [PATCH 1/2] fix: Support partitioning by DATE or Date --- sqlalchemy_bigquery/base.py | 2 +- tests/unit/test_table_options.py | 5 +++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index c531c102..81af9eaf 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -838,7 +838,7 @@ def _process_time_partitioning( field = time_partitioning.field if isinstance( table.columns[time_partitioning.field].type, - sqlalchemy.sql.sqltypes.DATE, + (sqlalchemy.sql.sqltypes.DATE, sqlalchemy.sql.sqltypes.Date), ): return f"PARTITION BY {field}" elif isinstance( diff --git a/tests/unit/test_table_options.py b/tests/unit/test_table_options.py index 2b757e04..8952f5d4 100644 --- a/tests/unit/test_table_options.py +++ b/tests/unit/test_table_options.py @@ -193,14 +193,15 @@ def test_table_time_partitioning_with_timestamp_dialect_option(faux_conn): ) -def test_table_time_partitioning_with_date_dialect_option(faux_conn): +@pytest.mark.parametrize("date_type", [sqlalchemy.DATE, sqlalchemy.Date]) +def test_table_time_partitioning_with_date_dialect_option(faux_conn, date_type): # 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), + sqlalchemy.Column("createdAt", date_type), bigquery_time_partitioning=TimePartitioning(field="createdAt"), ) From 5f2b84e97e7e0de966697f25fa2ce99ebe2f48e1 Mon Sep 17 00:00:00 2001 From: Brett Naul Date: Fri, 20 Sep 2024 10:45:57 -0400 Subject: [PATCH 2/2] Workaround for googleapis/python-bigquery-sqlalchemy#1117 --- sqlalchemy_bigquery/base.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index 81af9eaf..a5f9a2b2 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -344,8 +344,9 @@ def visit_label(self, *args, within_group_by=False, **kwargs): column_label = args[0] sql_keywords = {"GROUPING SETS", "ROLLUP", "CUBE"} for keyword in sql_keywords: - if keyword in str(column_label): - break + # if keyword in str(column_label): + # break + pass else: # for/else always happens unless break gets called kwargs["render_label_as_label"] = column_label