diff --git a/src/kpi_manager/service/database/KpiModel.py b/src/kpi_manager/service/database/KpiModel.py
index 16844fdc05991eb29a06ce50ab3ae9fdaac3162c..9ab98e4ef1212bc8b5ed54151f2244bc26fbc91e 100644
--- a/src/kpi_manager/service/database/KpiModel.py
+++ b/src/kpi_manager/service/database/KpiModel.py
@@ -15,7 +15,8 @@
 import logging
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy import Column, Integer, String, Float, Text, ForeignKey
-from sqlalchemy.ext.declarative import declarative_base
+# from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import registry
 from sqlalchemy.orm import sessionmaker, relationship
 
 
@@ -23,7 +24,8 @@ logging.basicConfig(level=logging.INFO)
 LOGGER = logging.getLogger(__name__)
 
 # Create a base class for declarative models
-Base = declarative_base()
+Base = registry().generate_base()
+# Base = declarative_base()
 
 class Kpi(Base):
     __tablename__ = 'kpi'