284 lines
9.4 KiB
Python
284 lines
9.4 KiB
Python
"""
|
||
从 192.168.1.100:4500 API 拉取指定数量相册作为测试数据（多线程并行版）
|
||
导入到 Rust imager 项目的 SQLite 数据库和 albums 目录
|
||
"""
|
||
import requests
|
||
import sqlite3
|
||
import os
|
||
import sys
|
||
import time
|
||
import uuid
|
||
import threading
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from datetime import datetime
|
||
|
||
API_BASE = "http://192.168.1.100:4500"
|
||
DB_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "imager.db")
|
||
ALBUMS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "albums")
|
||
|
||
# 全局 DB 写锁
|
||
db_lock = threading.Lock()
|
||
# 进度计数器
|
||
progress_lock = threading.Lock()
|
||
completed_count = [0]
|
||
total_count = [0]
|
||
|
||
# 线程本地 Session（复用连接）
|
||
thread_local = threading.local()
|
||
|
||
|
||
def get_session():
|
||
if not hasattr(thread_local, "session"):
|
||
thread_local.session = requests.Session()
|
||
thread_local.session.timeout = 60
|
||
return thread_local.session
|
||
|
||
|
||
def fetch_albums(limit=100):
|
||
"""获取相册列表"""
|
||
url = f"{API_BASE}/api/albums"
|
||
print(f"[1/4] 获取相册列表: {url}")
|
||
resp = requests.get(url, timeout=30)
|
||
data = resp.json()
|
||
|
||
if isinstance(data, dict) and "albums" in data:
|
||
albums = data["albums"]
|
||
elif isinstance(data, list):
|
||
albums = data
|
||
else:
|
||
print(f"未知响应格式: {type(data)}")
|
||
return []
|
||
|
||
albums = albums[:limit]
|
||
print(f" 获取到 {len(albums)} 个相册")
|
||
return albums
|
||
|
||
|
||
def fetch_album_images(album_id):
|
||
"""获取相册的图片列表"""
|
||
url = f"{API_BASE}/api/album/{album_id}"
|
||
try:
|
||
resp = requests.get(url, timeout=30)
|
||
data = resp.json()
|
||
if isinstance(data, dict) and "images" in data:
|
||
return data["images"]
|
||
elif isinstance(data, list):
|
||
return data
|
||
except Exception as e:
|
||
print(f" 获取图片列表失败 {album_id}: {e}")
|
||
return []
|
||
|
||
|
||
def download_single_image(img):
|
||
"""下载单张图片（线程安全）"""
|
||
img_id = img.get("id", str(uuid.uuid4()))
|
||
mime = img.get("mime_type", "image/jpeg")
|
||
ext = "jpg"
|
||
if mime and "/" in mime:
|
||
ext = mime.split("/")[-1]
|
||
|
||
album_id = str(img.get("_album_id", ""))
|
||
filename = f"{img_id}.{ext}"
|
||
db_filename = f"{album_id}/{filename}" # 数据库路径：{album_id}/{img_id}.ext
|
||
save_path = os.path.join(ALBUMS_DIR, album_id, filename)
|
||
|
||
# 跳过已下载
|
||
if os.path.exists(save_path) and os.path.getsize(save_path) > 0:
|
||
file_size = os.path.getsize(save_path)
|
||
width, height = get_image_dimensions(save_path)
|
||
return (img_id, db_filename, file_size, width, height, mime)
|
||
|
||
url = f"{API_BASE}/api/image/{img_id}"
|
||
try:
|
||
session = get_session()
|
||
resp = session.get(url, timeout=60)
|
||
if resp.status_code == 200:
|
||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||
with open(save_path, "wb") as f:
|
||
f.write(resp.content)
|
||
file_size = len(resp.content)
|
||
width, height = get_image_dimensions(save_path)
|
||
return (img_id, db_filename, file_size, width, height, resp.headers.get("content-type", mime))
|
||
except Exception as e:
|
||
pass
|
||
return None
|
||
|
||
|
||
def get_image_dimensions(path):
|
||
"""获取图片尺寸（二进制头解析）"""
|
||
try:
|
||
with open(path, "rb") as f:
|
||
header = f.read(32)
|
||
if header[:2] == b"\xff\xd8":
|
||
f.seek(2)
|
||
while True:
|
||
chunk = f.read(4)
|
||
if len(chunk) < 4:
|
||
break
|
||
if chunk[0] != 0xFF:
|
||
break
|
||
seg_type = chunk[1]
|
||
if seg_type == 0xDA:
|
||
break
|
||
seg_len = int.from_bytes(chunk[2:4], "big")
|
||
if seg_type in (0xC0, 0xC1, 0xC2):
|
||
data = f.read(seg_len - 2)
|
||
if len(data) >= 5:
|
||
return int.from_bytes(data[1:3], "big"), int.from_bytes(data[3:5], "big")
|
||
f.seek(seg_len - 2, 1)
|
||
return 0, 0
|
||
elif header[:8] == b"\x89PNG\r\n\x1a\n":
|
||
return int.from_bytes(header[16:20], "big"), int.from_bytes(header[20:24], "big")
|
||
elif header[:4] == b"RIFF" and header[8:12] == b"WEBP":
|
||
return int.from_bytes(header[24:26], "little"), int.from_bytes(header[26:28], "little")
|
||
elif header[:6] in (b"GIF87a", b"GIF89a"):
|
||
return int.from_bytes(header[6:8], "little"), int.from_bytes(header[8:10], "little")
|
||
return 0, 0
|
||
except Exception:
|
||
return 0, 0
|
||
|
||
|
||
def sync_album_to_db(conn, album, images_info):
|
||
"""将相册和图片写入数据库（调用方持有 db_lock）"""
|
||
cursor = conn.cursor()
|
||
now = datetime.utcnow().isoformat() + "Z"
|
||
album_id = album["id"]
|
||
title = album.get("title", "Untitled")
|
||
|
||
cursor.execute(
|
||
"""INSERT OR REPLACE INTO albums
|
||
(id, title, document, description, cover_image_id, image_count, created_at, updated_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
||
(
|
||
album_id, title, album_id,
|
||
album.get("description"),
|
||
images_info[0][0] if images_info else None,
|
||
len(images_info),
|
||
album.get("created_at", now),
|
||
now,
|
||
),
|
||
)
|
||
|
||
for img_id, filename, file_size, width, height, mime_type in images_info:
|
||
cursor.execute(
|
||
"""INSERT OR REPLACE INTO images
|
||
(id, album_id, filename, original_url, title, description, file_size, width, height, mime_type, created_at)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
||
(img_id, album_id, filename, f"{API_BASE}/api/image/{img_id}",
|
||
f"图片 {img_id[:8]}", None, file_size, width, height, mime_type or "image/jpeg", now),
|
||
)
|
||
|
||
conn.commit()
|
||
|
||
|
||
def process_album(album, conn, dl_workers):
|
||
"""处理单个相册：下载所有图片并写入数据库"""
|
||
aid = album["id"]
|
||
title = album.get("title", "Untitled")[:60]
|
||
|
||
# 获取图片列表
|
||
images = fetch_album_images(aid)
|
||
if not images:
|
||
with progress_lock:
|
||
completed_count[0] += 1
|
||
print(f"[{completed_count[0]}/{total_count[0]}] {title} — 无图片，跳过")
|
||
return 0
|
||
|
||
# 为每个图片打上 album_id 标记
|
||
for img in images:
|
||
img["_album_id"] = aid
|
||
|
||
# 并行下载
|
||
downloaded = []
|
||
with ThreadPoolExecutor(max_workers=dl_workers) as pool:
|
||
futures = {pool.submit(download_single_image, img): img for img in images}
|
||
for future in as_completed(futures):
|
||
result = future.result()
|
||
if result:
|
||
downloaded.append(result)
|
||
|
||
if not downloaded:
|
||
with progress_lock:
|
||
completed_count[0] += 1
|
||
print(f"[{completed_count[0]}/{total_count[0]}] {title} — 0 张下载成功")
|
||
return 0
|
||
|
||
# 写入数据库
|
||
with db_lock:
|
||
sync_album_to_db(conn, album, downloaded)
|
||
|
||
n = len(downloaded)
|
||
with progress_lock:
|
||
completed_count[0] += 1
|
||
print(f"[{completed_count[0]}/{total_count[0]}] {title} — {n} 张图片")
|
||
|
||
return n
|
||
|
||
|
||
def main():
|
||
total = int(sys.argv[1]) if len(sys.argv) > 1 else 100
|
||
dl_workers = int(sys.argv[2]) if len(sys.argv) > 2 else 8
|
||
album_workers = int(sys.argv[3]) if len(sys.argv) > 3 else 4
|
||
|
||
print(f"=== 从 API 导入 {total} 个相册（并行下载: {dl_workers}线程/相册, 并行相册: {album_workers}）===\n")
|
||
|
||
os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
|
||
os.makedirs(ALBUMS_DIR, exist_ok=True)
|
||
|
||
conn = sqlite3.connect(DB_PATH, check_same_thread=False)
|
||
conn.execute("PRAGMA journal_mode=WAL")
|
||
conn.execute("PRAGMA foreign_keys=ON")
|
||
conn.execute("PRAGMA synchronous=NORMAL")
|
||
|
||
# 获取已有相册
|
||
existing = set()
|
||
for row in conn.execute("SELECT id FROM albums"):
|
||
existing.add(row[0])
|
||
print(f" 数据库中已有 {len(existing)} 个相册")
|
||
|
||
# 获取相册列表
|
||
albums = fetch_albums(total)
|
||
if not albums:
|
||
print("未获取到相册，退出")
|
||
return
|
||
|
||
new_albums = [a for a in albums if a["id"] not in existing]
|
||
skipped = len(albums) - len(new_albums)
|
||
if skipped > 0:
|
||
print(f" 跳过 {skipped} 个已存在的相册")
|
||
print(f" 需要导入 {len(new_albums)} 个新相册\n")
|
||
|
||
if not new_albums:
|
||
print("没有新相册需要导入")
|
||
conn.close()
|
||
return
|
||
|
||
total_count[0] = len(new_albums)
|
||
total_images = 0
|
||
t0 = time.time()
|
||
|
||
# 并行处理相册
|
||
with ThreadPoolExecutor(max_workers=album_workers) as pool:
|
||
futures = {pool.submit(process_album, album, conn, dl_workers): album for album in new_albums}
|
||
for future in as_completed(futures):
|
||
n = future.result()
|
||
total_images += n
|
||
|
||
elapsed = time.time() - t0
|
||
|
||
# 重建 image_count
|
||
print(f"\n[4/4] 更新相册图片计数...")
|
||
conn.execute("""UPDATE albums SET image_count = (SELECT COUNT(*) FROM images WHERE images.album_id = albums.id)""")
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
print(f"\n=== 导入完成 ===")
|
||
print(f" 成功导入: {len(new_albums)} 个相册")
|
||
print(f" 总图片数: {total_images} 张")
|
||
print(f" 耗时: {elapsed:.1f} 秒")
|
||
print(f" 数据库: {DB_PATH}")
|
||
print(f" 图片目录: {ALBUMS_DIR}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |