Desktop migration: Rust backend → Python backend (#5302)

backend/database/advice.py

@@ -0,0 +1,115 @@
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import List, Dict, Any, Optional
+
+from google.cloud import firestore
+
+from ._client import db
+
+logger = logging.getLogger(__name__)
+
+USERS_COLLECTION = 'users'
+ADVICE_SUBCOLLECTION = 'advice'
+
+
+def _collection_ref(uid: str):
+    return db.collection(USERS_COLLECTION).document(uid).collection(ADVICE_SUBCOLLECTION)
+
+
+def create_advice(uid: str, data: Dict[str, Any]) -> Dict[str, Any]:
+    """Create a new advice document. Returns the created document with id."""
+    advice_id = str(uuid.uuid4())
+    now = datetime.now(timezone.utc)
+
+    doc_data = {
+        'content': data['content'],
+        'category': data.get('category', 'other'),
+        'confidence': data.get('confidence', 0.5),
+        'is_read': False,
+        'is_dismissed': False,
+        'created_at': now,
+    }
+    for optional_field in ('reasoning', 'source_app', 'context_summary', 'current_activity'):
+        if data.get(optional_field) is not None:
+            doc_data[optional_field] = data[optional_field]
+
+    _collection_ref(uid).document(advice_id).set(doc_data)
+
+    doc_data['id'] = advice_id
+    return doc_data
+
+
+def get_advice(
+    uid: str,
+    limit: int = 100,
+    offset: int = 0,
+    category: Optional[str] = None,
+    include_dismissed: bool = False,
+) -> List[Dict[str, Any]]:
+    """Query advice, ordered by created_at DESC."""
+    query = _collection_ref(uid).order_by('created_at', direction=firestore.Query.DESCENDING)
+
+    if not include_dismissed:
+        query = query.where(filter=firestore.FieldFilter('is_dismissed', '==', False))
+    if category:
+        query = query.where(filter=firestore.FieldFilter('category', '==', category))
+
+    query = query.offset(offset).limit(limit)
+
+    results = []
+    for doc in query.stream():
+        data = doc.to_dict()
+        data['id'] = doc.id
+        results.append(data)
+    return results
+
+
+def update_advice(uid: str, advice_id: str, data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """Update an advice document (is_read, is_dismissed). Returns updated doc."""
+    doc_ref = _collection_ref(uid).document(advice_id)
+
+    update_data = {'updated_at': datetime.now(timezone.utc)}
+    if 'is_read' in data:
+        update_data['is_read'] = data['is_read']
+    if 'is_dismissed' in data:
+        update_data['is_dismissed'] = data['is_dismissed']
+
+    try:
+        doc_ref.update(update_data)
+    except Exception as e:
+        if hasattr(e, 'code') and e.code == 404:
+            return None
+        raise
+
+    doc = doc_ref.get()
+    if doc.exists:
+        result = doc.to_dict()
+        result['id'] = doc.id
+        return result
+    return None
+
+
+def delete_advice(uid: str, advice_id: str) -> bool:
+    """Delete an advice document. Returns True on success."""
+    _collection_ref(uid).document(advice_id).delete()
+    return True
+
+
+def mark_all_advice_read(uid: str) -> int:
+    """Mark all unread, non-dismissed advice as read. Returns count of marked items."""
+    query = _collection_ref(uid).where(
+        filter=firestore.FieldFilter('is_dismissed', '==', False)
+    ).where(
+        filter=firestore.FieldFilter('is_read', '==', False)
+    ).limit(1000)
+
+    count = 0
+    now = datetime.now(timezone.utc)
+    for doc in query.stream():
+        try:
+            doc.reference.update({'is_read': True, 'updated_at': now})
+            count += 1
+        except Exception:
+            logger.warning('Failed to mark advice %s as read for uid=%s', doc.id, uid)
+    return count

backend/database/chat.py

@@ -468,10 +468,69 @@ def delete_chat_session(uid, chat_session_id):
     session_ref.delete()
 
 
-def add_message_to_chat_session(uid: str, chat_session_id: str, message_id: str):
+def get_chat_sessions(
+    uid: str, app_id: Optional[str] = None, limit: int = 50, offset: int = 0, starred: Optional[bool] = None
+):
+    """List chat sessions with optional filters.
+
+    Note: Client-side sort + slice because Firestore composite indexes would be
+    needed for every filter combination. Acceptable for desktop users (low session
+    counts). Revisit with server-side ordering if session counts grow large.
+    """
+    sessions_ref = db.collection('users').document(uid).collection('chat_sessions')
+    if app_id is not None:
+        sessions_ref = sessions_ref.where(filter=FieldFilter('plugin_id', '==', app_id))
+    if starred is not None:
+        sessions_ref = sessions_ref.where(filter=FieldFilter('starred', '==', starred))
+    sessions = [doc.to_dict() for doc in sessions_ref.stream()]
+    sessions.sort(key=lambda s: s.get('updated_at', s.get('created_at', datetime.min)), reverse=True)
+    return sessions[offset : offset + limit]
+
+
+def update_chat_session(uid: str, chat_session_id: str, update_data: dict):
+    """Partial update of a chat session."""
+    user_ref = db.collection('users').document(uid)
+    session_ref = user_ref.collection('chat_sessions').document(chat_session_id)
+    session_ref.update(update_data)
+
+
+@set_data_protection_level(data_arg_name='message_data')
+@prepare_for_write(data_arg_name='message_data', prepare_func=_prepare_data_for_write)
+def save_message(uid: str, message_data: dict):
+    """Save a message directly by document ID (for desktop CRUD)."""
+    user_ref = db.collection('users').document(uid)
+    user_ref.collection('messages').document(message_data['id']).set(message_data)
+    return message_data
+
+
+def delete_chat_session_messages(uid: str, chat_session_id: str):
+    """Delete all messages belonging to a chat session."""
+    user_ref = db.collection('users').document(uid)
+    messages_ref = user_ref.collection('messages').where(filter=FieldFilter('chat_session_id', '==', chat_session_id))
+    batch = db.batch()
+    count = 0
+    for doc in messages_ref.stream():
+        batch.delete(doc.reference)
+        count += 1
+        if count % 400 == 0:
+            batch.commit()
+            batch = db.batch()
+    if count % 400 != 0:
+        batch.commit()
+    logger.info(f"Deleted {count} messages for session {chat_session_id}")
+
+
+def add_message_to_chat_session(uid: str, chat_session_id: str, message_id: str, preview: str = None):
     user_ref = db.collection('users').document(uid)
     session_ref = user_ref.collection('chat_sessions').document(chat_session_id)
-    session_ref.update({"message_ids": firestore.ArrayUnion([message_id])})
+    update_data = {
+        "message_ids": firestore.ArrayUnion([message_id]),
+        "updated_at": datetime.now(timezone.utc),
+        "message_count": firestore.Increment(1),
+    }
+    if preview:
+        update_data["preview"] = preview[:200]
+    session_ref.update(update_data)
 
 
 def add_files_to_chat_session(uid: str, chat_session_id: str, file_ids: List[str]):

backend/database/conversations.py

@@ -220,6 +220,30 @@ def get_conversations(
     return conversations
 
 
+def count_conversations(uid: str, statuses: Optional[List[str]] = None) -> int:
+    """Count conversations matching status filters without fetching full documents."""
+    if statuses is None:
+        statuses = []
+    conversations_ref = db.collection('users').document(uid).collection(conversations_collection)
+    conversations_ref = conversations_ref.where(filter=FieldFilter('discarded', '==', False))
+    if statuses:
+        conversations_ref = conversations_ref.where(filter=FieldFilter('status', 'in', statuses))
+    count_query = conversations_ref.count()
+    results = count_query.get()
+    return results[0][0].value
+
+
+def stream_conversations(uid: str, statuses: Optional[List[str]] = None):
+    """Yield conversation docs as a stream for counting without loading all into memory."""
+    if statuses is None:
+        statuses = []
+    conversations_ref = db.collection('users').document(uid).collection(conversations_collection)
+    conversations_ref = conversations_ref.where(filter=FieldFilter('discarded', '==', False))
+    if statuses:
+        conversations_ref = conversations_ref.where(filter=FieldFilter('status', 'in', statuses))
+    yield from conversations_ref.stream()
+
+
 @prepare_for_read(decrypt_func=_prepare_conversation_for_read)
 def get_conversations_without_photos(
     uid: str,

backend/database/focus_sessions.py

@@ -0,0 +1,75 @@
+import logging
+import uuid
+from datetime import datetime, timedelta, timezone
+from typing import List, Dict, Any, Optional
+
+from google.cloud import firestore
+
+from ._client import db
+
+logger = logging.getLogger(__name__)
+
+USERS_COLLECTION = 'users'
+FOCUS_SESSIONS_SUBCOLLECTION = 'focus_sessions'
+
+
+def _collection_ref(uid: str):
+    return db.collection(USERS_COLLECTION).document(uid).collection(FOCUS_SESSIONS_SUBCOLLECTION)
+
+
+def create_focus_session(uid: str, data: Dict[str, Any]) -> Dict[str, Any]:
+    """Create a new focus session document. Returns the created document with id."""
+    session_id = str(uuid.uuid4())
+    now = datetime.now(timezone.utc)
+
+    doc_data = {
+        'status': data['status'],
+        'app_or_site': data['app_or_site'],
+        'description': data['description'],
+        'created_at': now,
+    }
+    if data.get('message') is not None:
+        doc_data['message'] = data['message']
+    if data.get('duration_seconds') is not None:
+        doc_data['duration_seconds'] = data['duration_seconds']
+
+    _collection_ref(uid).document(session_id).set(doc_data)
+
+    doc_data['id'] = session_id
+    return doc_data
+
+
+def get_focus_sessions(
+    uid: str,
+    limit: int = 100,
+    offset: int = 0,
+    date: Optional[str] = None,
+) -> List[Dict[str, Any]]:
+    """Query focus sessions, ordered by created_at DESC. Optional date filter (YYYY-MM-DD)."""
+    query = _collection_ref(uid).order_by('created_at', direction=firestore.Query.DESCENDING)
+
+    if date:
+        day_start = datetime.strptime(date, '%Y-%m-%d').replace(tzinfo=timezone.utc)
+        next_day_start = day_start + timedelta(days=1)
+        query = query.where(filter=firestore.FieldFilter('created_at', '>=', day_start))
+        query = query.where(filter=firestore.FieldFilter('created_at', '<', next_day_start))
+
+    query = query.offset(offset).limit(limit)
+
+    results = []
+    for doc in query.stream():
+        data = doc.to_dict()
+        data['id'] = doc.id
+        results.append(data)
+    return results
+
+
+def delete_focus_session(uid: str, session_id: str) -> bool:
+    """Delete a focus session document. Returns True on success."""
+    _collection_ref(uid).document(session_id).delete()
+    return True
+
+
+def get_focus_sessions_for_stats(uid: str, date: str) -> List[Dict[str, Any]]:
+    """Get up to 1000 sessions for a date, for stats computation."""
+    return get_focus_sessions(uid, limit=1000, offset=0, date=date)