104 lines
3.1 KiB
Python
Executable File
104 lines
3.1 KiB
Python
Executable File
"""
|
|
SQLite database module for storing forum data.
|
|
|
|
This module provides functionality to store parsed forum data in SQLite database.
|
|
"""
|
|
|
|
import sqlite3
|
|
from typing import Optional
|
|
from ..Datamodel import Thread, Message, Image
|
|
from .table import Table
|
|
from .thread_table import ThreadTable
|
|
from .message_table import MessageTable
|
|
from .image_table import ImageTable
|
|
|
|
|
|
class Database:
|
|
"""
|
|
SQLite database for storing forum thread data.
|
|
|
|
Attributes:
|
|
db_path: Path to the SQLite database file
|
|
connection: SQLite database connection
|
|
thread_table: Thread table object
|
|
message_table: Message table object
|
|
image_table: Image table object
|
|
"""
|
|
|
|
def __init__(self, db_path: str):
|
|
"""
|
|
Initialize the database and create tables.
|
|
|
|
Args:
|
|
db_path: Path to the SQLite database file
|
|
"""
|
|
self.db_path = db_path
|
|
self.connection = None
|
|
self.thread_table = None
|
|
self.message_table = None
|
|
self.image_table = None
|
|
self._create_tables()
|
|
|
|
def _create_tables(self):
|
|
"""Create database tables if they don't exist."""
|
|
self.connection = sqlite3.connect(self.db_path)
|
|
|
|
# Create table objects
|
|
self.thread_table = ThreadTable(self.connection)
|
|
self.message_table = MessageTable(self.connection)
|
|
self.image_table = ImageTable(self.connection)
|
|
|
|
# Create tables
|
|
self.thread_table.create()
|
|
self.message_table.create()
|
|
self.image_table.create()
|
|
|
|
def save_thread_recursive(self, thread: Thread) -> int:
|
|
"""
|
|
Recursively save a thread, all its messages, and all images.
|
|
|
|
Args:
|
|
thread: Thread object to save
|
|
|
|
Returns:
|
|
Number of objects that were actually saved (not skipped)
|
|
"""
|
|
saved_count = 0
|
|
|
|
# Save thread
|
|
if self.thread_table.save(thread):
|
|
saved_count += 1
|
|
print(f"✓ Saved thread: {thread.title}")
|
|
else:
|
|
print(f"⊘ Thread already exists: {thread.title}")
|
|
|
|
# Save messages
|
|
for message in thread.messages:
|
|
if self.message_table.save(message):
|
|
saved_count += 1
|
|
print(f"✓ Saved message: {message.id}")
|
|
else:
|
|
print(f"⊘ Message already exists: {message.id}")
|
|
|
|
# Save images for this message
|
|
for image in message.images:
|
|
if self.image_table.save(image):
|
|
saved_count += 1
|
|
print(f"✓ Saved image: {image.name}")
|
|
else:
|
|
print(f"⊘ Image already exists: {image.name}")
|
|
|
|
return saved_count
|
|
|
|
def close(self):
|
|
"""Close the database connection."""
|
|
if self.connection:
|
|
self.connection.close()
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit."""
|
|
self.close() |