diff --git a/business_objects/organization.py b/business_objects/organization.py index 2a0fcab..7e2e462 100644 --- a/business_objects/organization.py +++ b/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 ( diff --git a/cognition_objects/conversation.py b/cognition_objects/conversation.py index cc3d93d..8545919 100644 --- a/cognition_objects/conversation.py +++ b/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) diff --git a/cognition_objects/file_reference.py b/cognition_objects/file_reference.py index cec6da0..15b7a4e 100644 --- a/cognition_objects/file_reference.py +++ b/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, diff --git a/models.py b/models.py index f4cbe12..33b4592 100644 --- a/models.py +++ b/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):