Add the Update button: Rebrickable CSV dump sync
This commit is contained in:
jessikitty committed
1 parent
63c3ece271
commit
f6f1f4e2bc
1 file changed
+263
@@ -0,0 +1,263 @@
|
||||
"""
|
||||
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(),
|
||||
}
|
||||
Reference in new issue
Block a user