Custom cleanup conversations and files by org

business_objects/organization.py

@@ -5,7 +5,7 @@
 
 
 from ..session import session
-from ..models import Organization, Project
+from ..models import Organization, Project, User
 from ..business_objects import project, user, general
 from ..util import prevent_sql_injection
 
@@ -62,6 +62,10 @@ def get_organization_overview_stats(
         return values[0]
 
 
+def get_user_count(organization_id: str) -> int:
+    return session.query(User).filter(User.organization_id == organization_id).count()
+
+
 def __get_organization_overview_stats_query(organization_id: str):
     return f"""
     WITH labeled_records AS (

cognition_objects/conversation.py

@@ -25,6 +25,35 @@ def get(project_id: str, conversation_id: str) -> CognitionConversation:
     )
 
 
+def get_conversations_to_clean_up() -> List[Tuple[str, str, str]]:
+    query = """
+    SELECT cc.id, cc.project_id, o.id organization_id
+    FROM cognition.conversation cc
+    INNER JOIN cognition.project cp
+        ON cc.project_id = cp.id
+    INNER JOIN (
+        SELECT o.*, NOW() - INTERVAL '1 DAY' * conversation_lifespan_days conversation_delete_by
+        FROM PUBLIC.organization o
+    ) o
+        ON cp.organization_id = o.id AND cc.created_at <= o.conversation_delete_by"""
+    return general.execute_all(query)
+
+
+def get_conversation_files_to_clean_up() -> List[Tuple[str, str, str]]:
+    query = """
+    SELECT cc.id, cc.project_id, o.id organization_id
+    FROM cognition.conversation cc
+    INNER JOIN cognition.project cp
+        ON cc.project_id = cp.id
+    INNER JOIN (
+        SELECT o.*, NOW() - INTERVAL '1 DAY' * file_lifespan_days file_delete_by
+        FROM PUBLIC.organization o
+    ) o
+        ON cp.organization_id = o.id AND cc.created_at <= o.file_delete_by
+    WHERE NOT cc.archived AND cc.has_tmp_files """
+    return general.execute_all(query)
+
+
 def get_scoped(project_id: str, conversation_id: str, user_id) -> CognitionConversation:
     return (
         session.query(CognitionConversation)

cognition_objects/file_reference.py

@@ -1,7 +1,7 @@
 from ..business_objects import general
 from ..session import session
 from ..models import FileReference
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Tuple
 
 
 def get(org_id: str, hash: str, file_size_bytes) -> FileReference:
@@ -54,6 +54,18 @@ def get_all() -> List[FileReference]:
     return session.query(FileReference).all()
 
 
+def get_all_for_cleanup() -> List[Tuple[str, str]]:
+    query = """
+    SELECT fr.id, fr.organization_id
+    FROM cognition.file_reference fr
+    INNER JOIN (
+        SELECT o.*, NOW() - INTERVAL '1 DAY' * file_lifespan_days file_delete_by
+        FROM PUBLIC.organization o
+    ) o
+        ON fr.organization_id = o.id AND fr.last_used <= o.file_delete_by"""
+    return general.execute_all(query)
+
+
 def create(
     org_id: str,
     hash: str,

models.py

@@ -141,8 +141,6 @@ class Organization(Base):
     # database entry
     is_paying = Column(Boolean, default=False)
     created_at = Column(DateTime, default=sql.func.now())
-    gdpr_compliant = Column(Boolean, default=False)
-
     projects = parent_to_child_relationship(
         Tablenames.ORGANIZATION,
         Tablenames.PROJECT,
@@ -157,6 +155,8 @@ class Organization(Base):
 
     # designed as opt out to ensure "forgotten" doesn't result in issues
     log_admin_requests = Column(String, default=AdminLogLevel.NO_GET.value)
+    conversation_lifespan_days = Column(Integer)
+    file_lifespan_days = Column(Integer, default=14)
 
 
 class User(Base):