Casting our Own Types for SQLAlchemy
SQLAlchemy's type system is incredibly flexible, allowing us to map Python types to SQL database types in various ways. In this article, we'll explore how to customize this type mapping to handle more complex scenarios, making our code more expressive and type-safe.
We'll cover several advanced type mapping techniques including union types, type aliases, mapping multiple type annotations to Python types, whole column declarations as Python types, and using enums and literals. These techniques allow us to create more expressive and precise models that better represent our domain.
The Default Type Map
When using the Mapped annotation or when providing a type to the mapped_column function, internally SQLAlchemy maps these types to an internal SQLAlchemy type.
One thing to note is that the dictionary that maps these types is completely customizable.
The default type map is represented in the documentation as follows:
from typing import Any
from typing import Dict
from typing import Type
import datetime
import decimal
import uuid
from sqlalchemy import types
# default type mapping, deriving the type for mapped_column()
# from a Mapped[] annotation
type_map: Dict[Type[Any], TypeEngine[Any]] = {
bool: types.Boolean(),
bytes: types.LargeBinary(),
datetime.date: types.Date(),
datetime.datetime: types.DateTime(),
datetime.time: types.Time(),
datetime.timedelta: types.Interval(),
decimal.Decimal: types.Numeric(),
float: types.Float(),
int: types.Integer(),
str: types.String(),
uuid.UUID: types.Uuid(),
}Nullability
Like we've seen in earlier articles, SQLAlchemy will construct its column definitions using either the mapped_column function or the Mapped type annotation (or both), and as we've seen both are optional and partially interchangeable.
In the case that you do not need any of the parameters that are provided by the mapped_column such as default_insert or primary_key, you may use the Mapped annotation alone. When using the annotation, it is possible to also mark your types with the Optional annotation.
When SQLAlchemy generates the column definitions, it will use the Optional type annotation to generate the column as NULLABLE.
Note that the same functionality is possible with mapped_column(nullable=True).
from sqlalchemy import String, Integer, ForeignKey, Text, Date, Boolean
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from typing import Optional
from datetime import date
from sqlalchemy import create_engine
engine = create_engine('sqlite:///:memory:', echo=True)
class Manuscript(DeclarativeBase):
__tablename__ = 'manuscripts'
id: Mapped[int] = mapped_column(primary_key=True)
title: Mapped[str] = mapped_column(String(100), nullable=False)
# Required fields
creation_date: Mapped[date] = mapped_column(Date)
content: Mapped[str] = mapped_column(Text)
# Optional fields
language: Mapped[Optional[str]] = mapped_column(String(50))
rarity: Mapped[Optional[str]] = mapped_column(String(20), default="Common")
# Optional notes field
notes: Mapped[Optional[str]] = mapped_column(Text)If we generate the tables and attempt to create them, we can see the output of our typing definitions.
Base.metadata.create_all(bind=engine)And as we can see, the tables that have the Mapped[Optional[...]] type annotation, the CREATE TABLE statement that is being emitted shows that all of the columns that are not optional are created with the NOT NULL constraint, while the Optional columns are missing these constraints.
It is also possible to have "conflicting" annotations and mapped_column parameters, such as having a Mapped[Optional] as well as mapped_column(nullable=False).
In this case, when you generate the model, work with it within your program with the field being set to None, until the moment where you push it into your database, at which point you either populate the field, or make sure that the field populates itself, as to avoid database exceptions.
Creating our own Types
Before we've seen the default mapping that SQLAlchemy has behind the scenes, that dictionary itself is hardcoded and not modifiable, but that does not mean we do not have the means to extend it in our own code.
When the registry object coordinates the Declarative mapping process it will first consult a local, user defined dictionary of types which may be passed as a registry.type_annotation_map parameter when constructing our DeclarativeBase superclass.
For example, if we would want to create some custom type for a timestamp with timezone:
import datetime
from sqlalchemy import TIMESTAMP
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
class SomeSuperBase(DeclarativeBase):
type_annotation_map = {
datetime.datetime: TIMESTAMP(timezone=True)
}
class Discovery(SomeSuperBase):
__tablename__ = 'discoveries'
id: Mapped[int] = mapped_column(primary_key=True)
title: Mapped[str] = mapped_column(String(100), nullable=False)
description: Mapped[str] = mapped_column(Text)
# This will use the custom type mapping (TIMESTAMP with timezone)
discovered_at: Mapped[datetime.datetime]
# This will also use the custom type mapping
documented_at: Mapped[Optional[datetime.datetime]] = mapped_column(nullable=True)And if we would like to see how the table would be created in PostgreSQL we can use the following code snippet:
from sqlalchemy.schema import CreateTable
from sqlalchemy.dialects import postgresql
print(CreateTable(Discovery.__table__).compile(dialect=postgresql.dialect()))As we can see the discovered_at and documented_at fields are being generated using the TIMESTAMP WITH TIME ZONE type, which will keep the timezone of the server along with the timestamp.
Different Database Variants
If for some reason your backend has multiple databases, that share tables, or that your software needs to be able to be deployable on top of different databases, you can give your custom types different variants for different databases.
For example if we would want to deploy our software on either Postgres or SQL Server we could define our types using the following type_annotation_map:
import datetime
from sqlalchemy import TIMESTAMP, NVARCHAR, String
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from typing import Optional
class SomeSuperBase(DeclarativeBase):
type_annotation_map = {
datetime.datetime: TIMESTAMP(timezone=True),
str: String().with_variant(NVARCHAR, "mssql")
}
class Amulet(SomeSuperBase):
__tablename__ = 'amulets'
id: Mapped[int] = mapped_column(primary_key=True)
# These will use the custom string type mapping
name: Mapped[str]
material: Mapped[str]
inscription: Mapped[Optional[str]]
# This will use the custom datetime type mapping
created_at: Mapped[datetime.datetime]
power_level: Mapped[int] = mapped_column(default=1)
is_activated: Mapped[bool] = mapped_column(default=False)And when we inspect the table definition:
from sqlalchemy.schema import CreateTable
from sqlalchemy.dialects import mssql
print(CreateTable(Amulet.__table__).compile(dialect=mssql.dialect()))We can see that our str columns are defined as NVARCHAR.
Union Types
What if we have multiple possible types for a specific column? Well, in that case we need to make sure that there is one SQL type that we can use to represent all of the python types we want to use in our union.
An example of a union type might be:
from typing import Union, Optional
from sqlalchemy import JSON
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
# Union using the pipe operator
properties_list = list[int] | list[str]
# Using the Union annotation
effect_value = Union[float, str, bool]
class Base(DeclarativeBase):
type_annotation_map = {
properties_list: postgresql.JSONB,
effect_value: JSON,
}
class Catalyst(Base):
__tablename__ = "catalysts"
id: Mapped[int] = mapped_column(primary_key=True)
components: Mapped[list[str] | list[int]]
# uses JSON
potency: Mapped[effect_value]
# uses JSON and is also nullable=True
side_effect: Mapped[effect_value | None]
# these forms all use JSON as well due to the effect_value entry
reaction_time: Mapped[float | str | bool]
stability: Mapped[Union[float, str, bool]]
storage_requirements: Mapped[Optional[float | str | bool]]Let's see how our table definitions would look like if we migrated them to PostgreSQL like we did before:
from sqlalchemy.schema import CreateTable
from sqlalchemy.dialects import postgresql
print(CreateTable(Catalyst.__table__).compile(dialect=postgresql.dialect()))As we can see the components column is marked as JSONB even though we did not explicitly tell it that it is a JSONB column nor did we use the properties_list variable, but we did use the same union type that we used to define the properties_list -> JSONB type mapping.
Type Aliases
There are multiple ways to create type aliases in Python, and SQLAlchemy allows us to use them in our type_annotation_map to map our custom types to SQL native types.
To define types, we can do the following:
from typing import NewType, List, Set
text500 = NewType("text500", str)
type SmallInt = int
type BigInt = int
type StringArray = List[str]
type IntArray = List[int]
type TagSet = Set[str]The type statement was introduced in Python 3.12 and is required to use this language feature.
Note that type is a soft keyword, which means that it is only reserved in specific cases. If you decide to move to Python 3.12 and use the type keyword as a name for a variable, it should not break your code in most cases.
As you can see, in this example, we are merely defining our types, but using them in our type_annotation_map is quite simple and similar to what we've seen before.
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String, Integer, SmallInteger, BigInteger
from sqlalchemy.dialects.postgresql import ARRAY
class ExtendedTypesBase(DeclarativeBase):
type_annotation_map = {
text500: String(500),
SmallInt: SmallInteger,
BigInt: BigInteger,
StringArray: ARRAY(String),
IntArray: ARRAY(Integer),
# Store the Set[str] type as Array in the Database
# Treat as set in the application
TagSet: ARRAY(String)
}Using these types in a Declarative Table definition is also very straightforward:
class Ingredient(ExtendedTypesBase):
__tablename__ = "ingredients"
id: Mapped[int] = mapped_column(primary_key=True)
name: Mapped[str] = mapped_column(String(100), nullable=False)
# Using text500 for detailed description
description: Mapped[text500]
# Using SmallInt for attributes with limited ranges
rarity: Mapped[SmallInt]
potency: Mapped[SmallInt]
# Using BigInt for large numbers
discovery_year: Mapped[BigInt]
# Using array types for collections
alternative_names: Mapped[StringArray]
common_measurements: Mapped[IntArray]
# Using TagSet for categories and properties
magical_properties: Mapped[TagSet]
elements: Mapped[TagSet]And if we run the CreateTable function again we will see that our types are being translated to the database types that we mapped in our ExtendedTypesBase class:
from sqlalchemy.schema import CreateTable
from sqlalchemy.dialects import postgresql
print(CreateTable(Ingredient.__table__).compile(dialect=postgresql.dialect()))This technique of creating custom type mappings can greatly enhance the expressiveness and maintainability of your SQLAlchemy models, ensuring that your Python code and database schema remain in sync while providing rich type information to both your IDE and the runtime system.
Today I'm writing and architecting software using many different technologies, and I'm always looking for the next thing to learn.