compub/src/compub/models/geography.py

117 lines
3.4 KiB
Python

from typing import Any, Optional, TYPE_CHECKING
from pydantic import validator
from sqlalchemy.sql.schema import Column, Index
from sqlalchemy.sql.sqltypes import Unicode
from sqlalchemy_utils.primitives.country import Country
from sqlalchemy_utils.types import CountryType
from sqlmodel.main import Field, Relationship
from .base import AbstractBase
if TYPE_CHECKING:
from .companies import Company
from .courts import Court
__all__ = [
'StateProvince',
'City',
'Street',
'Address',
]
class StateProvince(AbstractBase, table=True):
__tablename__ = 'state_province'
__table_args__ = (
Index('ux_state_province_name_country', 'name', 'country', unique=True),
)
# Fields
country: str = Field(sa_column=Column(CountryType, nullable=False, index=True))
name: str = Field(max_length=255, sa_column=Column(Unicode(255), nullable=False, index=True))
# Relationships
cities: list['City'] = Relationship(
back_populates='state_province', sa_relationship_kwargs={'lazy': 'selectin'}
)
courts: list['Court'] = Relationship(
back_populates='state_province', sa_relationship_kwargs={'lazy': 'selectin'}
)
@validator('country', pre=True)
def country_as_uppercase_string(cls, v: Any) -> str:
if isinstance(v, Country):
return v.code
if isinstance(v, str):
return v.upper()
raise TypeError
class City(AbstractBase, table=True):
__tablename__ = 'city'
__table_args__ = (
Index('ux_city_name_zip_state', 'name', 'zip_code', 'state_province_id', unique=True),
)
# Fields
zip_code: str = Field(max_length=5, nullable=False, index=True)
name: str = Field(max_length=255, sa_column=Column(Unicode(255), nullable=False, index=True))
# Relationships
state_province_id: Optional[int] = Field(
foreign_key='state_province.id', default=None, nullable=False, index=True
)
state_province: Optional[StateProvince] = Relationship(
back_populates='cities', sa_relationship_kwargs={'lazy': 'selectin'}
)
streets: list['Street'] = Relationship(
back_populates='city', sa_relationship_kwargs={'lazy': 'selectin'}
)
class Street(AbstractBase, table=True):
__tablename__ = 'street'
# Fields
name: str = Field(max_length=255, sa_column=Column(Unicode(255), nullable=False, index=True))
# Relationships
city_id: Optional[int] = Field(
foreign_key='city.id', default=None, nullable=False, index=True
)
city: Optional[City] = Relationship(
back_populates='streets', sa_relationship_kwargs={'lazy': 'selectin'}
)
addresses: list['Address'] = Relationship(
back_populates='street', sa_relationship_kwargs={'lazy': 'selectin'}
)
class Address(AbstractBase, table=True):
__tablename__ = 'address'
# Fields
house_number: str = Field(max_length=8, nullable=False)
supplement: str = Field(max_length=255)
# Relationships
street_id: Optional[int] = Field(
foreign_key='street.id', default=None, nullable=False, index=True
)
street: Optional[Street] = Relationship(
back_populates='addresses', sa_relationship_kwargs={'lazy': 'selectin'}
)
companies: list['Company'] = Relationship(
back_populates='address', sa_relationship_kwargs={'lazy': 'selectin'}
)
courts: list['Court'] = Relationship(
back_populates='address', sa_relationship_kwargs={'lazy': 'selectin'}
)