"""Master catalog merchandising associations (related / upsell / cross-sell).""" from __future__ import annotations from typing import Any, Dict, Iterable, List, Optional, Sequence, Set import sqlalchemy as sa from sqlalchemy import select from sqlalchemy.dialects.postgresql import insert from sqlalchemy.orm import Session from db.association_rules import ( DEFAULT_CAP, LINK_TYPES, ORIGIN_AUTO, ORIGIN_MANUAL, KitchenCabinetAssociationRules, ProductSnapshot, build_catalog_index, clamp_cap, extract_manual_overrides_from_payload, normalize_category_l2, ) from db.models import MasterProduct, MasterProductAssociation, MasterProductPrice, MasterProductRelation DEFAULT_BATCH_SIZE = 250 def associations_by_source( session: Session, source_skus: Optional[Iterable[str]] = None, *, link_types: Optional[Iterable[str]] = None, ) -> Dict[str, Dict[str, List[Dict[str, Any]]]]: """Return {source_sku: {link_type: [row_dicts...]}} ordered by sort_order.""" stmt = select(MasterProductAssociation).order_by( MasterProductAssociation.source_sku, MasterProductAssociation.link_type, MasterProductAssociation.sort_order, MasterProductAssociation.target_sku, ) if source_skus is not None: sources = sorted({str(s).strip() for s in source_skus if str(s).strip()}) if not sources: return {} stmt = stmt.where(MasterProductAssociation.source_sku.in_(sources)) if link_types is not None: types = sorted({str(t).strip() for t in link_types if str(t).strip()}) if types: stmt = stmt.where(MasterProductAssociation.link_type.in_(types)) grouped: Dict[str, Dict[str, List[Dict[str, Any]]]] = {} for row in session.scalars(stmt).all(): by_type = grouped.setdefault(row.source_sku, {}) by_type.setdefault(row.link_type, []).append(_association_dict(row)) return grouped def association_fields_by_sku( session: Session, skus: Iterable[str], ) -> Dict[str, Dict[str, List[str]]]: """Compact {sku: {related: [...], upsell: [...], crosssell: [...]}} for exports.""" grouped = associations_by_source(session, skus) out: Dict[str, Dict[str, List[str]]] = {} for sku, by_type in grouped.items(): out[sku] = { link_type: [row["target_sku"] for row in rows] for link_type, rows in by_type.items() } return out def upsert_associations(session: Session, rows: List[Dict[str, Any]]) -> int: if not rows: return 0 written = 0 chunk_size = 1000 for start in range(0, len(rows), chunk_size): chunk = rows[start : start + chunk_size] stmt = insert(MasterProductAssociation).values(chunk) stmt = stmt.on_conflict_do_update( constraint="uq_master_product_association_source_target_type", set_={ "sort_order": stmt.excluded.sort_order, "origin": stmt.excluded.origin, "source_label": stmt.excluded.source_label, "updated_at": sa.func.now(), }, ) session.execute(stmt) written += len(chunk) return written def delete_auto_associations_for_sources( session: Session, source_skus: Iterable[str], *, link_types: Optional[Iterable[str]] = None, ) -> int: sources = sorted({str(s).strip() for s in source_skus if str(s).strip()}) if not sources: return 0 deleted = 0 chunk_size = 500 type_filter = None if link_types is not None: types = sorted({str(t).strip() for t in link_types if str(t).strip()}) if types: type_filter = types for start in range(0, len(sources), chunk_size): chunk = sources[start : start + chunk_size] stmt = sa.delete(MasterProductAssociation).where( MasterProductAssociation.source_sku.in_(chunk), MasterProductAssociation.origin == ORIGIN_AUTO, ) if type_filter: stmt = stmt.where(MasterProductAssociation.link_type.in_(type_filter)) result = session.execute(stmt) deleted += int(result.rowcount or 0) return deleted def replace_manual_associations_for_source( session: Session, source_sku: str, overrides: Dict[str, List[str]], *, source_label: str = "manual_override", ) -> int: """Replace manual rows for the given source+types; leaves auto rows untouched.""" sku = str(source_sku or "").strip() if not sku or not overrides: return 0 written = 0 for link_type, targets in overrides.items(): if link_type not in LINK_TYPES: continue session.execute( sa.delete(MasterProductAssociation).where( MasterProductAssociation.source_sku == sku, MasterProductAssociation.link_type == link_type, MasterProductAssociation.origin == ORIGIN_MANUAL, ) ) rows = [ { "source_sku": sku, "target_sku": target, "link_type": link_type, "sort_order": index, "origin": ORIGIN_MANUAL, "source_label": source_label, } for index, target in enumerate(targets) if target and target != sku ] written += upsert_associations(session, rows) return written def apply_manual_overrides_from_records( session: Session, records: Sequence[Dict[str, Any]], *, source_label: str = "master_sku", ) -> int: written = 0 for record in records: sku = str(record.get("sku") or "").strip() payload = record.get("payload") or {} if not sku: continue overrides = extract_manual_overrides_from_payload(payload) if overrides: written += replace_manual_associations_for_source( session, sku, overrides, source_label=source_label ) return written def discover_associations( session: Session, *, skus: Optional[Iterable[str]] = None, cap: int = DEFAULT_CAP, dry_run: bool = False, source_label: str = "association_discovery", bidirectional: bool = True, offset: int = 0, limit: Optional[int] = None, ) -> Dict[str, Any]: """Discover related/upsell/cross-sell for SKUs and optionally peer reverse links. Manual rows are preserved. Auto rows for touched source SKUs are replaced. Use offset/limit to process a slice of the sorted active SKU list. """ rules = KitchenCabinetAssociationRules(cap=clamp_cap(cap)) catalog = _load_catalog_snapshots(session, skus=skus) if not catalog: return { "status": "ok", "dry_run": dry_run, "sources": 0, "auto_rows": 0, "peer_sources": 0, "planned": [], "offset": offset, "limit": limit, "total_candidates": 0, } index = build_catalog_index(catalog) by_sku = index.by_sku variant_exclusions = _variant_exclusion_sets(session) if skus is None: all_sources = sorted(by_sku.keys()) else: all_sources = sorted( {str(s).strip() for s in skus if str(s).strip() and str(s).strip() in by_sku} ) total_candidates = len(all_sources) start = max(0, int(offset or 0)) if limit is None: source_list = all_sources[start:] else: source_list = all_sources[start : start + max(0, int(limit))] planned_by_source: Dict[str, Dict[str, List[str]]] = {} for source_sku in source_list: source = by_sku[source_sku] discovered = rules.discover_for_source(source, index) planned_by_source[source_sku] = { link_type: _filter_targets(source_sku, targets, variant_exclusions) for link_type, targets in discovered.items() } peer_sources: Set[str] = set() if bidirectional: for source_sku in list(source_list): source = by_sku[source_sku] reverse_map = rules.reverse_peers_for_source(source, index) for link_type, peer_skus in reverse_map.items(): for peer_sku in peer_skus: # Only append this source into the peer's auto list; ranking/cap happens below. peer_sources.add(peer_sku) bucket = planned_by_source.setdefault( peer_sku, {t: [] for t in LINK_TYPES} ) targets = bucket.setdefault(link_type, []) if source_sku not in targets: targets.append(source_sku) # Cap peer reverse lists; outbound source lists already capped by rules. for peer_sku in peer_sources: if peer_sku in source_list: # Already fully discovered outbound; reverse only adds extras into buckets. for link_type in LINK_TYPES: planned_by_source[peer_sku][link_type] = _filter_targets( peer_sku, planned_by_source[peer_sku].get(link_type) or [], variant_exclusions, )[: rules.cap] else: # Peer not in this batch: keep only reverse hits for this batch, capped. for link_type in LINK_TYPES: planned_by_source[peer_sku][link_type] = _filter_targets( peer_sku, planned_by_source[peer_sku].get(link_type) or [], variant_exclusions, )[: rules.cap] # Merge with existing manual targets (manual wins; auto fills remaining slots). existing = associations_by_source(session, planned_by_source.keys()) final_rows: List[Dict[str, Any]] = [] summary_planned: List[Dict[str, Any]] = [] for source_sku, by_type in sorted(planned_by_source.items()): manual_by_type = { link_type: [r["target_sku"] for r in rows if r.get("origin") == ORIGIN_MANUAL] for link_type, rows in (existing.get(source_sku) or {}).items() } # For peers outside this batch, merge reverse auto hits with existing auto rows # so we don't wipe their previously discovered outbound associations. existing_auto_by_type = { link_type: [r["target_sku"] for r in rows if r.get("origin") == ORIGIN_AUTO] for link_type, rows in (existing.get(source_sku) or {}).items() } preserve_existing_auto = source_sku not in source_list merged: Dict[str, List[str]] = {} for link_type in LINK_TYPES: manual = manual_by_type.get(link_type) or [] auto_new = [s for s in (by_type.get(link_type) or []) if s not in manual] if preserve_existing_auto: auto_existing = [ s for s in (existing_auto_by_type.get(link_type) or []) if s not in manual and s not in auto_new ] auto = auto_existing + auto_new else: auto = auto_new merged[link_type] = (manual + auto)[: rules.cap] for index_pos, target in enumerate(merged[link_type]): origin = ORIGIN_MANUAL if target in manual else ORIGIN_AUTO if origin == ORIGIN_AUTO: final_rows.append( { "source_sku": source_sku, "target_sku": target, "link_type": link_type, "sort_order": index_pos, "origin": ORIGIN_AUTO, "source_label": source_label, } ) summary_planned.append({"sku": source_sku, "associations": merged}) next_offset = start + len(source_list) has_more = next_offset < total_candidates if dry_run: return { "status": "ok", "dry_run": True, "sources": len(source_list), "peer_sources": len(peer_sources), "auto_rows": len(final_rows), "planned": summary_planned[:50], "planned_total": len(summary_planned), "offset": start, "limit": limit, "next_offset": next_offset, "has_more": has_more, "total_candidates": total_candidates, } # Only replace auto rows for sources we fully rediscovered in this batch. # Peer-only sources get a merge-style rewrite of their auto rows (already merged above). touched_sources = sorted(planned_by_source.keys()) deleted = delete_auto_associations_for_sources(session, touched_sources) upserted = upsert_associations(session, final_rows) return { "status": "ok", "dry_run": False, "sources": len(source_list), "peer_sources": len(peer_sources), "auto_deleted": deleted, "auto_rows": upserted, "planned_total": len(summary_planned), "offset": start, "limit": limit, "next_offset": next_offset, "has_more": has_more, "total_candidates": total_candidates, } def discover_associations_in_batches( session: Session, *, skus: Optional[Iterable[str]] = None, cap: int = DEFAULT_CAP, dry_run: bool = False, source_label: str = "association_discovery", bidirectional: bool = True, batch_size: int = DEFAULT_BATCH_SIZE, offset: int = 0, max_batches: Optional[int] = None, commit: bool = True, progress_callback: Optional[Any] = None, ) -> Dict[str, Any]: """Run discovery in batches with optional per-batch commits and progress callbacks.""" batch = max(1, int(batch_size or DEFAULT_BATCH_SIZE)) current = max(0, int(offset or 0)) batches_run = 0 totals = { "status": "ok", "dry_run": dry_run, "batch_size": batch, "offset_start": current, "sources": 0, "peer_sources": 0, "auto_rows": 0, "auto_deleted": 0, "batches": [], "total_candidates": None, } while True: if max_batches is not None and batches_run >= max_batches: break result = discover_associations( session, skus=skus, cap=cap, dry_run=dry_run, source_label=source_label, bidirectional=bidirectional, offset=current, limit=batch, ) if totals["total_candidates"] is None: totals["total_candidates"] = result.get("total_candidates", 0) totals["sources"] += int(result.get("sources") or 0) totals["peer_sources"] += int(result.get("peer_sources") or 0) totals["auto_rows"] += int(result.get("auto_rows") or 0) totals["auto_deleted"] += int(result.get("auto_deleted") or 0) batch_summary = { "batch": batches_run + 1, "offset": result.get("offset"), "sources": result.get("sources"), "peer_sources": result.get("peer_sources"), "auto_rows": result.get("auto_rows"), "next_offset": result.get("next_offset"), "has_more": result.get("has_more"), } totals["batches"].append(batch_summary) batches_run += 1 if progress_callback is not None: progress_callback( { **batch_summary, "total_candidates": totals["total_candidates"], "sources_done": totals["sources"], "auto_rows_done": totals["auto_rows"], } ) if not dry_run and commit: session.commit() if not result.get("has_more"): break current = int(result.get("next_offset") or (current + batch)) totals["batches_run"] = batches_run totals["offset_end"] = current if not totals["batches"] else totals["batches"][-1].get("next_offset") totals["has_more"] = bool(totals["batches"][-1]["has_more"]) if totals["batches"] else False return totals def compute_associations_hash(associations: Dict[str, List[str]]) -> str: import hashlib import json normalized = { link_type: list(associations.get(link_type) or []) for link_type in LINK_TYPES } raw = json.dumps(normalized, sort_keys=True, separators=(",", ":")) return hashlib.sha256(raw.encode("utf-8")).hexdigest() def _load_catalog_snapshots( session: Session, *, skus: Optional[Iterable[str]] = None, ) -> List[ProductSnapshot]: source_skus = sorted({str(s).strip() for s in (skus or []) if str(s).strip()}) if source_skus: source_products = session.scalars( select(MasterProduct) .where(MasterProduct.is_active.is_(True)) .where(MasterProduct.sku.in_(source_skus)) ).all() if not source_products: return [] source_collections = sorted({str(row.collection or "").strip() for row in source_products if str(row.collection or "").strip()}) products_by_sku = {row.sku: row for row in source_products if str(row.sku or "").strip()} if source_collections: collection_products = session.scalars( select(MasterProduct) .where(MasterProduct.is_active.is_(True)) .where(MasterProduct.collection.in_(source_collections)) ).all() for row in collection_products: sku = str(row.sku or "").strip() if sku: products_by_sku[sku] = row products = list(products_by_sku.values()) else: products = session.scalars( select(MasterProduct).where(MasterProduct.is_active.is_(True)) ).all() if not products: return [] product_ids = [int(row.id) for row in products if getattr(row, "id", None) is not None] prices = { row.sku: float(row.amount) for row in session.scalars( select(MasterProductPrice).where( MasterProductPrice.product_id.in_(product_ids), MasterProductPrice.price_type == "base", MasterProductPrice.amount.is_not(None), ) ).all() if row.amount is not None } return [ ProductSnapshot( sku=p.sku, collection=str(p.collection or "").strip(), category_l2=normalize_category_l2(p.category_l2), product_family=str(p.product_family or "").strip(), name=str(p.name or "").strip(), price=prices.get(p.sku), ) for p in products if str(p.sku or "").strip() ] def _variant_exclusion_sets(session: Session) -> Dict[str, Set[str]]: """Map each SKU to siblings it should not associate with (same configurable family).""" rows = session.scalars(select(MasterProductRelation)).all() children_by_parent: Dict[str, Set[str]] = {} parent_by_child: Dict[str, str] = {} for row in rows: children_by_parent.setdefault(row.parent_sku, set()).add(row.child_sku) parent_by_child[row.child_sku] = row.parent_sku exclusions: Dict[str, Set[str]] = {} for parent, children in children_by_parent.items(): family = set(children) | {parent} for sku in family: exclusions[sku] = family - {sku} return exclusions def _filter_targets( source_sku: str, targets: Sequence[str], exclusions: Dict[str, Set[str]], ) -> List[str]: blocked = exclusions.get(source_sku) or set() out: List[str] = [] seen: Set[str] = set() for target in targets: sku = str(target).strip() if not sku or sku == source_sku or sku in seen or sku in blocked: continue seen.add(sku) out.append(sku) return out def _association_dict(row: MasterProductAssociation) -> Dict[str, Any]: return { "source_sku": row.source_sku, "target_sku": row.target_sku, "link_type": row.link_type, "sort_order": row.sort_order, "origin": row.origin, "source_label": row.source_label, }