72 lines
2.6 KiB
Python
72 lines
2.6 KiB
Python
import os
|
|
import sqlite3
|
|
|
|
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'booru.db')
|
|
|
|
|
|
def get_conn():
|
|
return sqlite3.connect(DB_PATH)
|
|
|
|
|
|
def init_db():
|
|
with get_conn() as c:
|
|
c.execute("""CREATE TABLE IF NOT EXISTS images (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
post_id TEXT NOT NULL,
|
|
site TEXT NOT NULL,
|
|
filename TEXT NOT NULL UNIQUE,
|
|
tags TEXT NOT NULL DEFAULT '',
|
|
file_url TEXT NOT NULL,
|
|
post_url TEXT NOT NULL DEFAULT '',
|
|
preview_filename TEXT NOT NULL DEFAULT '',
|
|
downloaded_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
|
)""")
|
|
c.execute(
|
|
"CREATE UNIQUE INDEX IF NOT EXISTS idx_site_post ON images (site, post_id)"
|
|
)
|
|
# migrate existing DBs that predate the post_url column
|
|
try:
|
|
c.execute("ALTER TABLE images ADD COLUMN post_url TEXT NOT NULL DEFAULT ''")
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
# migrate existing DBs that predate the preview_filename column
|
|
try:
|
|
c.execute("ALTER TABLE images ADD COLUMN preview_filename TEXT NOT NULL DEFAULT ''")
|
|
except sqlite3.OperationalError:
|
|
pass
|
|
|
|
|
|
def image_exists(site, post_id):
|
|
with get_conn() as c:
|
|
return (
|
|
c.execute(
|
|
"SELECT 1 FROM images WHERE site=? AND post_id=?", (site, post_id)
|
|
).fetchone()
|
|
is not None
|
|
)
|
|
|
|
|
|
def insert_image(post_id, site, filename, tags, file_url, post_url, preview_filename=''):
|
|
with get_conn() as c:
|
|
c.execute(
|
|
"INSERT OR IGNORE INTO images (post_id, site, filename, tags, file_url, post_url, preview_filename) VALUES (?,?,?,?,?,?,?)",
|
|
(post_id, site, filename, tags, file_url, post_url, preview_filename),
|
|
)
|
|
|
|
|
|
def search_images(tag_query):
|
|
terms = tag_query.split() if tag_query.strip() else []
|
|
with get_conn() as c:
|
|
if not terms:
|
|
rows = c.execute(
|
|
"SELECT filename, preview_filename, post_url, tags FROM images ORDER BY id"
|
|
).fetchall()
|
|
else:
|
|
where = ' AND '.join(['tags LIKE ?'] * len(terms))
|
|
params = [f'%{t}%' for t in terms]
|
|
rows = c.execute(
|
|
f"SELECT filename, preview_filename, post_url, tags FROM images WHERE {where} ORDER BY id",
|
|
params,
|
|
).fetchall()
|
|
return [{'filename': r[0], 'preview_filename': r[1], 'post_url': r[2], 'tags': r[3]} for r in rows]
|