How to use Arrow type in FastAPI response schema?

2.6k Views Asked by At

I want to use Arrow type in FastAPI response because I am using it already in SQLAlchemy model (thanks to sqlalchemy_utils).

I prepared a small self-contained example with a minimal FastAPI app. I expect that this app return product1 data from database.

Unfortunately the code below gives exception:

Exception has occurred: FastAPIError
Invalid args for response field! Hint: check that <class 'arrow.arrow.Arrow'> is a valid pydantic field type
import sqlalchemy
import uvicorn
from arrow import Arrow
from fastapi import FastAPI
from pydantic import BaseModel
from sqlalchemy import Column, Integer, Text, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import ArrowType

app = FastAPI()

engine = sqlalchemy.create_engine('sqlite:///db.db')
Base = declarative_base()

class Product(Base):
    __tablename__ = "product"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=True)
    created_at = Column(ArrowType(timezone=True), nullable=False, server_default=func.now())

Base.metadata.create_all(engine)


Session = sessionmaker(bind=engine)
session = Session()

product1 = Product(name="ice cream")
product2 = Product(name="donut")
product3 = Product(name="apple pie")

session.add_all([product1, product2, product3])
session.commit()


class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True


@app.get('/', response_model=ProductResponse)
async def return_product():

    product = session.query(Product).filter(Product.id == 1).first()

    return product

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)

requirements.txt:

sqlalchemy==1.4.23
sqlalchemy_utils==0.37.8
arrow==1.1.1
fastapi==0.68.1
uvicorn==0.15.0

This error is already discussed in those FastAPI issues:

  1. https://github.com/tiangolo/fastapi/issues/1186
  2. https://github.com/tiangolo/fastapi/issues/2382

One possible workaround is to add this code (source):

from pydantic import BaseConfig
BaseConfig.arbitrary_types_allowed = True

It is enough to put it just above @app.get('/'..., but it can be put even before app = FastAPI()

The problem with this solution is that output of GET endpoint will be:

// 20210826001330
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": {
    "_datetime": "2021-08-25T21:38:01+00:00"
  }
}

instead of desired:

// 20210826001330
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": "2021-08-25T21:38:01+00:00"
}
6

There are 6 best solutions below

1
On BEST ANSWER

The solution is to monkeypatch pydantic's ENCODERS_BY_TYPE so it knows how to convert Arrow object so it can be accepted by json format:

from arrow import Arrow
from pydantic.json import ENCODERS_BY_TYPE
ENCODERS_BY_TYPE |= {Arrow: str}

Setting BaseConfig.arbitrary_types_allowed = True is also necessary.

Result:

// 20220514022717
// http://localhost:8000/

{
  "id": 1,
  "name": "ice cream",
  "created_at": "2022-05-14T00:20:11+00:00"
}

Full code:

import sqlalchemy
import uvicorn
from arrow import Arrow
from fastapi import FastAPI
from pydantic import BaseModel
from sqlalchemy import Column, Integer, Text, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import ArrowType

from pydantic.json import ENCODERS_BY_TYPE
ENCODERS_BY_TYPE |= {Arrow: str}

from pydantic import BaseConfig
BaseConfig.arbitrary_types_allowed = True

app = FastAPI()

engine = sqlalchemy.create_engine('sqlite:///db.db')
Base = declarative_base()

class Product(Base):
    __tablename__ = "product"
    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(Text, nullable=True)
    created_at = Column(ArrowType(timezone=True), nullable=False, server_default=func.now())

Base.metadata.create_all(engine)


Session = sessionmaker(bind=engine)
session = Session()

product1 = Product(name="ice cream")
product2 = Product(name="donut")
product3 = Product(name="apple pie")

session.add_all([product1, product2, product3])
session.commit()


class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True


@app.get('/', response_model=ProductResponse)
async def return_product():

    product = session.query(Product).filter(Product.id == 1).first()

    return product

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)
0
On

Add a custom function with the @validator decorator that returns the desired _datetime of the object:

class ProductResponse(BaseModel):
    id: int
    name: str
    created_at: Arrow

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True

    @validator("created_at")
    def format_datetime(cls, value):
        return value._datetime

Tested on local, seems to be working:

$ curl -s localhost:8000 | jq
{
  "id": 1,
  "name": "ice cream",
  "created_at": "2021-12-02T08:25:10+00:00"
}
1
On

Here is a code example where you do not need class Config and can work for any type by creating your own subclass with validators:

from psycopg2.extras import DateTimeTZRange as DateTimeTZRangeBase
from sqlalchemy.dialects.postgresql import TSTZRANGE
from sqlmodel import (
    Column,
    Field,
    Identity,
    SQLModel,
)

from pydantic.json import ENCODERS_BY_TYPE

