131 lines
3.9 KiB
Python
Executable File
131 lines
3.9 KiB
Python
Executable File
"""
|
|
Message table class for database operations.
|
|
|
|
This module provides the MessageTable class that handles message-specific
|
|
database operations for storing and retrieving Message objects.
|
|
"""
|
|
|
|
import sqlite3
|
|
import json
|
|
from typing import Optional
|
|
from datetime import datetime
|
|
from ..Datamodel.Message import Message
|
|
|
|
|
|
class MessageTable:
|
|
"""
|
|
Table class for managing message data in the database.
|
|
|
|
Attributes:
|
|
connection: SQLite database connection
|
|
cursor: Database cursor for executing queries
|
|
"""
|
|
|
|
def __init__(self, connection: sqlite3.Connection):
|
|
"""
|
|
Initialize the message table with a database connection.
|
|
|
|
Args:
|
|
connection: SQLite database connection
|
|
"""
|
|
self.connection = connection
|
|
self.cursor = self.connection.cursor()
|
|
|
|
def create(self):
|
|
"""Create the messages table if it doesn't exist."""
|
|
self.cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id TEXT PRIMARY KEY,
|
|
thread_id TEXT NOT NULL,
|
|
author TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
text_content TEXT NOT NULL,
|
|
timestamp TEXT NOT NULL,
|
|
reply_links TEXT,
|
|
ref_links TEXT,
|
|
FOREIGN KEY (thread_id) REFERENCES threads (id) ON DELETE CASCADE
|
|
)
|
|
''')
|
|
self.connection.commit()
|
|
|
|
def exists(self, message_id: str) -> bool:
|
|
"""
|
|
Check if a message exists in the database.
|
|
|
|
Args:
|
|
message_id: Message ID to check
|
|
|
|
Returns:
|
|
True if message exists, False otherwise
|
|
"""
|
|
self.cursor.execute('SELECT 1 FROM messages WHERE id = ?', (message_id,))
|
|
return self.cursor.fetchone() is not None
|
|
|
|
def save(self, message: Message) -> bool:
|
|
"""
|
|
Save a message to the database.
|
|
|
|
Args:
|
|
message: Message object to save
|
|
thread_id: Thread ID that this message belongs to
|
|
|
|
Returns:
|
|
True if message was saved, False if it already exists
|
|
"""
|
|
if self.exists(message.id):
|
|
return False
|
|
|
|
self.cursor.execute('''
|
|
INSERT INTO messages (id, thread_id, author, content, text_content, timestamp, reply_links, ref_links)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (
|
|
message.id,
|
|
message.ref_id,
|
|
message.author,
|
|
message.content,
|
|
message.text_content,
|
|
message.timestamp.isoformat(),
|
|
json.dumps(message.reply_links) if message.reply_links else None,
|
|
json.dumps(message.ref_links) if message.ref_links else None
|
|
))
|
|
|
|
self.connection.commit()
|
|
return True
|
|
|
|
def load(self, message_id: str) -> Optional[Message]:
|
|
"""
|
|
Load a message from the database by ID.
|
|
|
|
Args:
|
|
message_id: Message ID to load
|
|
|
|
Returns:
|
|
Message object if found, None otherwise
|
|
"""
|
|
self.cursor.execute('''
|
|
SELECT id, thread_id, author, content, text_content, timestamp, reply_links, ref_links
|
|
FROM messages WHERE id = ?
|
|
''', (message_id,))
|
|
|
|
row = self.cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
|
|
return Message(
|
|
id=row[0],
|
|
ref_id=row[1],
|
|
author=row[2],
|
|
content=row[3],
|
|
text_content=row[4],
|
|
timestamp=datetime.fromisoformat(row[5]),
|
|
reply_links=json.loads(row[6]) if row[6] else list(),
|
|
ref_links=json.loads(row[7]) if row[7] else list()
|
|
)
|
|
|
|
def commit(self):
|
|
"""Commit pending changes to the database."""
|
|
self.connection.commit()
|
|
|
|
def close(self):
|
|
"""Close the cursor."""
|
|
self.cursor.close() |