341 lines
13 KiB
Python
341 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
SQLite database for neopig metadata.
|
|
|
|
Stores:
|
|
- Crawl jobs (target, keywords, timestamps, stats)
|
|
- Media records (md5_hash, source URLs, metadata)
|
|
- Analysis results (Qwen 3 VL outputs)
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from typing import List, Dict, Any, Optional, Set
|
|
|
|
import aiosqlite
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Database:
|
|
"""
|
|
Async SQLite database for neopig metadata.
|
|
"""
|
|
|
|
def __init__(self, db_path: str = "neopig.db"):
|
|
self.db_path = db_path
|
|
self._initialized = False
|
|
|
|
async def init(self) -> None:
|
|
"""Initialize database schema."""
|
|
if self._initialized:
|
|
return
|
|
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
# Crawl jobs table
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS crawl_jobs (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
target_uri TEXT NOT NULL,
|
|
keywords TEXT, -- JSON array
|
|
mode TEXT DEFAULT 'images',
|
|
status TEXT DEFAULT 'running',
|
|
started_at TEXT NOT NULL,
|
|
completed_at TEXT,
|
|
stats TEXT -- JSON object
|
|
)
|
|
""")
|
|
|
|
# Media records table (main deduped storage)
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS media (
|
|
md5_hash TEXT PRIMARY KEY,
|
|
media_type TEXT, -- 'image', 'video', 'audio'
|
|
mime_type TEXT,
|
|
file_size INTEGER,
|
|
keywords TEXT, -- JSON array
|
|
alt_text TEXT,
|
|
title TEXT,
|
|
first_seen_at TEXT NOT NULL,
|
|
analysis_status TEXT DEFAULT 'pending', -- 'pending', 'analyzed', 'invalid', 'error'
|
|
analysis_result TEXT -- JSON from Qwen 3 VL
|
|
)
|
|
""")
|
|
|
|
# Media sources table (tracks all contexts where media was found)
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS media_sources (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
md5_hash TEXT NOT NULL,
|
|
media_uri TEXT NOT NULL,
|
|
page_uri TEXT,
|
|
page_title TEXT,
|
|
page_description TEXT,
|
|
page_keywords TEXT,
|
|
page_content TEXT,
|
|
alt_text TEXT,
|
|
link_text TEXT,
|
|
detail_page_uri TEXT,
|
|
detail_title TEXT,
|
|
detail_content TEXT,
|
|
searchable_text TEXT, -- Combined metadata for full-text search
|
|
crawl_job_id INTEGER,
|
|
discovered_at TEXT NOT NULL,
|
|
FOREIGN KEY (md5_hash) REFERENCES media(md5_hash),
|
|
FOREIGN KEY (crawl_job_id) REFERENCES crawl_jobs(id),
|
|
UNIQUE(md5_hash, media_uri, page_uri)
|
|
)
|
|
""")
|
|
|
|
# Indexes
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_media_type ON media(media_type)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_media_analysis ON media(analysis_status)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_hash ON media_sources(md5_hash)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_job ON media_sources(crawl_job_id)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_media_uri ON media_sources(media_uri)")
|
|
await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_page_uri ON media_sources(page_uri)")
|
|
|
|
await db.commit()
|
|
|
|
self._initialized = True
|
|
logger.info(f"Database initialized: {self.db_path}")
|
|
|
|
async def create_crawl_job(
|
|
self,
|
|
target_uri: str,
|
|
keywords: List[str] = None,
|
|
mode: str = "images"
|
|
) -> int:
|
|
"""Create a new crawl job and return its ID."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
cursor = await db.execute(
|
|
"""
|
|
INSERT INTO crawl_jobs (target_uri, keywords, mode, started_at)
|
|
VALUES (?, ?, ?, ?)
|
|
""",
|
|
(
|
|
target_uri,
|
|
json.dumps(keywords or []),
|
|
mode,
|
|
datetime.now(timezone.utc).isoformat()
|
|
)
|
|
)
|
|
await db.commit()
|
|
return cursor.lastrowid
|
|
|
|
async def complete_crawl_job(self, job_id: int, stats: Dict[str, Any]) -> None:
|
|
"""Mark a crawl job as complete."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
UPDATE crawl_jobs
|
|
SET status = 'completed', completed_at = ?, stats = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
datetime.now(timezone.utc).isoformat(),
|
|
json.dumps(stats),
|
|
job_id
|
|
)
|
|
)
|
|
await db.commit()
|
|
|
|
async def create_media_record(
|
|
self,
|
|
md5_hash: str,
|
|
media_uri: str,
|
|
page_uri: str,
|
|
crawl_job_id: int,
|
|
media_type: str = "image",
|
|
mime_type: str = "",
|
|
file_size: int = 0,
|
|
page_title: str = "",
|
|
page_description: str = "",
|
|
page_keywords: str = "",
|
|
page_content: str = "",
|
|
alt_text: str = "",
|
|
link_text: str = "",
|
|
detail_page_uri: str = "",
|
|
detail_title: str = "",
|
|
detail_content: str = "",
|
|
searchable_text: str = "",
|
|
) -> None:
|
|
"""Create a new media record and add source context.
|
|
|
|
Skeleton key approach: stores both embedding context (page_title, page_content from listing)
|
|
and detail context (detail_title, detail_content from detail page) for maximum searchability.
|
|
|
|
For blogs: page_content contains the post text so images are searchable by post content.
|
|
For galleries: detail_content contains the detail page text for richer metadata.
|
|
"""
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
# Insert or ignore media record (just the hash and basic info)
|
|
await db.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO media
|
|
(md5_hash, media_type, mime_type, file_size, first_seen_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(md5_hash, media_type, mime_type, file_size, now)
|
|
)
|
|
|
|
# Always add source with full context (unique on media_uri + page_uri)
|
|
await db.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO media_sources
|
|
(md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, discovered_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, now)
|
|
)
|
|
|
|
await db.commit()
|
|
|
|
async def add_media_source(
|
|
self,
|
|
md5_hash: str,
|
|
media_uri: str,
|
|
page_uri: str,
|
|
page_title: str = "",
|
|
page_description: str = "",
|
|
page_keywords: str = "",
|
|
page_content: str = "",
|
|
alt_text: str = "",
|
|
link_text: str = "",
|
|
detail_page_uri: str = "",
|
|
detail_title: str = "",
|
|
detail_content: str = "",
|
|
searchable_text: str = "",
|
|
crawl_job_id: int = None
|
|
) -> None:
|
|
"""Add another source context for an existing media hash.
|
|
|
|
Skeleton key approach: stores both embedding context (page_title, page_content from listing)
|
|
and detail context (detail_title, detail_content from detail page) for maximum searchability.
|
|
|
|
For blogs: page_content contains the post text so images are searchable by post content.
|
|
For galleries: detail_content contains the detail page text for richer metadata.
|
|
"""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO media_sources
|
|
(md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, discovered_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, datetime.now(timezone.utc).isoformat())
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_pending_analysis(self, limit: int = 100) -> List[Dict[str, Any]]:
|
|
"""Get media items pending analysis."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute(
|
|
"""
|
|
SELECT md5_hash, media_type, mime_type, keywords, alt_text, title
|
|
FROM media
|
|
WHERE analysis_status = 'pending'
|
|
LIMIT ?
|
|
""",
|
|
(limit,)
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def update_analysis(
|
|
self,
|
|
md5_hash: str,
|
|
status: str,
|
|
result: Dict[str, Any] = None
|
|
) -> None:
|
|
"""Update analysis status and result for a media item."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
await db.execute(
|
|
"""
|
|
UPDATE media
|
|
SET analysis_status = ?, analysis_result = ?
|
|
WHERE md5_hash = ?
|
|
""",
|
|
(status, json.dumps(result) if result else None, md5_hash)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_media_by_keyword(self, keyword: str, limit: int = 100) -> List[Dict[str, Any]]:
|
|
"""Search media by keyword."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
# Search in keywords JSON array and alt_text/title
|
|
cursor = await db.execute(
|
|
"""
|
|
SELECT m.*, GROUP_CONCAT(ms.source_url) as source_urls
|
|
FROM media m
|
|
LEFT JOIN media_sources ms ON m.md5_hash = ms.md5_hash
|
|
WHERE m.keywords LIKE ?
|
|
OR m.alt_text LIKE ?
|
|
OR m.title LIKE ?
|
|
GROUP BY m.md5_hash
|
|
LIMIT ?
|
|
""",
|
|
(f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', limit)
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def check_media_uri_exists(self, media_uri: str, page_uri: str) -> Optional[str]:
|
|
"""
|
|
Check if a media URI from a specific page has already been crawled.
|
|
|
|
Returns:
|
|
The md5_hash if exists, None otherwise
|
|
"""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
cursor = await db.execute(
|
|
"SELECT md5_hash FROM media_sources WHERE media_uri = ? AND page_uri = ?",
|
|
(media_uri, page_uri)
|
|
)
|
|
row = await cursor.fetchone()
|
|
return row[0] if row else None
|
|
|
|
async def get_crawled_media_uris(self) -> Set[str]:
|
|
"""Get all media URIs that have been crawled."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
cursor = await db.execute(
|
|
"SELECT DISTINCT media_uri FROM media_sources"
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return {row[0] for row in rows}
|
|
|
|
async def get_stats(self) -> Dict[str, Any]:
|
|
"""Get database statistics."""
|
|
async with aiosqlite.connect(self.db_path) as db:
|
|
stats = {}
|
|
|
|
# Total media
|
|
cursor = await db.execute("SELECT COUNT(*) FROM media")
|
|
stats['total_media'] = (await cursor.fetchone())[0]
|
|
|
|
# By type
|
|
cursor = await db.execute(
|
|
"SELECT media_type, COUNT(*) FROM media GROUP BY media_type"
|
|
)
|
|
stats['by_type'] = {row[0]: row[1] for row in await cursor.fetchall()}
|
|
|
|
# By analysis status
|
|
cursor = await db.execute(
|
|
"SELECT analysis_status, COUNT(*) FROM media GROUP BY analysis_status"
|
|
)
|
|
stats['by_analysis'] = {row[0]: row[1] for row in await cursor.fetchall()}
|
|
|
|
# Total sources
|
|
cursor = await db.execute("SELECT COUNT(*) FROM media_sources")
|
|
stats['total_sources'] = (await cursor.fetchone())[0]
|
|
|
|
# Crawl jobs
|
|
cursor = await db.execute("SELECT COUNT(*) FROM crawl_jobs")
|
|
stats['total_jobs'] = (await cursor.fetchone())[0]
|
|
|
|
return stats
|