refac

backend/open_webui/models/automations.py

@@ -21,6 +21,7 @@ class Automation(Base):
 
     id = Column(Text, primary_key=True)
     user_id = Column(Text, nullable=False)
+    folder_id = Column(Text, nullable=True)
     name = Column(Text, nullable=False)
     data = Column(JSON, nullable=False)  # {prompt, model_id, rrule}
     meta = Column(JSON, nullable=True)
@@ -31,7 +32,10 @@ class Automation(Base):
     created_at = Column(BigInteger, nullable=False)
     updated_at = Column(BigInteger, nullable=False)
 
-    __table_args__ = (Index('ix_automation_next_run', 'next_run_at'),)
+    __table_args__ = (
+        Index('ix_automation_next_run', 'next_run_at'),
+        Index('ix_automation_user_folder', 'user_id', 'folder_id'),
+    )
 
 
 class AutomationRun(Base):
@@ -72,6 +76,7 @@ class AutomationModel(BaseModel):
 
     id: str
     user_id: str
+    folder_id: Optional[str] = None
     name: str
     data: dict
     meta: Optional[dict] = None
@@ -96,6 +101,7 @@ class AutomationRunModel(BaseModel):
 
 class AutomationForm(BaseModel):
     name: str
+    folder_id: Optional[str] = None
     data: AutomationData
     meta: Optional[dict] = None
     is_active: Optional[bool] = True
@@ -129,6 +135,7 @@ class AutomationTable:
             row = Automation(
                 id=str(uuid4()),
                 user_id=user_id,
+                folder_id=form.folder_id,
                 name=form.name,
                 data=form.data.model_dump(),
                 meta=form.meta,
@@ -164,6 +171,7 @@ class AutomationTable:
         user_id: str,
         query: Optional[str] = None,
         status: Optional[str] = None,
+        folder_id: Optional[str] = None,
         skip: int = 0,
         limit: int = 30,
         db: Optional[AsyncSession] = None,
@@ -171,6 +179,9 @@ class AutomationTable:
         async with get_async_db_context(db) as db:
             stmt = select(Automation).filter_by(user_id=user_id)
 
+            if folder_id is not None:
+                stmt = stmt.filter(Automation.folder_id == (folder_id or None))
+
             if query:
                 search = f'%{query}%'
                 # Search in name and prompt inside JSON data
@@ -216,6 +227,7 @@ class AutomationTable:
             if not row:
                 return None
             row.name = form.name
+            row.folder_id = form.folder_id
             row.data = form.data.model_dump()
             row.meta = form.meta
             if form.is_active is not None:
@@ -225,6 +237,23 @@ class AutomationTable:
             await db.commit()
             return AutomationModel.model_validate(row)
 
+    async def clear_folder_ids(
+        self,
+        user_id: str,
+        folder_ids: list[str],
+        db: Optional[AsyncSession] = None,
+    ) -> int:
+        if not folder_ids:
+            return 0
+        async with get_async_db_context(db) as db:
+            result = await db.execute(
+                update(Automation)
+                .where(Automation.user_id == user_id, Automation.folder_id.in_(folder_ids))
+                .values(folder_id=None, updated_at=int(time.time_ns()))
+            )
+            await db.commit()
+            return result.rowcount or 0
+
     async def toggle(
         self,
         id: str,