From 9a417bfc5e8c079faa6aa7320b2e959694e7c2aa Mon Sep 17 00:00:00 2001
From: Yeuoly <45712896+Yeuoly@users.noreply.github.com>
Date: Wed, 26 Feb 2025 18:45:12 +0800
Subject: [PATCH] fix: update database query and model definitions (#14415)

---
 api/controllers/service_api/wraps.py |  2 +-
 api/models/model.py                  | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py
index c746944be1..bc5dba9a78 100644
--- a/api/controllers/service_api/wraps.py
+++ b/api/controllers/service_api/wraps.py
@@ -154,7 +154,7 @@ def validate_dataset_token(view=None):
             )  # TODO: only owner information is required, so only one is returned.
             if tenant_account_join:
                 tenant, ta = tenant_account_join
-                account = Account.query.filter_by(id=ta.account_id).first()
+                account = db.session.query(Account).filter(Account.id == ta.account_id).first()
                 # Login admin
                 if account:
                     account.current_tenant = tenant
diff --git a/api/models/model.py b/api/models/model.py
index 7591e66464..cb099d5654 100644
--- a/api/models/model.py
+++ b/api/models/model.py
@@ -604,7 +604,7 @@ class InstalledApp(Base):
         return tenant
 
 
-class Conversation(Base):
+class Conversation(db.Model):  # type: ignore[name-defined]
     __tablename__ = "conversations"
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="conversation_pkey"),
@@ -839,7 +839,7 @@ class Conversation(Base):
         return self.override_model_configs is not None
 
 
-class Message(Base):
+class Message(db.Model):  # type: ignore[name-defined]
     __tablename__ = "messages"
     __table_args__ = (
         PrimaryKeyConstraint("id", name="message_pkey"),
@@ -1190,7 +1190,7 @@ class Message(Base):
         )
 
 
-class MessageFeedback(Base):
+class MessageFeedback(db.Model):  # type: ignore[name-defined]
     __tablename__ = "message_feedbacks"
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="message_feedback_pkey"),
@@ -1217,7 +1217,7 @@ class MessageFeedback(Base):
         return account
 
 
-class MessageFile(Base):
+class MessageFile(db.Model):  # type: ignore[name-defined]
     __tablename__ = "message_files"
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="message_file_pkey"),
@@ -1258,7 +1258,7 @@ class MessageFile(Base):
     created_at: Mapped[datetime] = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
 
 
-class MessageAnnotation(Base):
+class MessageAnnotation(db.Model):  # type: ignore[name-defined]
     __tablename__ = "message_annotations"
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="message_annotation_pkey"),
@@ -1327,7 +1327,7 @@ class AppAnnotationHitHistory(db.Model):  # type: ignore[name-defined]
         return account
 
 
-class AppAnnotationSetting(Base):
+class AppAnnotationSetting(db.Model):  # type: ignore[name-defined]
     __tablename__ = "app_annotation_settings"
     __table_args__ = (
         db.PrimaryKeyConstraint("id", name="app_annotation_settings_pkey"),