how to set up ORM models based on dataclasses

user_11061257 0
Viewed 738

I have a bunch of dataclasses and want to use them as ORM Models for my database. To reach my goal i made decorator @sa_orm which maps each decorated dataclass to metadata.

Problems:

  1. some fields are dataclass.
  2. some fields are list of dataclass objects

OneToOne (add child_id column to parent table) relationship solve first problem, but to solve second one i need set OneToMany relationship (add parent_id to child table).

Here is simplified code:

from dataclasses import dataclass, field, is_dataclass
from sqlalchemy import Table, MetaData, Column, Integer, String, create_engine, ForeignKey
from sqlalchemy.orm import sessionmaker, mapper, relationship
from sqlalchemy.ext.declarative import declarative_base
from typing import Optional, List
metadata = MetaData()
Base = declarative_base(metadata=metadata)

def sa_orm(metadata):
    def decorator(cls):
        columns = [Column('id', Integer, primary_key=True, autoincrement=True)]
        children = {}
        for field_name, meta in cls.__dataclass_fields__.items():
            try:
                field_class = meta.type.__args__[0]  # If Optional or List
            except AttributeError:
                field_class = meta.type  # If not Optional or List
            if is_dataclass(field_class):
                field_class_name = field_class.__name__
                columns.append(Column(f"{field_class_name}_id".lower(), Integer, ForeignKey(f'{field_class_name}.id'.lower())))
                children[field_name] = relationship(field_class)
            else:
                columns.append(Column(field_name, String))
        table = Table(cls.__name__.lower(), metadata,
                      *columns)
        mapper(cls, table, properties=children)

        return cls
    return decorator


@sa_orm(metadata)
@dataclass
class Profession:
    name: Optional[str] = field(default_factory=str)

@sa_orm(metadata)
@dataclass
class Person:
    profession: Optional[Profession] = field(default=None)
    name: str = field(default_factory=str)

@sa_orm(metadata)
@dataclass
class Company:
    employee: List[Person] = field(default_factory=list)
    name: str = field(default_factory=str)

engine = create_engine(f'postgresql://user:pass@localhost:5432', echo=True)

metadata.create_all(engine)

session = sessionmaker(bind=engine)()

developer = Profession(name="developer")
alex = Person(name="Alexander", profession=developer)
peter = Person(name="Peter")
google = Company(name="Google", employee=[alex, peter])  # OneToMany

session.add(alex)  # this works - insert alex to db, insert developer to db
session.add(peter)  # also works - insert peter to db
session.commit()

session.add(google) # this fails
session.close()

Is there a way to set column like "company_id" to person table from Company decorator or another way to link these tables without changing dataclasses?

I have hundreds of such dataclasses scattered across dozens of files, maybe there is better solution for setting database based on dataclasses?

python-dataclasses sqlalchemy python
1 Answers

SqlAlchemy appears to support using dataclasses via either Imperative or Declarative methods.

Since you seem to be using declarative_base, here is the example from the declarative docs that may help you achieve what you are looking for:

from __future__ import annotations

from dataclasses import dataclass
from dataclasses import field
from typing import List
from typing import Optional

from sqlalchemy import Column
from sqlalchemy import ForeignKey
from sqlalchemy import Integer
from sqlalchemy import String
from sqlalchemy import Table
from sqlalchemy.orm import registry
from sqlalchemy.orm import relationship

mapper_registry = registry()


@mapper_registry.mapped
@dataclass
class User:
    __table__ = Table(
        "user",
        mapper_registry.metadata,
        Column("id", Integer, primary_key=True),
        Column("name", String(50)),
        Column("fullname", String(50)),
        Column("nickname", String(12)),
    )
    id: int = field(init=False)
    name: Optional[str] = None
    fullname: Optional[str] = None
    nickname: Optional[str] = None
    addresses: List[Address] = field(default_factory=list)

    __mapper_args__ = {   # type: ignore
        "properties" : {
            "addresses": relationship("Address")
        }
    }

@mapper_registry.mapped
@dataclass
class Address:
    __table__ = Table(
        "address",
        mapper_registry.metadata,
        Column("id", Integer, primary_key=True),
        Column("user_id", Integer, ForeignKey("user.id")),
        Column("email_address", String(50)),
    )
    id: int = field(init=False)
    user_id: int = field(init=False)
    email_address: Optional[str] = None
user_30333860
Related