Fixing Sqlalchemy UUID mapping when using BigQuery.

sqlalchemy_model_bigquery.py

# I encountered this error when using `Mapped[uuid.UUID]` in my sqlalchemy 2.0 data models in conjuntion with BigQuery:
# sqlalchemy.exc.DatabaseError: (google.cloud.bigquery.dbapi.exceptions.DatabaseError) 400 POST https://bigquery.googleapis.com/bigquery/v2... Type not found: CHAR at [3:14]

# requirements:
# pip install sqlalchemy-bigquery=1.11.0.dev2

# Here is the fix:

from operator import attrgetter
from sqlalchemy.types import TypeDecorator, String
from sqlalchemy.dialects.mssql import UNIQUEIDENTIFIER
from sqlalchemy.dialects.postgresql import UUID
import uuid


# see also https://docs.sqlalchemy.org/en/20/core/custom_types.html#backend-agnostic-guid-type
class BigQueryUUID(TypeDecorator):
    """Platform-independent GUID type.

    Uses PostgreSQL's UUID type or MSSQL's UNIQUEIDENTIFIER,
    otherwise uses CHAR(32), storing as stringified hex values.

    """

    impl = String  # using String instead of CHAR, because BigQuery does not have type CHAR
    cache_ok = True

    _default_type = String(32)  # using String instead of CHAR, because BigQuery does not have type CHAR
    _uuid_as_str = attrgetter("hex")

    def load_dialect_impl(self, dialect):
        if dialect.name == "postgresql":
            return dialect.type_descriptor(UUID())
        elif dialect.name == "mssql":
            return dialect.type_descriptor(UNIQUEIDENTIFIER())
        else:
            return dialect.type_descriptor(self._default_type)

    def process_bind_param(self, value, dialect):
        if value is None:
            return value
        elif dialect.name in ("postgresql", "mssql"):
            return str(value)
        else:
            if not isinstance(value, uuid.UUID):
                value = uuid.UUID(value)
            return self._uuid_as_str(value)

    def process_result_value(self, value, dialect):
        if value is None:
            return value
        else:
            if not isinstance(value, uuid.UUID):
                value = uuid.UUID(value)
            return value


# Data model definition:

from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    # use this type mapping
    type_annotation_map = {
        uuid.UUID: BigQueryUUID,
    }


class ExampleTable(Base):
    id: Mapped[uuid.UUID] = mapped_column(primary_key=True, default=uuid.uuid4)
    content: Mapped[str | None]