Skip to content

Instantly share code, notes, and snippets.

@zzzeek
Created February 25, 2021 17:56
Show Gist options
  • Save zzzeek/ee7f5925959b218a7fdc9a637b5efcfb to your computer and use it in GitHub Desktop.
Save zzzeek/ee7f5925959b218a7fdc9a637b5efcfb to your computer and use it in GitHub Desktop.
"""this is getting close but SQLAlchemy is still not instrumenting __init__ such
that it can intercept Pydantic's operations, in particular that list coming
in which it wants to convert to instrumented list.
this probably could be made to work with some more effort but it will
be quite hacky in the end, not really worth it unless there was
some explicit support in pydantic.
"""
from typing import List
from typing import Optional
from pydantic import BaseModel
from pydantic import constr
from sqlalchemy import Column
from sqlalchemy import ForeignKey
from sqlalchemy import Integer
from sqlalchemy import String
from sqlalchemy import Table
from sqlalchemy.engine.create import create_engine
from sqlalchemy.ext.instrumentation import InstrumentationManager
from sqlalchemy.orm import registry
from sqlalchemy.orm import relationship
from sqlalchemy.orm.attributes import del_attribute
from sqlalchemy.orm.attributes import get_attribute
from sqlalchemy.orm.attributes import set_attribute
from sqlalchemy.orm.instrumentation import is_instrumented
from sqlalchemy.orm.session import Session
reg = registry()
class MyClassState(InstrumentationManager):
def install_state(self, class_, instance, state):
object.__setattr__(instance, "_sa_custom_instrumentation", state)
def manager_getter(self, class_):
def get(cls):
return cls._default_class_manager
return get
def state_getter(self, class_):
def find(instance):
return object.__getattribute__(
instance, "_sa_custom_instrumentation"
)
return find
class PydanticAdapter(BaseModel):
__slots__ = ("__weakref__", "_sa_custom_instrumentation")
__sa_instrumentation_manager__ = MyClassState
def __getattr__(self, key):
if is_instrumented(self, key):
return get_attribute(self, key)
else:
return super().__getattr__(self, key)
def __setattr__(self, key, value):
if is_instrumented(self, key):
set_attribute(self, key, value)
else:
super().__setattr__(self, key, value)
def __delattr__(self, key):
if is_instrumented(self, key):
del_attribute(self, key)
else:
super().__delattr__(key)
@reg.mapped
class CompanyModel(PydanticAdapter):
__table__ = Table(
"companies",
reg.metadata,
Column("id", Integer, primary_key=True, nullable=False),
Column(
"public_key", String(20), index=True, nullable=False, unique=True
),
Column("name", String(63), unique=True),
)
id: Optional[int]
public_key: constr(max_length=20)
name: constr(max_length=63)
domains: List[constr(max_length=255)]
employees: "List[Employee]"
__mapper_args__ = {"properties": {"employees": relationship("Employee")}}
@reg.mapped
class Employee(PydanticAdapter):
__table__ = Table(
"employees",
reg.metadata,
Column("id", Integer, primary_key=True, nullable=False),
Column("company_id", ForeignKey("companies.id")),
Column("name", String(63), unique=True),
)
id: Optional[int]
name: constr(max_length=63)
CompanyModel.update_forward_refs()
reg.configure()
co_model = CompanyModel(
public_key="foobar",
name="Testing",
domains=["example.com", "foobar.com"],
)
co_model.employees = []
print(co_model)
co_model.employees.append(Employee(name="spongebob"))
e = create_engine("sqlite://", echo=True)
reg.metadata.create_all(e)
s = Session(e)
s.add(co_model)
s.commit()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment