"""Ingest collaborator splits from an XLSX file. Usage (CLI flags): poetry run python scripts/ingest_splits.py --file-path ingest.xlsx --ticket-id JIRA-123 [--dry-run true|false] (default: true) Usage (environment variables): FILE_PATH=ingest.xlsx TICKET_ID=JIRA-123 DRY_RUN=false poetry run python scripts/ingest_splits.py """ from collections import defaultdict from dataclasses import dataclass import logging from pathlib import Path import sys from typing import Any, Literal import openpyxl from pydantic import Field, ValidationError, model_validator from pydantic_settings import BaseSettings, SettingsConfigDict from sqlalchemy import select from sqlalchemy.orm.session import Session from collaborator.connectors import mysql, snowflake from collaborator.constants.split import RateType from collaborator.constants.split_type import SplitTypeId from collaborator.models.rds.collaborator import Collaborator from collaborator.models.rds.recipient import Recipient # noqa: F401 from collaborator.models.rds.split import Split from collaborator.models.rds.split_type import SplitType # noqa: F401 from collaborator.models.snowflake.account_payment_term import AccountPaymentTerm from collaborator.models.snowflake.product import Product from collaborator.models.snowflake.track import Track from collaborator.schemas import BaseSchema logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") logging.getLogger("snowflake.connector").setLevel(logging.WARNING) log = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- REQUIRED_HEADERS = { "Vendor ID", "Collaborator ID", "Collaborator Name", "Collaborator Split", "Split Type", "productid", } # tuid is optional: files may contain only product-ID-level rows class IngestConfig(BaseSettings): """Script configuration — populated from CLI flags or environment variables.""" model_config = SettingsConfigDict( cli_parse_args=True, env_ignore_empty=True, ) file_path: Path = Field(description="Path to the XLSX ingest file") ticket_id: str | None = Field( default=None, description="Ticket ID used as created_by / updated_by" ) dry_run: bool = Field( default=True, description="Print summary only; make no DB changes" ) environment: str = Field(default="dev", alias="Environment") build_tag: str = Field(default="ingest_splits", alias="BUILD_TAG") @model_validator(mode="after") def require_ticket_id_for_live_run(self) -> "IngestConfig": """Require ticket_id when dry_run is False.""" if not self.dry_run and not self.ticket_id: raise ValueError("ticket_id is required when dry_run is False") return self # --------------------------------------------------------------------------- # Row schema # --------------------------------------------------------------------------- class SplitRow(BaseSchema): """A single parsed and validated row from the ingest XLSX.""" vendor_id: int = Field(alias="Vendor ID") collaborator_id: int | None = Field(alias="Collaborator ID", default=None) collaborator_name: str | None = Field(alias="Collaborator Name", default=None) split_rate: float = Field(alias="Collaborator Split", gt=0, le=1) split_type: Literal["NET", "GROSS"] = Field(alias="Split Type") tuid: str | None = Field(alias="tuid", default=None) product_id: int | None = Field(alias="productid", default=None) @model_validator(mode="before") @classmethod def coerce_types(cls, values: dict) -> dict: """Normalise raw cell values before field-level validation.""" # tuid may come in as int from openpyxl; convert to str tuid = values.get("tuid") if tuid is not None and not isinstance(tuid, str): values["tuid"] = str(int(tuid)) # split_rate: accept whole-number percentages (e.g. 100 → 1.0, 50 → 0.5) rate = values.get("Collaborator Split") if isinstance(rate, (int, float)) and rate > 1: values["Collaborator Split"] = rate / 100 # split_type must be uppercase string split_type = values.get("Split Type") if isinstance(split_type, str): values["Split Type"] = split_type.strip().upper() # collaborator_name: strip whitespace name = values.get("Collaborator Name") if isinstance(name, str): values["Collaborator Name"] = name.strip() or None # product_id: coerce float cell values (e.g. 1.0 → 1) pid = values.get("productid") if pid is not None and not isinstance(pid, int): values["productid"] = int(pid) return values @model_validator(mode="after") def must_have_identifier(self) -> "SplitRow": """Require at least one of tuid or product_id.""" if not self.tuid and not self.product_id: raise ValueError("Each row must have either a tuid or a Product ID") return self # --------------------------------------------------------------------------- # XLSX loading # --------------------------------------------------------------------------- def load_xlsx(path: Path) -> list[dict[str, Any]]: """Load raw rows from an XLSX file using openpyxl. Returns a list of dicts keyed by the header row values. Empty/trailing rows are dropped. tuid values are preserved as their native cell type (int or str). """ wb = openpyxl.load_workbook(path, read_only=True, data_only=True) ws = wb.active all_rows = list(ws.iter_rows(values_only=True)) wb.close() if not all_rows: raise ValueError(f"File is empty: {path}") headers = [str(h).strip() if h is not None else "" for h in all_rows[0]] raw_rows = [] for row in all_rows[1:]: if any(cell is not None for cell in row): raw_rows.append(dict(zip(headers, row, strict=False))) return raw_rows # --------------------------------------------------------------------------- # Schema validation (file structure + row parsing) # --------------------------------------------------------------------------- def validate_headers(raw_rows: list[dict[str, Any]]) -> None: """Raise ValueError if any expected column is missing.""" if not raw_rows: raise ValueError("File contains no data rows") actual = set(raw_rows[0].keys()) missing = REQUIRED_HEADERS - actual if missing: raise ValueError(f"Missing required columns: {sorted(missing)}") def parse_rows(raw_rows: list[dict[str, Any]]) -> list[SplitRow]: """Parse and validate every row, collecting all errors before raising.""" errors: list[str] = [] rows: list[SplitRow] = [] for i, raw in enumerate(raw_rows, start=2): # row 1 is the header if raw.get("Collaborator Split") is None: log.warning("Row %d: skipping — 'Collaborator Split' is empty", i) continue try: rows.append(SplitRow.model_validate(raw)) except ValidationError as exc: errors.append(f"Row {i}: {exc.error_count()} error(s) — {exc.errors()}") if errors: raise ValueError( f"Validation failed for {len(errors)} row(s):\n" + "\n".join(errors) ) return rows def validate_single_vendor(rows: list[SplitRow]) -> int: """Ensure all rows share a single vendor_id and return it.""" vendor_ids = {row.vendor_id for row in rows} if len(vendor_ids) != 1: raise ValueError( f"File must contain exactly one vendor ID; found: {sorted(vendor_ids)}" ) return vendor_ids.pop() # --------------------------------------------------------------------------- # DB helpers — return plain data structures, not ORM objects, so that callers # are not affected by session expiry after the session context closes. # --------------------------------------------------------------------------- @mysql.db_session def _fetch_collaborator_vendor_map( collaborator_ids: set[int], session: Session ) -> dict[int, int]: """Return {collaborator_id: vendor_id} for the given IDs.""" rows = session.execute( select(Collaborator.collaborator_id, Collaborator.vendor_id).where( Collaborator.collaborator_id.in_(collaborator_ids) ) ).all() return {collaborator_id: vendor_id for collaborator_id, vendor_id in rows} @mysql.db_session def _fetch_existing_split_map(tuids: set[str], session: Session) -> dict[str, set[int]]: """Return {tuid: set(collaborator_ids)} for existing SplitTypeId.TRACK splits.""" rows = session.execute( select(Split.identifier, Split.collaborator_id).where( Split.identifier.in_(tuids), Split.split_type_id == SplitTypeId.TRACK, ) ).all() result: dict[str, set[int]] = defaultdict(set) for identifier, collaborator_id in rows: result[identifier].add(collaborator_id) return dict(result) @snowflake.db_session def _fetch_tuid_vendor_map(tuids: set[str], session: Session) -> dict[str, int]: """Return {tuid_str: vendor_id} for TUIDs found in DIM_TRACK. Uses an inner join to Product so each TUID maps to its product's vendor. TUIDs absent from the result do not exist in Snowflake. """ int_tuids = [int(t) for t in tuids if t.isdigit()] rows = session.execute( select(Track.tuid, Product.vendor_id) .select_from(Track) .join(Product, Track.upc == Product.upc) .where(Track.tuid.in_(int_tuids)) ).all() return {str(tuid): vendor_id for tuid, vendor_id in rows} @snowflake.db_session def _fetch_product_ids_for_vendor( product_ids: set[int], vendor_id: int, session: Session ) -> set[int]: """Return the subset of product IDs that belong to the vendor.""" rows = session.execute( select(Product.product_id).where( Product.product_id.in_(product_ids), Product.vendor_id == vendor_id, ) ).all() return {pid for (pid,) in rows} @snowflake.db_session def _fetch_tuid_product_pairs( product_ids: set[int], session: Session ) -> list[tuple[str, int]]: """Return (tuid_str, product_id) pairs for tracks matching the given product IDs. Implements the join: DIM_RELEASE → DIM_TRACK on DIM_TRACK.upc = DIM_RELEASE.releaseid, filtered by DIM_RELEASE.product_id. """ rows = session.execute( select(Track.tuid, Product.product_id) .select_from(Product) .join(Track, Track.upc == Product.upc) .where(Product.product_id.in_(product_ids), Track.tuid != 0) ).all() return [(str(tuid), product_id) for tuid, product_id in rows] # --------------------------------------------------------------------------- # DB-backed validation + UPC expansion # --------------------------------------------------------------------------- def validate_collaborators(rows: list[SplitRow], vendor_id: int) -> None: """Verify all collaborator IDs in the file exist and belong to the vendor.""" collaborator_ids = { r.collaborator_id for r in rows if r.collaborator_id is not None } if not collaborator_ids: return found = _fetch_collaborator_vendor_map(collaborator_ids) errors = [] for cid in sorted(collaborator_ids): if cid not in found: errors.append(f"Collaborator ID {cid} not found") elif found[cid] != vendor_id: errors.append( f"Collaborator ID {cid} belongs to vendor {found[cid]}, not {vendor_id}" ) if errors: raise ValueError("Collaborator validation failed:\n" + "\n".join(errors)) def resolve_product_ids( rows: list[SplitRow], vendor_id: int ) -> tuple[list[SplitRow], int]: """Expand product-ID-only rows to TUID rows; validate product IDs belong to the vendor. Returns (rows, skipped_count) where rows replaces all product-ID-only rows with one row per track under that product, and skipped_count is the number of product-ID rows dropped because the product had no eligible tracks. Rows that already have a tuid are unchanged. """ product_id_rows = [r for r in rows if r.product_id and not r.tuid] if not product_id_rows: return rows, 0 product_ids = {r.product_id for r in product_id_rows} valid_ids = _fetch_product_ids_for_vendor(product_ids, vendor_id) missing = product_ids - valid_ids if missing: raise ValueError( f"Product IDs not found for vendor {vendor_id}: {sorted(missing)}" ) track_pairs = _fetch_tuid_product_pairs(valid_ids) tuids_by_product: dict[int, list[str]] = defaultdict(list) for tuid_str, product_id in track_pairs: tuids_by_product[product_id].append(tuid_str) result = [r for r in rows if r.tuid] # preserve tuid-only rows unchanged skipped_ids: list[int] = [] for row in product_id_rows: pid = row.product_id assert pid is not None tuid_list = tuids_by_product.get(pid, []) if not tuid_list: skipped_ids.append(pid) continue for tuid_str in tuid_list: result.append(row.model_copy(update={"tuid": tuid_str})) if skipped_ids: log.warning( "%d product ID(s) had no eligible tracks and were skipped: %s", len(skipped_ids), sorted(skipped_ids), ) return result, len(skipped_ids) def validate_tracks(rows: list[SplitRow], vendor_id: int) -> None: """Verify all TUIDs in the file exist in Snowflake and belong to the vendor.""" tuids = {r.tuid for r in rows if r.tuid} if not tuids: return tuid_vendor_map = _fetch_tuid_vendor_map(tuids) missing = sorted(tuids - tuid_vendor_map.keys()) if missing: raise ValueError(f"TUIDs not found in Snowflake: {missing}") invalid = sorted(t for t, vid in tuid_vendor_map.items() if vid != vendor_id) if invalid: raise ValueError(f"TUIDs do not belong to vendor {vendor_id}: {invalid}") # --------------------------------------------------------------------------- # Change computation # --------------------------------------------------------------------------- @dataclass class IngestSummary: """Summary of what the ingest will do (or did, if dry_run=False).""" vendor_id: int input_rows: int # rows parsed from the file (pre-expansion) expanded_rows: int # rows after product IDs are resolved to TUIDs # collaborator counts existing_collaborators_matched: int new_collaborators: int # track-level counts tracks_with_existing_splits_updated: int tracks_with_new_splits: int tracks_with_unaffected_splits: int # existing splits for other collaborators # split-level counts existing_splits_updated: int new_splits_created: int products_skipped_no_tracks: int vendor_currency: str dry_run: bool def compute_changes( rows: list[SplitRow], vendor_id: int, input_rows: int, vendor_currency: str, products_skipped_no_tracks: int = 0, newly_created_collabs: int = 0, dry_run: bool = True, ) -> IngestSummary: """Compute what the ingest would do by comparing file rows against existing splits. All rows must already have tuid set (call resolve_product_ids before this). """ collaborator_ids = { r.collaborator_id for r in rows if r.collaborator_id is not None } new_collaborator_names = { r.collaborator_name for r in rows if r.collaborator_id is None and r.collaborator_name } tuids = {r.tuid for r in rows if r.tuid} existing_by_tuid = _fetch_existing_split_map(tuids) # {tuid: set(collaborator_ids)} # Group known-collaborator rows by TUID → set of collaborator IDs ingest_by_tuid: dict[str, set[int]] = defaultdict(set) for r in rows: if r.tuid and r.collaborator_id is not None: ingest_by_tuid[r.tuid].add(r.collaborator_id) # Group new-collaborator rows by TUID → set of names (deduplicates same name per TUID) new_collab_by_tuid: dict[str, set[str]] = defaultdict(set) for r in rows: if r.tuid and r.collaborator_id is None and r.collaborator_name: new_collab_by_tuid[r.tuid].add(r.collaborator_name) tracks_with_existing_updated = 0 tracks_with_new = 0 tracks_with_unaffected = 0 existing_splits_updated = 0 new_splits_created = 0 for tuid in tuids: existing_collaborator_ids = existing_by_tuid.get(tuid, set()) tuid_collaborator_ids = ingest_by_tuid.get(tuid, set()) tuid_new_collab_names = new_collab_by_tuid.get(tuid, set()) overlapping = tuid_collaborator_ids & existing_collaborator_ids new_for_track = tuid_collaborator_ids - existing_collaborator_ids unaffected = existing_collaborator_ids - tuid_collaborator_ids existing_splits_updated += len(overlapping) new_splits_created += len(new_for_track) + len(tuid_new_collab_names) if overlapping: tracks_with_existing_updated += 1 if new_for_track or tuid_new_collab_names: tracks_with_new += 1 if unaffected: tracks_with_unaffected += 1 return IngestSummary( vendor_id=vendor_id, input_rows=input_rows, expanded_rows=len(rows), existing_collaborators_matched=len(collaborator_ids) - newly_created_collabs, new_collaborators=len(new_collaborator_names) + newly_created_collabs, tracks_with_existing_splits_updated=tracks_with_existing_updated, tracks_with_new_splits=tracks_with_new, tracks_with_unaffected_splits=tracks_with_unaffected, existing_splits_updated=existing_splits_updated, new_splits_created=new_splits_created, products_skipped_no_tracks=products_skipped_no_tracks, vendor_currency=vendor_currency, dry_run=dry_run, ) # --------------------------------------------------------------------------- # Write step # --------------------------------------------------------------------------- @mysql.db_session def _upsert_splits(rows: list[SplitRow], ticket_id: str, session: Session) -> None: """Upsert SplitTypeId.TRACK splits for the given rows. Existing splits are updated in-place; new splits are added. Splits for collaborators not present in *rows* are left untouched. All rows must have both tuid and collaborator_id set before calling this. """ tuids = {r.tuid for r in rows if r.tuid} collaborator_ids = { r.collaborator_id for r in rows if r.collaborator_id is not None } existing: dict[tuple[str, int], Split] = { (s.identifier, s.collaborator_id): s for s in session.execute( select(Split).where( Split.identifier.in_(tuids), Split.collaborator_id.in_(collaborator_ids), Split.split_type_id == SplitTypeId.TRACK, ) ) .scalars() .all() } for row in rows: if not row.tuid or row.collaborator_id is None: continue key = (row.tuid, row.collaborator_id) if key in existing: s = existing[key] s.split_rate = row.split_rate s.rate_type = RateType(row.split_type) s.updated_by = ticket_id else: session.add( Split( identifier=row.tuid, collaborator_id=row.collaborator_id, split_rate=row.split_rate, split_type_id=SplitTypeId.TRACK, rate_type=RateType(row.split_type), created_by=ticket_id, ) ) session.commit() # --------------------------------------------------------------------------- # Collaborator creation helpers # --------------------------------------------------------------------------- @snowflake.db_session def _fetch_vendor_currency(vendor_id: int, session: Session) -> str: """Return the payment currency for a vendor from Snowflake ACCOUNT_PAYMENT_TERM.""" rows = session.execute( select(AccountPaymentTerm.currency_code).where( AccountPaymentTerm.account_id == vendor_id ) ).all() if not rows: raise ValueError(f"No account_payment_term found for vendor {vendor_id}") if len(rows) > 1: raise ValueError( f"Multiple account_payment_term rows found for vendor {vendor_id}" ) return rows[0][0] @mysql.db_session def _create_missing_collaborators( rows: list[SplitRow], vendor_id: int, ticket_id: str, currency: str, session: Session, ) -> tuple[list[SplitRow], int]: """Create collaborators for rows with no collaborator_id. Idempotent: if a COLLABORATOR with the same name already exists for the vendor it is reused rather than duplicated. Returns (updated_rows, newly_created_count). """ names_needed = { r.collaborator_name for r in rows if r.collaborator_id is None and r.collaborator_name } if not names_needed: return rows, 0 name_to_id: dict[str, int] = { c.name: c.collaborator_id for c in session.execute( select(Collaborator).where( Collaborator.vendor_id == vendor_id, Collaborator.name.in_(names_needed), Collaborator.collaborator_type == "COLLABORATOR", ) ) .scalars() .all() } newly_created = 0 for name in sorted(names_needed - set(name_to_id)): collab = Collaborator( name=name, vendor_id=vendor_id, currency=currency, collaborator_type="COLLABORATOR", performance_rights=True, created_by=ticket_id, ) session.add(collab) session.flush() # populate collaborator_id before commit name_to_id[name] = collab.collaborator_id newly_created += 1 session.commit() updated_rows = [ ( row.model_copy( update={"collaborator_id": name_to_id[row.collaborator_name]} ) if row.collaborator_id is None and row.collaborator_name in name_to_id else row ) for row in rows ] return updated_rows, newly_created # TODO: implement fetch_vendor_has_direct_payments when direct payments support is added. # --------------------------------------------------------------------------- # Dry-run output # --------------------------------------------------------------------------- def print_summary(summary: IngestSummary) -> None: """Print a human-readable ingest summary.""" mode = "DRY RUN" if summary.dry_run else "APPLIED" log.info("=" * 60) log.info( " ".join( [ f"Ingest summary [{mode}]", f"vendor_id={summary.vendor_id}", f"currency={summary.vendor_currency}", ] ) ) log.info("=" * 60) # Rows — show expansion inline if product IDs were resolved to TUIDs if summary.expanded_rows != summary.input_rows: log.info( f" Rows in file : {summary.input_rows:,}" f" → {summary.expanded_rows:,} after product ID expansion" ) else: log.info(f" Rows in file : {summary.input_rows:,}") # Collaborators — one combined line collab_total = summary.existing_collaborators_matched + summary.new_collaborators log.info( f" Collaborators : {collab_total:,} total" f" ({summary.existing_collaborators_matched:,} existing," f" {summary.new_collaborators:,} to create)" ) if summary.products_skipped_no_tracks: log.info( f" Products skipped (no tracks) : {summary.products_skipped_no_tracks:,}" ) # Track-level breakdown — always shown so the counts can be cross-checked log.info(f" Tracks — splits added : {summary.tracks_with_new_splits:,}") log.info( f" Tracks — splits updated : {summary.tracks_with_existing_splits_updated:,}" ) log.info( f" Tracks — other splits intact : {summary.tracks_with_unaffected_splits:,}" ) # Split totals — suppress zeros to highlight what's actually changing if summary.new_splits_created or summary.existing_splits_updated: log.info(f" Splits to create : {summary.new_splits_created:,}") log.info( f" Splits to update : {summary.existing_splits_updated:,}" ) else: log.info(" No split changes.") log.info("=" * 60) if summary.dry_run: log.info(" Re-run with --dry-run=false to apply changes.") # --------------------------------------------------------------------------- # Orchestration # --------------------------------------------------------------------------- def run(config: IngestConfig) -> IngestSummary: """Run the ingest pipeline, returning a summary of changes.""" log.info(f"Loading file: {config.file_path}") raw_rows = load_xlsx(config.file_path) log.info("Validating file structure…") validate_headers(raw_rows) rows = parse_rows(raw_rows) vendor_id = validate_single_vendor(rows) input_rows = len(rows) log.info(f"File OK — {input_rows} rows, vendor_id={vendor_id}") log.info("Validating collaborators…") validate_collaborators(rows, vendor_id) log.info("Resolving product IDs to TUIDs…") rows, products_skipped_no_tracks = resolve_product_ids(rows, vendor_id) log.info("Validating tracks…") validate_tracks(rows, vendor_id) newly_created_collabs = 0 currency = _fetch_vendor_currency(vendor_id) if not config.dry_run and any(r.collaborator_id is None for r in rows): log.info("Creating missing collaborators…") rows, newly_created_collabs = _create_missing_collaborators( rows, vendor_id, config.ticket_id, currency ) log.info("Computing changes…") summary = compute_changes( rows, vendor_id, input_rows, currency, products_skipped_no_tracks, newly_created_collabs=newly_created_collabs, dry_run=config.dry_run, ) if not config.dry_run: log.info("Applying changes…") _upsert_splits(rows, config.ticket_id) log.info("Done.") return summary def main() -> None: """Entry point for the ingest script.""" config = IngestConfig() # type: ignore[call-arg] # fields populated from sys.argv log.info(f"ticket_id={config.ticket_id} dry_run={config.dry_run}") summary = run(config) print_summary(summary) if __name__ == "__main__": try: main() except (ValueError, NotImplementedError) as exc: log.exception(str(exc)) sys.exit(1)