diff --git a/platypush/entities/_base.py b/platypush/entities/_base.py index fb38fa46..1e8fed93 100644 --- a/platypush/entities/_base.py +++ b/platypush/entities/_base.py @@ -1,11 +1,20 @@ import inspect import pathlib from datetime import datetime -from typing import Mapping, Type +from typing import Mapping, Type, Tuple import pkgutil -from sqlalchemy import Column, Index, Integer, String, DateTime, JSON, UniqueConstraint -from sqlalchemy.orm import declarative_base +from sqlalchemy import ( + Column, + Index, + Integer, + String, + DateTime, + JSON, + UniqueConstraint, + inspect as schema_inspect, +) +from sqlalchemy.orm import declarative_base, ColumnProperty Base = declarative_base() entities_registry: Mapping[Type['Entity'], Mapping] = {} @@ -40,6 +49,12 @@ class Entity(Base): 'polymorphic_on': type, } + @classmethod + @property + def columns(cls) -> Tuple[ColumnProperty]: + inspector = schema_inspect(cls) + return tuple(inspector.mapper.column_attrs) + def _discover_entity_types(): from platypush.context import get_plugin diff --git a/platypush/entities/_engine.py b/platypush/entities/_engine.py index f747027e..52d7d3d4 100644 --- a/platypush/entities/_engine.py +++ b/platypush/entities/_engine.py @@ -4,7 +4,7 @@ from threading import Thread, Event from time import time from typing import Iterable, List -from sqlalchemy import and_, or_, inspect as schema_inspect +from sqlalchemy import and_, or_ from sqlalchemy.orm import Session from ._base import Entity @@ -93,8 +93,7 @@ class EntitiesEngine(Thread): self, entities: List[Entity], existing_entities: List[Entity] ) -> List[Entity]: def merge(entity: Entity, existing_entity: Entity) -> Entity: - inspector = schema_inspect(entity.__class__) - columns = [col.key for col in inspector.mapper.column_attrs] + columns = [col.key for col in entity.columns] for col in columns: if col not in ('id', 'created_at'): setattr(existing_entity, col, getattr(entity, col))