feat(sqlalchemy-bigquery): support BigQuery SAFE_CAST and SQLAlchemy 2.0 try_cast (#18953)

Author: chalmerlowe

packages/sqlalchemy-bigquery/README.rst

@@ -84,7 +84,7 @@ Installations when processing large datasets
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 When handling large datasets, you may see speed increases by also installing the
-`bqstorage` dependencies. See the instructions above about creating a virtual 
+`bqstorage` dependencies. See the instructions above about creating a virtual
 environment and then install `sqlalchemy-bigquery` using the `bqstorage` extras:
 
 .. code-block:: console
@@ -117,7 +117,7 @@ Project
 Authentication
 ^^^^^^^^^^^^^^
 
-Follow the `Google Cloud library guide <https://google-cloud-python.readthedocs.io/en/latest/core/auth.html>`_ for authentication. 
+Follow the `Google Cloud library guide <https://google-cloud-python.readthedocs.io/en/latest/core/auth.html>`_ for authentication.
 
 Alternatively, you can choose either of the following approaches:
 
@@ -131,10 +131,10 @@ Alternatively, you can choose either of the following approaches:
 * pass the credentials in ``create_engine()`` as a Python dictionary using the ``credentials_info`` parameter:
 
 .. code-block:: python
-    
+
     # provide credentials as a Python dictionary
     credentials_info = {
-        "type": "service_account", 
+        "type": "service_account",
         "project_id": "your-service-account-project-id"
     }
     engine = create_engine('bigquery://', credentials_info=credentials_info)
@@ -351,6 +351,22 @@ To create an integer-range partitioned table
     )
 
 
+Casting with ``SAFE_CAST``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To generate a BigQuery ``SAFE_CAST`` expression (which returns ``NULL`` instead
+of raising an error when a cast fails), use ``safe_cast`` from
+``sqlalchemy_bigquery`` (or ``sqlalchemy.try_cast`` on SQLAlchemy 2.0+):
+
+.. code-block:: python
+
+    from sqlalchemy import Integer, select
+    from sqlalchemy_bigquery import safe_cast
+
+    stmt = select(safe_cast(table.c.string_col, Integer))
+
+
+
 Threading and Multiprocessing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

packages/sqlalchemy-bigquery/docs/README.rst

@@ -84,7 +84,7 @@ Installations when processing large datasets
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 When handling large datasets, you may see speed increases by also installing the
-`bqstorage` dependencies. See the instructions above about creating a virtual 
+`bqstorage` dependencies. See the instructions above about creating a virtual
 environment and then install `sqlalchemy-bigquery` using the `bqstorage` extras:
 
 .. code-block:: console
@@ -117,7 +117,7 @@ Project
 Authentication
 ^^^^^^^^^^^^^^
 
-Follow the `Google Cloud library guide <https://google-cloud-python.readthedocs.io/en/latest/core/auth.html>`_ for authentication. 
+Follow the `Google Cloud library guide <https://google-cloud-python.readthedocs.io/en/latest/core/auth.html>`_ for authentication.
 
 Alternatively, you can choose either of the following approaches:
 
@@ -131,10 +131,10 @@ Alternatively, you can choose either of the following approaches:
 * pass the credentials in ``create_engine()`` as a Python dictionary using the ``credentials_info`` parameter:
 
 .. code-block:: python
-    
+
     # provide credentials as a Python dictionary
     credentials_info = {
-        "type": "service_account", 
+        "type": "service_account",
         "project_id": "your-service-account-project-id"
     }
     engine = create_engine('bigquery://', credentials_info=credentials_info)
@@ -351,6 +351,22 @@ To create an integer-range partitioned table
     )
 
 
+Casting with ``SAFE_CAST``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+To generate a BigQuery ``SAFE_CAST`` expression (which returns ``NULL`` instead
+of raising an error when a cast fails), use ``safe_cast`` from
+``sqlalchemy_bigquery`` (or ``sqlalchemy.try_cast`` on SQLAlchemy 2.0+):
+
+.. code-block:: python
+
+    from sqlalchemy import Integer, select
+    from sqlalchemy_bigquery import safe_cast
+
+    stmt = select(safe_cast(table.c.string_col, Integer))
+
+
+
 Threading and Multiprocessing
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/__init__.py

@@ -42,7 +42,7 @@
     TIME,
     TIMESTAMP,
 )
-from .base import BigQueryDialect, dialect
+from .base import BigQueryDialect, SafeCast, dialect, safe_cast
 from .version import __version__
 
 # Now that support for Python 3.7, 3.8 and 3.9 has been removed, we don't expect the
@@ -77,8 +77,10 @@
     "RECORD",
     "STRING",
     "STRUCT",
+    "SafeCast",
     "TIME",
     "TIMESTAMP",
+    "safe_cast",
 ]
 
 try:

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py

@@ -25,6 +25,7 @@
 import re
 import uuid
 from decimal import Decimal
+from typing import Any
 
 import google.api_core.exceptions
 import packaging.version
@@ -587,6 +588,13 @@ def visit_not_regexp_match_op_binary(self, binary, operator, **kw):
     def visit_mod_binary(self, binary, operator, **kw):
         return f"MOD({self.process(binary.left, **kw)}, {self.process(binary.right, **kw)})"
 
+    def visit_try_cast(self, cast: elements.Cast, **kwargs: Any) -> str:
+        """Compile a ``SAFE_CAST`` / ``try_cast`` expression for BigQuery."""
+        return "SAFE_CAST(%s AS %s)" % (
+            self.process(cast.clause, **kwargs),
+            self.process(cast.typeclause, **kwargs),
+        )
+
 
 class BigQueryTypeCompiler(GenericTypeCompiler):
     def visit_INTEGER(self, type_, **kw):
