diff --git a/packages/sqlalchemy-bigquery/README.rst b/packages/sqlalchemy-bigquery/README.rst index 0c31785fa9fe..84f456d4b5be 100644 --- a/packages/sqlalchemy-bigquery/README.rst +++ b/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 `_ for authentication. +Follow the `Google Cloud library guide `_ 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 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/packages/sqlalchemy-bigquery/docs/README.rst b/packages/sqlalchemy-bigquery/docs/README.rst index 0c31785fa9fe..84f456d4b5be 100644 --- a/packages/sqlalchemy-bigquery/docs/README.rst +++ b/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 `_ for authentication. +Follow the `Google Cloud library guide `_ 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 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/__init__.py b/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/__init__.py index 1d3bdb51e082..97c9476699d8 100644 --- a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/__init__.py +++ b/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: diff --git a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py b/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py index db11393d97ae..cc0bfe06a270 100644 --- a/packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py +++ b/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) diff --git a/packages/sqlalchemy-bigquery/tests/unit/test_compiler.py b/packages/sqlalchemy-bigquery/tests/unit/test_compiler.py index 4fae78e7b1e6..3616945ff1db 100644 --- a/packages/sqlalchemy-bigquery/tests/unit/test_compiler.py +++ b/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