Initial commit
This commit is contained in:
commit
f4cf78603a
25 changed files with 2895 additions and 0 deletions
271
src/treetrail/models.py
Normal file
271
src/treetrail/models.py
Normal file
|
|
@ -0,0 +1,271 @@
|
|||
|
||||
from typing import Annotated, Any, Literal
|
||||
from datetime import datetime
|
||||
import uuid
|
||||
|
||||
from sqlalchemy.ext.mutable import MutableDict
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.orm import joinedload, QueryableAttribute
|
||||
from geoalchemy2 import Geometry, WKBElement # type: ignore
|
||||
from sqlmodel import (SQLModel, Field, String, Relationship, JSON,
|
||||
select)
|
||||
import pandas as pd
|
||||
import geopandas as gpd # type: ignore
|
||||
|
||||
from treetrail.utils import pandas_query, geopandas_query
|
||||
from treetrail.config import Map, conf, App
|
||||
from treetrail.database import db_session
|
||||
|
||||
class BaseModel(SQLModel):
|
||||
@classmethod
|
||||
def selectinload(cls) -> list[Literal['*'] | QueryableAttribute[Any]]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
async def get_df(cls, **kwargs) -> pd.DataFrame:
|
||||
return await cls._get_df(pandas_query, **kwargs)
|
||||
|
||||
@classmethod
|
||||
async def get_gdf(cls, **kwargs) -> gpd.GeoDataFrame:
|
||||
return await cls._get_df(geopandas_query, model=cls, **kwargs) # type: ignore
|
||||
|
||||
@classmethod
|
||||
async def _get_df(cls, method, *,
|
||||
where=None, with_related=True, with_only_columns=[],
|
||||
simplify_tolerance: float | None=None,
|
||||
preserve_topology: bool | None=None,
|
||||
**kwargs) -> pd.DataFrame | gpd.GeoDataFrame:
|
||||
async with db_session() as session:
|
||||
if len(with_only_columns) == 0:
|
||||
query = select(cls)
|
||||
else:
|
||||
columns = set(with_only_columns)
|
||||
# TODO: user SQLModel model_fields instead of __table__
|
||||
columns.add(*(col.name for col in cls.__table__.primary_key.columns)) # type: ignore
|
||||
query = select(*(getattr(cls, col) for col in columns))
|
||||
if where is not None:
|
||||
query = query.where(where)
|
||||
## Get the joined tables
|
||||
joined_tables = cls.selectinload()
|
||||
if with_related and len(joined_tables) > 0:
|
||||
query = query.options(*(joinedload(jt) for jt in joined_tables))
|
||||
df = await session.run_sync(method, query, **kwargs)
|
||||
if method is geopandas_query and simplify_tolerance is not None:
|
||||
df['geom'] = df['geom'].simplify(
|
||||
simplify_tolerance / conf.geo.simplify_geom_factor,
|
||||
preserve_topology=(conf.geo.simplify_preserve_topology
|
||||
if preserve_topology is None
|
||||
else preserve_topology)
|
||||
)
|
||||
## Chamge column names to reflect the joined tables
|
||||
## Leave the first columns unchanged, as their names come straight
|
||||
## from the model's fields
|
||||
joined_columns = list(df.columns[len(cls.model_fields):])
|
||||
renames: dict[str, str] = {}
|
||||
## Match colum names with the joined tables
|
||||
## Important: this assumes that orders of the joined tables
|
||||
## and their columns is preserved by pandas' read_sql
|
||||
for joined_table in joined_tables:
|
||||
target = joined_table.property.target # type: ignore
|
||||
target_name = target.name
|
||||
for col in target.columns:
|
||||
## Pop the column from the colujmn list and make a new name
|
||||
renames[joined_columns.pop(0)] = f'{target.schema}_{target_name}_{col.name}'
|
||||
df.rename(columns=renames, inplace=True)
|
||||
## Finally, set the index of the df as the index of cls
|
||||
df.set_index([c.name for c in cls.__table__.primary_key.columns], # type: ignore
|
||||
inplace=True)
|
||||
return df
|
||||
|
||||
class TreeTrail(BaseModel, table=True):
|
||||
__tablename__: str = 'tree_trail' # type: ignore
|
||||
|
||||
tree_id: uuid.UUID | None = Field(
|
||||
default=None,
|
||||
foreign_key='tree.id',
|
||||
primary_key=True
|
||||
)
|
||||
trail_id: int | None = Field(
|
||||
default=None,
|
||||
foreign_key='trail.id',
|
||||
primary_key=True
|
||||
)
|
||||
|
||||
|
||||
class Trail(BaseModel, table=True):
|
||||
__tablename__: str = "trail" # type: ignore
|
||||
|
||||
id: int = Field(primary_key=True)
|
||||
name: str
|
||||
description: str
|
||||
create_date: datetime = Field(default_factory=datetime.now)
|
||||
geom: Annotated[str, WKBElement] = Field(
|
||||
sa_type=Geometry('LINESTRING', srid=4326, dimension=2),
|
||||
)
|
||||
photo: str = Field(sa_type=String(250)) # type: ignore
|
||||
trees: list['Tree'] = Relationship(
|
||||
link_model=TreeTrail,
|
||||
back_populates="trails")
|
||||
viewable_role_id: str | None = Field(foreign_key='role.name', index=True)
|
||||
viewable_role: 'Role' = Relationship(back_populates='viewable_trails')
|
||||
|
||||
# __mapper_args__ = {"eager_defaults": True}
|
||||
|
||||
life_stages = ('Y', 'MA', 'M', 'OM', 'A')
|
||||
|
||||
|
||||
class Tree(BaseModel, table=True):
|
||||
__tablename__: str = "tree" # type: ignore
|
||||
|
||||
id: uuid.UUID | None = Field(
|
||||
default_factory=uuid.uuid1,
|
||||
primary_key=True,
|
||||
index=True,
|
||||
nullable=False,
|
||||
)
|
||||
create_date: datetime = Field(default_factory=datetime.now)
|
||||
# ALTER TABLE tree ADD CONSTRAINT tree_plant_id_fkey FOREIGN KEY (plantekey_id) REFERENCES plant(id); # noqa: E501
|
||||
plantekey_id: str = Field(foreign_key='plant.id')
|
||||
geom: Annotated[str, WKBElement] = Field(
|
||||
sa_type=Geometry('POINT', srid=4326, dimension=2))
|
||||
photo: str | None = Field(sa_type=String(250)) # type: ignore
|
||||
height: float | None
|
||||
comments: str | None
|
||||
# ALTER TABLE public.tree ADD contributor_id varchar(50) NULL;
|
||||
# ALTER TABLE public.tree ADD CONSTRAINT contributor_fk FOREIGN KEY (contributor_id) REFERENCES public."user"(username);
|
||||
contributor_id: str = Field(foreign_key='user.username', index=True)
|
||||
contributor: 'User' = Relationship()
|
||||
viewable_role_id: str | None = Field(foreign_key='role.name', index=True)
|
||||
viewable_role: 'Role' = Relationship(back_populates='viewable_trees')
|
||||
|
||||
# CREATE EXTENSION hstore;
|
||||
# ALTER TABLE tree ADD COLUMN data JSONB;
|
||||
data: dict = Field(sa_type=MutableDict.as_mutable(JSONB), # type: ignore
|
||||
default_factory=dict) # type: ignore
|
||||
trails: list[Trail] = Relationship(
|
||||
link_model=TreeTrail,
|
||||
back_populates="trees")
|
||||
|
||||
__mapper_args__ = {"eager_defaults": True}
|
||||
|
||||
@classmethod
|
||||
def get_tree_insert_params(
|
||||
cls,
|
||||
plantekey_id: str,
|
||||
lng, lat,
|
||||
username,
|
||||
details: dict,
|
||||
) -> dict:
|
||||
params = {
|
||||
'plantekey_id': plantekey_id,
|
||||
'geom': f'POINT({lng} {lat})',
|
||||
'contributor_id': username
|
||||
}
|
||||
## Consume some details in their respective field...
|
||||
if p:=details.pop('comments', None):
|
||||
params['comments'] = p
|
||||
if p:=details.pop('height', None):
|
||||
params['height'] = p
|
||||
# ... and store the rest in data
|
||||
params['data'] = {k: v for k, v in details.items() if v}
|
||||
return params
|
||||
|
||||
|
||||
class UserRoleLink(SQLModel, table=True):
|
||||
__tablename__: str = 'roles_users' # type: ignore
|
||||
user_id: str | None = Field(
|
||||
default=None,
|
||||
foreign_key='user.username',
|
||||
primary_key=True
|
||||
)
|
||||
role_id: str | None = Field(
|
||||
default=None,
|
||||
foreign_key='role.name',
|
||||
primary_key=True
|
||||
)
|
||||
|
||||
|
||||
class UserBase(BaseModel):
|
||||
username: str = Field(sa_type=String(50), primary_key=True) # type: ignore
|
||||
full_name: str | None = None
|
||||
email: str | None = None
|
||||
|
||||
|
||||
class User(UserBase, table=True):
|
||||
__tablename__: str = "user" # type: ignore
|
||||
roles: list["Role"] = Relationship(back_populates="users",
|
||||
link_model=UserRoleLink)
|
||||
password: str
|
||||
disabled: bool = False
|
||||
|
||||
|
||||
class UserWithRoles(UserBase):
|
||||
roles: list['Role']
|
||||
|
||||
|
||||
class Role(BaseModel, table=True):
|
||||
__tablename__: str = "role" # type: ignore
|
||||
name: str = Field(sa_type=String(50), primary_key=True) # type: ignore
|
||||
users: list[User] = Relationship(back_populates="roles",
|
||||
link_model=UserRoleLink)
|
||||
viewable_trees: list[Tree] = Relationship(back_populates='viewable_role')
|
||||
viewable_zones: list['Zone'] = Relationship(back_populates='viewable_role')
|
||||
viewable_trails: list[Trail] = Relationship(back_populates='viewable_role')
|
||||
|
||||
|
||||
class POI(BaseModel, table=True):
|
||||
__tablename__: str = "poi" # type: ignore
|
||||
|
||||
id: int = Field(primary_key=True)
|
||||
name: str = Field(sa_column=String(200)) # type: ignore
|
||||
description: str | None = None
|
||||
create_date: datetime = Field(default_factory=datetime.now)
|
||||
geom: Annotated[str, WKBElement] = Field(
|
||||
sa_type=Geometry('POINTZ', srid=4326, dimension=3))
|
||||
photo: str = Field(sa_column=String(250)) # type: ignore
|
||||
type: str = Field(sa_column=String(25)) # type: ignore
|
||||
data: dict = Field(sa_type=MutableDict.as_mutable(JSONB), # type: ignore
|
||||
default_factory=dict) # type: ignore
|
||||
|
||||
|
||||
class Zone(BaseModel, table=True):
|
||||
__tablename__: str = "zone" # type: ignore
|
||||
id: int = Field(primary_key=True)
|
||||
name: str = Field(sa_type=String(200)) # type:ignore
|
||||
description: str
|
||||
create_date: datetime = Field(default_factory=datetime.now)
|
||||
geom: Annotated[str, WKBElement] = Field(
|
||||
sa_type=Geometry('MULTIPOLYGON', srid=4326))
|
||||
photo: str | None = Field(sa_type=String(250)) # type:ignore
|
||||
type: str = Field(sa_type=String(30)) # type:ignore
|
||||
data: dict | None = Field(sa_type=MutableDict.as_mutable(JSONB), # type:ignore
|
||||
default_factory=dict) # type:ignore
|
||||
viewable_role_id: str | None = Field(foreign_key='role.name', index=True)
|
||||
viewable_role: 'Role' = Relationship(back_populates='viewable_zones')
|
||||
|
||||
|
||||
class MapStyle(BaseModel, table=True):
|
||||
__tablename__: str = "map_style" # type: ignore
|
||||
|
||||
id: int = Field(primary_key=True)
|
||||
layer: str = Field(sa_type=String(100), nullable=False) # type:ignore
|
||||
paint: dict[str, Any] | None = Field(sa_type=JSON(none_as_null=True)) # type:ignore
|
||||
layout: dict[str, Any] | None = Field(sa_type=JSON(none_as_null=True)) # type:ignore
|
||||
|
||||
|
||||
class VersionedComponent(BaseModel):
|
||||
version: str
|
||||
|
||||
|
||||
class BaseMapStyles(BaseModel):
|
||||
embedded: list[str]
|
||||
external: dict[str, str]
|
||||
|
||||
|
||||
class Bootstrap(BaseModel):
|
||||
client: VersionedComponent
|
||||
server: VersionedComponent
|
||||
app: App
|
||||
user: UserWithRoles | None # type:ignore
|
||||
map: Map
|
||||
baseMapStyles: BaseMapStyles
|
||||
Loading…
Add table
Add a link
Reference in a new issue