@@ -1382,6 +1390,32 @@ def get_view_definition(self, connection, view_name, schema=None, **kw):
         return view.view_query
 
 
+_BaseTryCast = getattr(elements, "TryCast", elements.Cast)
+
+
+class SafeCast(_BaseTryCast):
+    """Represent a BigQuery ``SAFE_CAST`` expression."""
+
+    __visit_name__ = "try_cast"
+    inherit_cache = True
+
+
+def safe_cast(expression: Any, type_: Any) -> SafeCast:
+    """Produce a BigQuery ``SAFE_CAST`` expression.
+
+    ``SAFE_CAST`` is identical to ``CAST``, except that it returns ``NULL``
+    instead of raising an error when a cast fails.
+
+    Args:
+        expression: A SQL expression, column, or literal to be cast.
+        type_: The target SQLAlchemy ``TypeEngine`` class or instance.
+
+    Returns:
+        A ``SafeCast`` clause element that compiles to ``SAFE_CAST(...)``.
+    """
+    return SafeCast(expression, type_)
+
+
 class unnest(sqlalchemy.sql.functions.GenericFunction):
     def __init__(self, *args, **kwargs):
         expr = kwargs.pop("expr", None)

packages/sqlalchemy-bigquery/tests/unit/test_compiler.py

@@ -21,6 +21,8 @@
 import sqlalchemy.exc
 from sqlalchemy.sql.functions import cube, grouping_sets, rollup
 
+from sqlalchemy_bigquery import SafeCast, safe_cast
+
 from .conftest import setup_table, sqlalchemy_2_0_or_higher, sqlalchemy_before_2_0
 
 
@@ -487,3 +489,78 @@ def test_float_literal_pyformat_bindparam():
     stmt = sqlalchemy.select(sqlalchemy.literal(15.7563))
     compiled = stmt.compile(dialect=dialect)
     assert "%(param_1:FLOAT64)s" in str(compiled)
+
+
+SAFE_CAST_TYPE_CASES = [
+    pytest.param(sqlalchemy.Integer, "INT64", id="integer_class"),
+    pytest.param(sqlalchemy.Integer(), "INT64", id="integer_instance"),
+    pytest.param(sqlalchemy.String, "STRING", id="string"),
+    pytest.param(sqlalchemy.Float, "FLOAT64", id="float"),
+    pytest.param(sqlalchemy.Boolean, "BOOL", id="boolean"),
+    pytest.param(sqlalchemy.Date, "DATE", id="date"),
+    pytest.param(sqlalchemy.Numeric, "NUMERIC", id="numeric"),
+]
+
+
+@pytest.mark.parametrize("target_type,expected_bq_type", SAFE_CAST_TYPE_CASES)
+def test_safe_cast_column(faux_conn, table, target_type, expected_bq_type):
+    expr = safe_cast(table.c.foo, target_type)
+    assert isinstance(expr, SafeCast)
+    assert expr.inherit_cache is True
+    expected_type_cls = (
+        target_type if isinstance(target_type, type) else type(target_type)
+    )
+    assert isinstance(expr.type, expected_type_cls)
+
+    stmt = sqlalchemy.select(expr.label("casted"))
+    found_sql = stmt.compile(faux_conn).string
+    expected_sql = (
+        f"SELECT SAFE_CAST(`table1`.`foo` AS {expected_bq_type}) AS `casted` \n"
+        "FROM `table1`"
+    )
+    assert found_sql == expected_sql
+
+
+def test_safe_cast_literal(faux_conn):
+    stmt = sqlalchemy.select(safe_cast(sqlalchemy.literal("123"), sqlalchemy.Integer))
+    found_sql = stmt.compile(faux_conn, compile_kwargs={"literal_binds": True}).string
+    assert found_sql == "SELECT SAFE_CAST('123' AS INT64) AS `anon_1`"
+
+
+def test_safe_cast_null(faux_conn):
+    stmt = sqlalchemy.select(
+        safe_cast(sqlalchemy.null(), sqlalchemy.Integer).label("val")
+    )
+    found_sql = stmt.compile(faux_conn).string
+    assert found_sql == "SELECT SAFE_CAST(NULL AS INT64) AS `val`"
+
+
+def test_safe_cast_nested(faux_conn, table):
+    inner = safe_cast(table.c.foo, sqlalchemy.String)
+    outer = safe_cast(inner, sqlalchemy.Integer).label("casted")
+    stmt = sqlalchemy.select(outer)
+    found_sql = stmt.compile(faux_conn).string
+    expected_sql = (
+        "SELECT SAFE_CAST(SAFE_CAST(`table1`.`foo` AS STRING) AS INT64) AS `casted` \n"
+        "FROM `table1`"
+    )
+    assert found_sql == expected_sql
+
+
+@pytest.mark.skipif(
+    not hasattr(sqlalchemy, "try_cast"),
+    reason="sqlalchemy.try_cast requires SQLAlchemy 2.0.14 or higher",
+)
+@pytest.mark.parametrize("target_type,expected_bq_type", SAFE_CAST_TYPE_CASES)
+def test_sqlalchemy_try_cast_compiles_to_safe_cast(
+    faux_conn, table, target_type, expected_bq_type
+):
+    stmt = sqlalchemy.select(
+        sqlalchemy.try_cast(table.c.foo, target_type).label("casted")
+    )
+    found_sql = stmt.compile(faux_conn).string
+    expected_sql = (
+        f"SELECT SAFE_CAST(`table1`.`foo` AS {expected_bq_type}) AS `casted` \n"
+        "FROM `table1`"
+    )
+    assert found_sql == expected_sql