264 lines
8.6 KiB
Python
264 lines
8.6 KiB
Python
"""
|
|
The Update button: pull Rebrickable's parts catalogue into the local database.
|
|
|
|
Rebrickable publishes its whole catalogue as gzipped CSV dumps, regenerated
|
|
monthly, free to use with a credit. Those are what this uses:
|
|
|
|
https://cdn.rebrickable.com/media/downloads/part_categories.csv.gz
|
|
https://cdn.rebrickable.com/media/downloads/parts.csv.gz
|
|
|
|
Two files and about 60,000 parts, rather than thousands of paged API calls.
|
|
No API key is needed for the dumps - the key is only used later, when a
|
|
catalogue row becomes a label and we want the authoritative part image.
|
|
|
|
The sync runs on a background thread and writes its progress to a CatalogSync
|
|
row, so the browser polls instead of holding a request open for minutes.
|
|
|
|
Created by: Jess Rogerson (yelling commands at Claude.AI)
|
|
"""
|
|
|
|
import csv
|
|
import gzip
|
|
import io
|
|
import logging
|
|
import re
|
|
import threading
|
|
from datetime import datetime
|
|
|
|
import requests
|
|
|
|
from app import db
|
|
from app.models.part_catalog_item import (CatalogSync, PartCatalogCategory,
|
|
PartCatalogItem)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
BASE_URL = 'https://cdn.rebrickable.com/media/downloads'
|
|
CATEGORIES_FILE = 'part_categories.csv.gz'
|
|
PARTS_FILE = 'parts.csv.gz'
|
|
|
|
DOWNLOAD_TIMEOUT = 120
|
|
BATCH_SIZE = 2000
|
|
|
|
# Only one sync at a time, whatever the UI does.
|
|
_sync_lock = threading.Lock()
|
|
_sync_thread = None
|
|
|
|
DIMENSION_RE = re.compile(
|
|
r'(\d+(?:\.\d+)?\s*x\s*\d+(?:\.\d+)?(?:\s*x\s*\d+(?:\.\d+)?)?)',
|
|
re.IGNORECASE)
|
|
|
|
|
|
def guess_dimensions(name):
|
|
"""Pull '2 x 4' out of 'Brick 2 x 4'. Empty is a fine answer."""
|
|
match = DIMENSION_RE.search(name or '')
|
|
if not match:
|
|
return None
|
|
return re.sub(r'\s*x\s*', ' x ', match.group(1).strip(), flags=re.I)
|
|
|
|
|
|
def _download_csv(filename):
|
|
"""Stream a gzipped CSV off the CDN and return it as a list of dicts."""
|
|
url = f'{BASE_URL}/{filename}'
|
|
logger.info('Downloading %s', url)
|
|
response = requests.get(url, timeout=DOWNLOAD_TIMEOUT, stream=True)
|
|
response.raise_for_status()
|
|
raw = gzip.decompress(response.content)
|
|
text = io.StringIO(raw.decode('utf-8', errors='replace'))
|
|
return list(csv.DictReader(text))
|
|
|
|
|
|
def is_running():
|
|
latest = CatalogSync.latest()
|
|
return bool(latest and latest.is_running)
|
|
|
|
|
|
def start_sync(app, user_id=None):
|
|
"""
|
|
Kick off a sync on a background thread.
|
|
|
|
Returns the CatalogSync row, or the already-running one if a sync is
|
|
in flight. Never raises just because the network is down - the row
|
|
records the failure instead.
|
|
"""
|
|
global _sync_thread
|
|
|
|
with _sync_lock:
|
|
existing = CatalogSync.latest()
|
|
if existing and existing.is_running:
|
|
return existing
|
|
|
|
run = CatalogSync(status='running', stage='Starting', progress=0,
|
|
started_by_id=user_id)
|
|
db.session.add(run)
|
|
db.session.commit()
|
|
run_id = run.id
|
|
|
|
_sync_thread = threading.Thread(
|
|
target=_run_sync, args=(app, run_id), daemon=True,
|
|
name=f'catalog-sync-{run_id}')
|
|
_sync_thread.start()
|
|
return run
|
|
|
|
|
|
def _run_sync(app, run_id):
|
|
"""Thread body. Owns its own app context and session."""
|
|
with app.app_context():
|
|
run = db.session.get(CatalogSync, run_id)
|
|
if run is None:
|
|
return
|
|
|
|
def step(stage, progress):
|
|
run.stage = stage
|
|
run.progress = progress
|
|
db.session.commit()
|
|
logger.info('Catalog sync %s: %s (%s%%)', run_id, stage, progress)
|
|
|
|
try:
|
|
step('Downloading part categories', 5)
|
|
category_rows = _download_csv(CATEGORIES_FILE)
|
|
|
|
step('Saving part categories', 15)
|
|
run.categories_added = _sync_categories(category_rows)
|
|
db.session.commit()
|
|
|
|
step('Downloading parts catalogue (about 60,000 parts)', 25)
|
|
part_rows = _download_csv(PARTS_FILE)
|
|
|
|
step(f'Saving {len(part_rows):,} parts', 45)
|
|
added, updated = _sync_parts(part_rows, run, step)
|
|
|
|
run.parts_added = added
|
|
run.parts_updated = updated
|
|
run.status = 'done'
|
|
run.stage = 'Finished'
|
|
run.progress = 100
|
|
run.finished_at = datetime.utcnow()
|
|
run.message = (f'{added:,} new parts, {updated:,} updated, '
|
|
f'{run.categories_added:,} new categories.')
|
|
db.session.commit()
|
|
logger.info('Catalog sync %s finished: %s', run_id, run.message)
|
|
|
|
except Exception as exc:
|
|
db.session.rollback()
|
|
logger.exception('Catalog sync %s failed', run_id)
|
|
try:
|
|
run = db.session.get(CatalogSync, run_id)
|
|
run.status = 'failed'
|
|
run.stage = 'Failed'
|
|
run.finished_at = datetime.utcnow()
|
|
run.message = f'{type(exc).__name__}: {exc}'
|
|
db.session.commit()
|
|
except Exception:
|
|
db.session.rollback()
|
|
|
|
|
|
def _sync_categories(rows):
|
|
"""Upsert Rebrickable's part categories. Returns how many were new."""
|
|
existing = {c.id: c for c in PartCatalogCategory.query.all()}
|
|
added = 0
|
|
for row in rows:
|
|
try:
|
|
cat_id = int(row['id'])
|
|
except (KeyError, TypeError, ValueError):
|
|
continue
|
|
name = (row.get('name') or '').strip()[:200]
|
|
if not name:
|
|
continue
|
|
current = existing.get(cat_id)
|
|
if current is None:
|
|
db.session.add(PartCatalogCategory(id=cat_id, name=name))
|
|
added += 1
|
|
elif current.name != name:
|
|
current.name = name
|
|
db.session.commit()
|
|
return added
|
|
|
|
|
|
def _sync_parts(rows, run, step):
|
|
"""
|
|
Upsert the parts catalogue in batches.
|
|
|
|
Reads the existing rows into memory once rather than issuing 60,000
|
|
SELECTs, then uses bulk insert/update mappings. On SQLite this takes
|
|
seconds instead of many minutes.
|
|
"""
|
|
existing = {
|
|
item.part_num: item
|
|
for item in PartCatalogItem.query.all()
|
|
}
|
|
existing_snapshot = {
|
|
part_num: (item.id, item.name, item.category_id, item.part_material,
|
|
item.dimensions)
|
|
for part_num, item in existing.items()
|
|
}
|
|
del existing
|
|
|
|
now = datetime.utcnow()
|
|
to_insert = []
|
|
to_update = []
|
|
seen = set()
|
|
|
|
total = max(len(rows), 1)
|
|
for index, row in enumerate(rows):
|
|
part_num = (row.get('part_num') or '').strip()[:32]
|
|
if not part_num or part_num in seen:
|
|
continue
|
|
seen.add(part_num)
|
|
|
|
name = (row.get('name') or '').strip()[:300] or part_num
|
|
material = (row.get('part_material') or '').strip()[:40] or None
|
|
try:
|
|
category_id = int(row['part_cat_id'])
|
|
except (KeyError, TypeError, ValueError):
|
|
category_id = None
|
|
dimensions = guess_dimensions(name)
|
|
|
|
previous = existing_snapshot.get(part_num)
|
|
if previous is None:
|
|
to_insert.append({
|
|
'part_num': part_num, 'name': name,
|
|
'category_id': category_id, 'part_material': material,
|
|
'dimensions': dimensions, 'has_label': False,
|
|
'synced_at': now,
|
|
})
|
|
else:
|
|
row_id, old_name, old_cat, old_material, old_dims = previous
|
|
if (old_name, old_cat, old_material, old_dims) != (
|
|
name, category_id, material, dimensions):
|
|
to_update.append({
|
|
'id': row_id, 'name': name, 'category_id': category_id,
|
|
'part_material': material, 'dimensions': dimensions,
|
|
'synced_at': now,
|
|
})
|
|
|
|
if index % 10000 == 0:
|
|
step(f'Saving parts ({index:,} of {len(rows):,})',
|
|
45 + int(45 * index / total))
|
|
|
|
step(f'Writing {len(to_insert):,} new parts', 92)
|
|
for start in range(0, len(to_insert), BATCH_SIZE):
|
|
db.session.bulk_insert_mappings(
|
|
PartCatalogItem, to_insert[start:start + BATCH_SIZE])
|
|
db.session.commit()
|
|
|
|
step(f'Updating {len(to_update):,} changed parts', 97)
|
|
for start in range(0, len(to_update), BATCH_SIZE):
|
|
db.session.bulk_update_mappings(
|
|
PartCatalogItem, to_update[start:start + BATCH_SIZE])
|
|
db.session.commit()
|
|
|
|
return len(to_insert), len(to_update)
|
|
|
|
|
|
def catalog_stats():
|
|
"""Counts for the catalogue page header."""
|
|
return {
|
|
'parts': PartCatalogItem.query.count(),
|
|
'categories': PartCatalogCategory.query.count(),
|
|
'labelled': PartCatalogItem.query.filter_by(has_label=True).count(),
|
|
'mapped': PartCatalogCategory.query.filter(
|
|
PartCatalogCategory.drawer_category_id.isnot(None)).count(),
|
|
'last_sync': CatalogSync.latest(),
|
|
}
|