How can I make DateTimeTZRange work?

[arielpontes]

First Check

• [X] I added a very descriptive title to this issue.
• [X] I used the GitHub search to find a similar issue and didn't find it.
• [X] I searched the SQLModel documentation, with the integrated search.
• [X] I already searched in Google "How to X in SQLModel" and didn't find any information.
• [X] I already read and followed all the tutorial in the docs and didn't find an answer.
• [X] I already checked if it is not related to SQLModel but to Pydantic.
• [X] I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• [X] I commit to help with one of those options 👆

Example Code

from psycopg2.extras import DateTimeTZRange
from pydantic import conlist
from sqlalchemy import Column
from sqlalchemy.dialects.postgresql import TSTZRANGE
from sqlmodel import Field

from app.models.base import BaseModel

class MyModel(BaseModel, table=True):
    # time_range: conlist(datetime, min_items=2, max_items=2) = Field(
    time_range: DateTimeTZRange = Field(
        sa_column=Column(TSTZRANGE()), nullable=False, index=False
    )

    class Config:
        arbitrary_types_allowed = True

Description

I have a model that has a DateTimeTZRange field. When I try to save an instance of this model, I get the following error:

sqlalchemy.exc.DBAPIError: (sqlalchemy.dialects.postgresql.asyncpg.Error) <class 'asyncpg.exceptions.DataError'>: invalid input for query argument $1: DateTimeTZRange(datetime.datetime(2019, ... (list, tuple or Range object expected (got type <class 'psycopg2._range.DateTimeTZRange'>))
[SQL: INSERT INTO my_model (time_range) VALUES (%s) RETURNING my_model.id]
[parameters: (DateTimeTZRange(datetime.datetime(2019, 4, 2, 23, 0, tzinfo=<UTC>), datetime.datetime(2019, 4, 2, 23, 30, tzinfo=<UTC>), '[)'))]

When I use conlist rather than DateTimeTZRange, I don't get errors but the new objects always have None in their time_range. I couldn't find any documentation on how to use custom PostgreSQL fields in SQLModel, and the documentation I found on sqlalchemy doesn't seem to work smoothly on SQLModel.

Operating System

macOS

Operating System Details

No response

SQLModel Version

0.0.4

Python Version

3.9.9

Additional Context

No response

[byrman]

I experienced no problems saving instances with the conlist version of your code. How did you initialize your model? This, for example, worked for me: MyModel(time_range=[datetime.now(timezone.utc), datetime.max])

[andrew222651]

Which version of PostgreSQL are you using?

[Zaffer]

This is a bit tricky due to Pydantic need validators etc. Basically the DateTimeTZRange type does not have the required method to convert its data to a dict, so you need to do that manually.

• The code I made below adds to Pydantic's ENCODERS_BY_TYPE
• It then creates a sub-class of DateTimeTZRange to include required validators
• This subclass also modifies the schema for presentation in the api docs correctly (with example)
• Because you are sub-classing the type you don't need arbitrary_types_allowed = True or any class Config:

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,
    )

Pydantic issues:

• https://github.com/tiangolo/fastapi/issues/1285
• https://github.com/samuelcolvin/pydantic/issues/951
• https://github.com/samuelcolvin/pydantic/issues/380

Pydantic reading:

• https://pydantic-docs.helpmanual.io/usage/schema/#modifying-schema-in-custom-fields
• https://pydantic-docs.helpmanual.io/usage/types/#custom-data-types
• https://pydantic-docs.helpmanual.io/usage/validators/

[shaug]

We had a similar question for our project, but for a DateRange adapter instead.

We adapted @Zaffer's suggestion, and it worked splendidly for persisting a value. But when the instance was refreshed/reloaded from the db, the field's resulting value on the SQLModel-derived instance was not the expected custom DateRange subclass, but rather a sqlalchemy.dialects.postgresql.ranges.Range instance (i.e., an instance of the DATERANGE used to specify the Column type).

I'm still trying to understand the mechanisms used to hydrate a row of SQL data into a Pydantic instance, but I would've expected it would attempt to cover the loaded data into the specified type on the SQLModel class.

Is this a bug, a misconfiguration, or a misunderstanding of how SQLModel/Pydantic/SQLAlchemy work?

from datetime import date

from psycopg2.extras import DateRange as DateRangeBase
from sqlalchemy.dialects.postgresql import DATERANGE

def _to_date(value: str) -> date:
    vals = map(
        int, value.strip().replace("[", "").replace(")", "").split("-", maxsplit=2)
    )
    return date(*vals)

class DateRange(DateRangeBase):
    """
    A range of dates, used for persisting in a PostgreSQL database.
    """

    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        if isinstance(v, DateRangeBase):
            return v
        if isinstance(v, str):
            lower, upper = map(_to_date, v.split(sep=",", maxsplit=1))
            bounds = v[:1] + v[-1:]
            return DateRange(lower, upper, bounds)
        raise TypeError("Type must be a string or DateRange")

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

class EventBase(SQLModel):
    date_range: DateRange = Field(sa_column=Column(DATERANGE(), nullable=False))