Create utils/collections.py

utils/collections.py

@@ -0,0 +1,296 @@
+# utils/collections.py
+
+import sqlite3
+from typing import List, Dict, Optional
+import streamlit as st
+from datetime import datetime
+from threading import Lock
+
+# Create a lock for database operations
+db_lock = Lock()
+
+def create_collection_tables(conn: sqlite3.Connection) -> None:
+    """
+    Create the necessary tables for the collections system.
+    """
+    with db_lock:
+        try:
+            # Collections table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS collections (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    name TEXT NOT NULL UNIQUE,
+                    description TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            ''')
+            
+            # Document-Collections relationship table
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS document_collections (
+                    document_id INTEGER,
+                    collection_id INTEGER,
+                    added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    PRIMARY KEY (document_id, collection_id),
+                    FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE,
+                    FOREIGN KEY (collection_id) REFERENCES collections (id) ON DELETE CASCADE
+                )
+            ''')
+            
+            # Index for faster lookups
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_collection_name ON collections(name)')
+            conn.execute('CREATE INDEX IF NOT EXISTS idx_doc_collection ON document_collections(collection_id)')
+            
+            conn.commit()
+            
+        except sqlite3.Error as e:
+            st.error(f"Database error: {e}")
+            raise
+
+def create_collection(conn: sqlite3.Connection, name: str, description: str = "") -> Optional[int]:
+    """
+    Create a new collection.
+    Returns the ID of the new collection or None if creation fails.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT INTO collections (name, description)
+                VALUES (?, ?)
+            ''', (name, description))
+            conn.commit()
+            return cursor.lastrowid
+        except sqlite3.IntegrityError:
+            st.error(f"Collection '{name}' already exists!")
+            return None
+        except sqlite3.Error as e:
+            st.error(f"Error creating collection: {e}")
+            return None
+
+def get_collections(conn: sqlite3.Connection) -> List[Dict]:
+    """
+    Get all collections with their document counts.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    c.id,
+                    c.name,
+                    c.description,
+                    c.created_at,
+                    COUNT(dc.document_id) as doc_count
+                FROM collections c
+                LEFT JOIN document_collections dc ON c.id = dc.collection_id
+                GROUP BY c.id
+                ORDER BY c.name
+            ''')
+            
+            collections = []
+            for row in cursor.fetchall():
+                collections.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'description': row[2],
+                    'created_at': row[3],
+                    'doc_count': row[4]
+                })
+            return collections
+            
+        except sqlite3.Error as e:
+            st.error(f"Error retrieving collections: {e}")
+            return []
+
+def get_collection_documents(conn: sqlite3.Connection, collection_id: int) -> List[Dict]:
+    """
+    Get all documents in a specific collection.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    d.id,
+                    d.name,
+                    d.upload_date,
+                    d.content
+                FROM documents d
+                JOIN document_collections dc ON d.id = dc.document_id
+                WHERE dc.collection_id = ?
+                ORDER BY d.upload_date DESC
+            ''', (collection_id,))
+            
+            documents = []
+            for row in cursor.fetchall():
+                documents.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'upload_date': row[2],
+                    'content': row[3]
+                })
+            return documents
+            
+        except sqlite3.Error as e:
+            st.error(f"Error retrieving collection documents: {e}")
+            return []
+
+def add_document_to_collection(
+    conn: sqlite3.Connection,
+    document_id: int,
+    collection_id: int
+) -> bool:
+    """
+    Add a document to a collection.
+    Returns True if successful, False otherwise.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                INSERT OR IGNORE INTO document_collections (document_id, collection_id)
+                VALUES (?, ?)
+            ''', (document_id, collection_id))
+            conn.commit()
+            return True
+            
+        except sqlite3.Error as e:
+            st.error(f"Error adding document to collection: {e}")
+            return False
+
+def remove_document_from_collection(
+    conn: sqlite3.Connection,
+    document_id: int,
+    collection_id: int
+) -> bool:
+    """
+    Remove a document from a collection.
+    Returns True if successful, False otherwise.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                DELETE FROM document_collections
+                WHERE document_id = ? AND collection_id = ?
+            ''', (document_id, collection_id))
+            conn.commit()
+            return True
+            
+        except sqlite3.Error as e:
+            st.error(f"Error removing document from collection: {e}")
+            return False
+
+def delete_collection(conn: sqlite3.Connection, collection_id: int) -> bool:
+    """
+    Delete a collection and all its document associations.
+    Returns True if successful, False otherwise.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            # Delete the collection (cascade will handle document_collections)
+            cursor.execute('DELETE FROM collections WHERE id = ?', (collection_id,))
+            conn.commit()
+            return True
+            
+        except sqlite3.Error as e:
+            st.error(f"Error deleting collection: {e}")
+            return False
+
+def update_collection(
+    conn: sqlite3.Connection,
+    collection_id: int,
+    name: Optional[str] = None,
+    description: Optional[str] = None
+) -> bool:
+    """
+    Update a collection's details.
+    Returns True if successful, False otherwise.
+    """
+    with db_lock:
+        try:
+            updates = []
+            params = []
+            
+            if name is not None:
+                updates.append("name = ?")
+                params.append(name)
+            if description is not None:
+                updates.append("description = ?")
+                params.append(description)
+                
+            if not updates:
+                return True
+                
+            updates.append("updated_at = CURRENT_TIMESTAMP")
+            params.append(collection_id)
+            
+            query = f'''
+                UPDATE collections 
+                SET {', '.join(updates)}
+                WHERE id = ?
+            '''
+            
+            cursor = conn.cursor()
+            cursor.execute(query, params)
+            conn.commit()
+            return True
+            
+        except sqlite3.Error as e:
+            st.error(f"Error updating collection: {e}")
+            return False
+
+def get_document_collections(conn: sqlite3.Connection, document_id: int) -> List[Dict]:
+    """
+    Get all collections that a document belongs to.
+    """
+    with db_lock:
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+                SELECT 
+                    c.id,
+                    c.name,
+                    c.description,
+                    dc.added_at
+                FROM collections c
+                JOIN document_collections dc ON c.id = dc.collection_id
+                WHERE dc.document_id = ?
+                ORDER BY c.name
+            ''', (document_id,))
+            
+            collections = []
+            for row in cursor.fetchall():
+                collections.append({
+                    'id': row[0],
+                    'name': row[1],
+                    'description': row[2],
+                    'added_at': row[3]
+                })
+            return collections
+            
+        except sqlite3.Error as e:
+            st.error(f"Error retrieving document collections: {e}")
+            return []
+
+# Example usage in your main app:
+"""
+# Initialize tables
+conn = create_connection('your_database.db')
+create_collection_tables(conn)
+
+# Create a new collection
+collection_id = create_collection(conn, "RFP 2024", "RFP documents for 2024 projects")
+
+# Add a document to the collection
+add_document_to_collection(conn, document_id, collection_id)
+
+# Get all collections
+collections = get_collections(conn)
+
+# Get documents in a collection
+documents = get_collection_documents(conn, collection_id)
+"""