From fe0f3202fef81fb8a223de37f2caad8973c7d10b Mon Sep 17 00:00:00 2001
From: Fabio Manganiello <info@fabiomanganiello.com>
Date: Tue, 5 Apr 2022 23:04:19 +0200
Subject: [PATCH] columns should be a property of the Entity object

---
 platypush/entities/_base.py   | 21 ++++++++++++++++++---
 platypush/entities/_engine.py |  5 ++---
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/platypush/entities/_base.py b/platypush/entities/_base.py
index fb38fa460..1e8fed934 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 f747027ee..52d7d3d4a 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))