ENCODERS_BY_TYPE |= {DateTimeTZRangeBase: str}


class DateTimeTZRange(DateTimeTZRangeBase):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        if isinstance(v, str):
            lower = v.split(", ")[0][1:].strip().strip()
            upper = v.split(", ")[1][:-1].strip().strip()
            bounds = v[:1] + v[-1:]
            return DateTimeTZRange(lower, upper, bounds)
        elif isinstance(v, DateTimeTZRangeBase):
            return v
        raise TypeError("Type must be string or DateTimeTZRange")

    @classmethod
    def __modify_schema__(cls, field_schema):
        field_schema.update(type="string", example="[2022,01,01, 2022,02,02)")


class EventBase(SQLModel):
    __tablename__ = "event"
    timestamp_range: DateTimeTZRange = Field(
        sa_column=Column(
            TSTZRANGE(),
            nullable=False,
        ),
    )


class Event(EventBase, table=True):
    id: int | None = Field(
        default=None,
        sa_column_args=(Identity(always=True),),
        primary_key=True,
        nullable=False,
    )

link to Github issue: https://github.com/tiangolo/sqlmodel/issues/235#issuecomment-1162063590

0
On

Here is CustomArrowType for validation and serializing on Pydantic v1 and v2

# pydantic < 2.0
import arrow
from pydantic import BaseModel


class ArrowPydanticV1(arrow.Arrow):
    @classmethod
    def __get_validators__(cls):
        yield cls.pydantic_validate

    @classmethod
    def __modify_schema__(cls, field_schema):
        field_schema.update(
            examples=["2024-02-06 13:38:18", "2024-02-06 13:38:30+00:00"],
        )

    @classmethod
    def pydantic_validate(cls, v):
        try:
            arr = arrow.get(v)
            return arr
        except Exception as e:
            raise ValueError(f"Arrow could not parse {v!r}: {e!r}")

    def __repr__(self):
        return f"PydanticV1Arrow({super().__repr__()})"
# pydantic >= 2.0
from typing import Any

import arrow

from pydantic import (
    BaseModel,
    GetCoreSchemaHandler,
    GetJsonSchemaHandler,
)
from pydantic.json_schema import JsonSchemaValue
from pydantic_core import core_schema
from typing_extensions import Annotated


class ArrowPydanticV2(arrow.Arrow):
    @classmethod
    def __get_pydantic_core_schema__(
        cls,
        _source_type: Any,
        _handler: GetCoreSchemaHandler,
    ) -> core_schema.CoreSchema:
        
        def validate_by_arrow(value) -> arrow.Arrow:
            try:
                arr = arrow.get(value)
                return arr
            except Exception as e:
                raise ValueError(f"Arrow can not parse")

        def arrow_serialization(value: Any, _, info) -> str | arrow.Arrow:
            if info.mode == "json":
                return value.format("YYYY-MM-DDTHH:mm:ss.SSSSSSZZ")
            return value   

        return core_schema.no_info_after_validator_function(
            function=validate_by_arrow,
            schema=core_schema.str_schema(),
            serialization=core_schema.wrap_serializer_function_ser_schema(arrow_serialization, info_arg=True),
        )

class Model(BaseModel):
    datetime: ArrowPydanticV2

# test
m = Model(datetime="2024-02-06 11:38:18+00:00")
print(m)
print(m.datetime)
print(m.model_dump(mode="python"))
print(m.model_dump_json())
print(m.model_dump(mode="json"))

>>>
datetime=<Arrow [2024-02-06T11:38:18+00:00]>
2024-02-06T11:38:18+00:00
{'datetime': {'datetime': <Arrow [2024-02-06T11:38:18+00:00]>}}
{"datetime":"2024-02-06 11:38:18.000000+00:00"}
{'datetime': '2024-02-06T11:38:18Z'}
0
On

Another simple aproach for Pydantic V2 with Annotated types

# Optional, you can just use `float` instead
TimeStamp = Annotated[
    float,
    WithJsonSchema(
        {"type": "float", "example": arrow.now().timestamp()}, mode="serialization"
    ),
]

PydanticArrow = Annotated[
    arrow.Arrow,
    PlainSerializer(lambda x: x.timestamp(), return_type=TimeStamp, when_used="json"),
]


class SimpleSchema(BaseModel):
    id: UUID 
    create_time: PydanticArrow
0
On

Recently I've faced similar issue and answer provided by @Karol Zlot seems to be obsolete - FastAPI was throwing JSON Schema error:

ValueError: Value not declarable with JSON Schema, field: name='created_at' type=ArrowType required=True

Below code seems to work:

import datetime

class ArrowType(datetime):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        return v._datetime

class Domain(DomainBase):
    id: int
    created_at: ArrowType
    updated_at: ArrowType