Merge pull request #696 from Mause/try-cast

feat: try_cast operator

Author: Mause

duckdb_engine/__init__.py

@@ -14,14 +14,16 @@
 )
 
 import duckdb
+import sqlalchemy
 from sqlalchemy import pool, text
 from sqlalchemy import types as sqltypes
 from sqlalchemy import util
 from sqlalchemy.dialects.postgresql import UUID
-from sqlalchemy.dialects.postgresql.base import PGDialect, PGInspector
+from sqlalchemy.dialects.postgresql.base import PGDialect, PGInspector, PGTypeCompiler
 from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
 from sqlalchemy.engine.default import DefaultDialect
 from sqlalchemy.engine.url import URL
+from sqlalchemy.ext.compiler import compiles
 
 from .config import apply_config, get_core_config
 from .datatypes import ISCHEMA_NAMES, register_extension_types
@@ -357,3 +359,18 @@ def get_multi_columns(
         columns = self._get_columns_info(rows, domains, enums, schema)  # type: ignore[attr-defined]
 
         return columns.items()
+
+
+if sqlalchemy.__version__ >= "2.0.14":
+    from sqlalchemy import TryCast  # type: ignore[attr-defined]
+
+    @compiles(TryCast, "duckdb")  # type: ignore[misc]
+    def visit_try_cast(
+        instance: TryCast,
+        compiler: PGTypeCompiler,
+        **kw: Any,
+    ) -> str:
+        return "TRY_CAST({} AS {})".format(
+            compiler.process(instance.clause, **kw),
+            compiler.process(instance.typeclause, **kw),
+        )

duckdb_engine/tests/test_basic.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import zlib
-from datetime import timedelta
+from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Any, Generic, Optional, TypeVar, cast
 
@@ -13,6 +13,7 @@
 from pytest import LogCaptureFixture, fixture, importorskip, mark, raises
 from sqlalchemy import (
     Column,
+    DateTime,
     ForeignKey,
     Integer,
     Interval,
@@ -22,6 +23,7 @@
     Table,
     create_engine,
     inspect,
+    select,
     text,
     types,
 )
@@ -408,3 +410,14 @@ def test_do_ping(tmp_path: Path, caplog: LogCaptureFixture) -> None:
         assert any(
             "Pool pre-ping on connection" in message for message in caplog.messages
         )
+
+
+def test_try_cast(engine: Engine) -> None:
+    try_cast = importorskip("sqlalchemy", "2.0.14").try_cast
+
+    with engine.connect() as conn:
+        query = select(try_cast("2022-01-01", DateTime))
+        assert conn.execute(query).one() == (datetime(2022, 1, 1),)
+
+        query = select(try_cast("not a date", DateTime))
+        assert conn.execute(query).one() == (None,)