How to create an instance of an sqlmodel with a custom type?

65 Views Asked by At

I'm using SQL Server's geography type functionality that is not natively supported by sqlalchemy, hence I am using UserDefinedType:

import sqlalchemy
from sqlalchemy import Column, FunctionElement, Integer, String, func
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.types import UserDefinedType
from sqlmodel import Field, SQLModel


class STAsText(FunctionElement):
    """
    Convert a geography into a text representation.
    The SQL Server provided way is a method of the column, e.g. @element.STAsText().
    That is not supported by sqlalchemy so we need to define a function that can be called like STAsText(element). See
    https://learn.microsoft.com/en-us/sql/t-sql/spatial-geography/spatial-types-geography?view=sql-server-ver16#b-returning-the-intersection-of-two-geography-instances
    """
    def __init__(self, column):
        self.column = column


@compiles(STAsText, 'mssql')
def compile_stastext(element, compiler, **kw):
    return '%s.STAsText()' % compiler.process(element.column)


class Geography(UserDefinedType):
    """
    Class for defining an SQL model with SQL server's geography type:
    https://learn.microsoft.com/en-us/sql/t-sql/spatial-geography/spatial-types-geography
    """
    def get_col_spec(self, **kw):
        return "geography"

    def bind_processor(self, dialect):
        def process(value):
            point_string = value
            return func.geography.STGeomFromText(point_string, 4326)

        return process

    def result_processor(self, dialect, coltype):
        def process(value):
            point_string = STAsText(value)
            return point_string

        return process


class SessionModel(SQLModel, table=True):
    __tablename__ = "sessions"

    # For the location field
    class Config:
        arbitrary_types_allowed = True

    session_key: int = Field(sa_column=Column(Integer, primary_key=True, autoincrement=True))
    location = Field(sa_column=Column(type_=Geography, nullable=False))

I verified that this model creates a correctly defined table that I can interact with manually with for example select sessions.location.STAsText() from sessions;.

However when I try to create an instance to insert in code like SessionModel(location='POINT(1.0, 2.0)', the created instance has a None location. How do I fix it? I tried adding a pydantic validator, or using the bind_expression and column_expression in the Geography class, with no luck.

0

There are 0 best solutions below