File size: 6,659 Bytes
de1e3fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
"""Delete the seeded demo data (generated-dot 'dogs' from scripts.make_sample_images) and everything
attached to it, leaving the real loaded datasets untouched.

Demo records are the ones created directly by scripts/seed.py (NOT via the batch loader), so they
have ``dataset_id IS NULL``. We remove:
  * every known/unknown dog with ``dataset_id IS NULL`` (Rex, Maple, the demo found dogs) +
    their pictures, embeddings, breed predictions, cases, matches, notifications, and media files;
  * the demo OWNER users of those dogs — but ONLY non-admin users who own no *dataset* dogs, so the
    admin account and the (dataset_id-NULL but real) detached dataset owners are preserved.

Idempotent. Dry-run by default.

Usage (from backend/):  python -m scripts.purge_seed_data --delete
"""
from __future__ import annotations

import argparse

from sqlalchemy import and_, delete, func, or_, select

from app.db import SessionLocal
from app.models import (
    BreedPrediction, Case, Embedding, KnownDog, Match, Notification, Picture, UnknownDog, User,
)
from app.models.base import SubjectType, UserRole
from app.storage import get_storage


def run(do_delete: bool) -> None:
    db = SessionLocal()
    try:
        known_ids = [i for (i,) in db.execute(select(KnownDog.id).where(KnownDog.dataset_id.is_(None)))]
        unknown_ids = [i for (i,) in db.execute(select(UnknownDog.id).where(UnknownDog.dataset_id.is_(None)))]

        # Demo owners: owners of these known dogs that are NOT admin and own no dataset dog.
        owner_ids = {
            o for (o,) in db.execute(
                select(KnownDog.owner_id).where(KnownDog.dataset_id.is_(None)).distinct()
            )
        }
        demo_user_ids = []
        for uid in owner_ids:
            u = db.get(User, uid)
            if u is None or u.role == UserRole.admin:
                continue
            keeps = db.execute(
                select(func.count()).select_from(KnownDog).where(
                    KnownDog.owner_id == uid, KnownDog.dataset_id.isnot(None)
                )
            ).scalar_one()
            if keeps == 0:
                demo_user_ids.append(uid)

        # Cases / pictures / matches tied to the demo dogs or demo owners.
        case_conds = []
        if known_ids:
            case_conds.append(Case.known_dog_id.in_(known_ids))
        if unknown_ids:
            case_conds.append(Case.unknown_dog_id.in_(unknown_ids))
        if demo_user_ids:
            case_conds.append(Case.person_id.in_(demo_user_ids))
        case_ids = (
            [i for (i,) in db.execute(select(Case.id).where(or_(*case_conds)))] if case_conds else []
        )

        pic_conds = []
        if known_ids:
            pic_conds.append(and_(Picture.subject_type == SubjectType.known, Picture.subject_id.in_(known_ids)))
        if unknown_ids:
            pic_conds.append(and_(Picture.subject_type == SubjectType.unknown, Picture.subject_id.in_(unknown_ids)))
        pic_rows = (
            db.execute(select(Picture.id, Picture.file_path, Picture.thumb_path).where(or_(*pic_conds))).all()
            if pic_conds else []
        )
        pic_ids = [r[0] for r in pic_rows]
        media_keys = [k for r in pic_rows for k in (r[1], r[2]) if k]

        match_conds = []
        if case_ids:
            match_conds.append(Match.case_id.in_(case_ids))
            match_conds.append(Match.candidate_case_id.in_(case_ids))
        if known_ids:
            match_conds.append(and_(Match.candidate_type == SubjectType.known, Match.candidate_id.in_(known_ids)))
        if unknown_ids:
            match_conds.append(and_(Match.candidate_type == SubjectType.unknown, Match.candidate_id.in_(unknown_ids)))
        match_ids = (
            [i for (i,) in db.execute(select(Match.id).where(or_(*match_conds)))] if match_conds else []
        )

        print("Seed data to remove:")
        print(f"  known dogs:   {len(known_ids)}  {known_ids}")
        print(f"  unknown dogs: {len(unknown_ids)}  {unknown_ids}")
        print(f"  demo owners:  {len(demo_user_ids)}  "
              f"{[db.get(User, u).email for u in demo_user_ids]}")
        print(f"  cases: {len(case_ids)}  matches: {len(match_ids)}  pictures: {len(pic_ids)}  "
              f"media files: {len(set(media_keys))}")

        if not do_delete:
            print("\nDRY RUN — nothing deleted. Re-run with --delete.")
            return

        def _del(stmt):
            return db.execute(stmt).rowcount or 0

        # Children first (FK order), mirroring services.datasets.purge_dataset.
        if demo_user_ids or case_ids or match_ids:
            nconds = []
            if demo_user_ids:
                nconds.append(Notification.user_id.in_(demo_user_ids))
            if case_ids:
                nconds.append(Notification.case_id.in_(case_ids))
            if match_ids:
                nconds.append(Notification.match_id.in_(match_ids))
            if nconds:
                _del(delete(Notification).where(or_(*nconds)))
        if pic_ids:
            _del(delete(BreedPrediction).where(BreedPrediction.picture_id.in_(pic_ids)))
            _del(delete(Embedding).where(Embedding.picture_id.in_(pic_ids)))
        if match_ids:
            _del(delete(Match).where(Match.id.in_(match_ids)))
        if demo_user_ids:
            db.execute(Match.__table__.update().where(Match.reviewed_by.in_(demo_user_ids)).values(reviewed_by=None))
        if pic_ids:
            _del(delete(Picture).where(Picture.id.in_(pic_ids)))
        if case_ids:
            _del(delete(Case).where(Case.id.in_(case_ids)))
        if unknown_ids:
            _del(delete(UnknownDog).where(UnknownDog.id.in_(unknown_ids)))
        if known_ids:
            _del(delete(KnownDog).where(KnownDog.id.in_(known_ids)))
        if demo_user_ids:
            _del(delete(User).where(User.id.in_(demo_user_ids)))
        db.commit()

        # Files after commit (best-effort; leftovers are recoverable via scripts.prune_media).
        storage = get_storage()
        removed = 0
        for key in dict.fromkeys(media_keys):
            try:
                storage.delete(key)
                removed += 1
            except Exception:  # noqa: BLE001
                pass
        print(f"\nDeleted seed data. Media files removed: {removed}. Admin + dataset dogs untouched.")
    finally:
        db.close()


def main() -> None:
    ap = argparse.ArgumentParser(description="Purge seeded demo (generated-dot) data.")
    ap.add_argument("--delete", action="store_true", help="Actually delete (default: dry run)")
    args = ap.parse_args()
    run(args.delete)


if __name__ == "__main__":
    main()