Skip to content

Commit 494ebd9

Browse files
UlthranUlthran
authored
Merge pull request #43 from PennChopMicrobiomeProgram/codex/add-utility-to-remove-an-isolate
Add isolate removal utility and CLI subcommand
2 parents 10af68a + 3367b69 commit 494ebd9

3 files changed

Lines changed: 228 additions & 0 deletions

File tree

‎marc_db/cli.py‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
from marc_db import __version__
44
from marc_db.db import create_database, get_session, get_marc_db_url
55
from marc_db.ingest import ingest_from_tsvs
6+
from marc_db.remove import remove_isolate
67
from marc_db.mock import fill_mock_db
78

89

@@ -13,6 +14,7 @@ def main():
1314
" init \tInitialize a new database.\n"
1415
" mock_db \tFill mock values into an empty db (for testing).\n"
1516
" ingest \tIngest data from TSV files into the database.\n"
17+
" remove \tRemove an isolate and associated records.\n"
1618
)
1719

1820
parser = argparse.ArgumentParser(
@@ -100,6 +102,25 @@ def main():
100102
yes=args_ingest.yes,
101103
session=get_session(db_url),
102104
)
105+
elif args.command == "remove":
106+
parser_remove = argparse.ArgumentParser(
107+
prog="marc_db remove",
108+
usage="%(prog)s --sample-id SAMPLE_ID",
109+
description="Remove an isolate and its associated data.",
110+
)
111+
parser_remove.add_argument(
112+
"--sample-id", required=True, help="SampleID of the isolate to remove."
113+
)
114+
parser_remove.add_argument(
115+
"--yes", action="store_true", help="Skip confirmation prompt."
116+
)
117+
args_remove = parser_remove.parse_args(remaining)
118+
create_database(db_url)
119+
remove_isolate(
120+
sample_id=args_remove.sample_id,
121+
yes=args_remove.yes,
122+
session=get_session(db_url),
123+
)
103124
else:
104125
parser.print_help()
105126
sys.stderr.write("Unrecognized command.\n")

‎marc_db/remove.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
from typing import Callable, Optional
2+
3+
from sqlalchemy import select
4+
from sqlalchemy.orm import Session
5+
6+
from marc_db.db import get_session
7+
from marc_db.models import (
8+
Aliquot,
9+
Antimicrobial,
10+
Assembly,
11+
AssemblyQC,
12+
Contaminant,
13+
Isolate,
14+
TaxonomicAssignment,
15+
)
16+
17+
18+
def _summarize_isolate(session: Session, sample_id: str) -> dict:
19+
assembly_ids = select(Assembly.id).where(Assembly.isolate_id == sample_id)
20+
return {
21+
"aliquots": session.query(Aliquot)
22+
.filter(Aliquot.isolate_id == sample_id)
23+
.count(),
24+
"assemblies": session.query(Assembly)
25+
.filter(Assembly.isolate_id == sample_id)
26+
.count(),
27+
"assembly_qc": session.query(AssemblyQC)
28+
.filter(AssemblyQC.assembly_id.in_(assembly_ids))
29+
.count(),
30+
"taxonomic_assignments": session.query(TaxonomicAssignment)
31+
.filter(TaxonomicAssignment.assembly_id.in_(assembly_ids))
32+
.count(),
33+
"contaminants": session.query(Contaminant)
34+
.filter(Contaminant.assembly_id.in_(assembly_ids))
35+
.count(),
36+
"antimicrobials": session.query(Antimicrobial)
37+
.filter(Antimicrobial.assembly_id.in_(assembly_ids))
38+
.count(),
39+
}
40+
41+
42+
def remove_isolate(
43+
*,
44+
sample_id: str,
45+
yes: bool = False,
46+
session: Optional[Session] = None,
47+
input_fn: Callable[[str], str] = input,
48+
):
49+
"""Remove a single isolate and its associated records."""
50+
51+
created_session = False
52+
if session is None:
53+
session = get_session()
54+
created_session = True
55+
56+
trans = session.begin_nested() if session.in_transaction() else session.begin()
57+
try:
58+
isolate = session.get(Isolate, sample_id)
59+
if isolate is None:
60+
print(f"No isolate found with SampleID {sample_id}.")
61+
trans.rollback()
62+
return
63+
64+
counts = _summarize_isolate(session, sample_id)
65+
66+
if not yes:
67+
print(f"Isolate {sample_id} will be removed with the following records:")
68+
for label, count in counts.items():
69+
print(f" {label.replace('_', ' ')}: {count}")
70+
answer = input_fn("Proceed with deletion? [y/N]: ").strip().lower()
71+
if answer not in {"y", "yes"}:
72+
trans.rollback()
73+
print("Removal cancelled.")
74+
return
75+
76+
assembly_ids = select(Assembly.id).where(Assembly.isolate_id == sample_id)
77+
78+
session.query(Antimicrobial).filter(
79+
Antimicrobial.assembly_id.in_(assembly_ids)
80+
).delete(synchronize_session=False)
81+
session.query(Contaminant).filter(
82+
Contaminant.assembly_id.in_(assembly_ids)
83+
).delete(synchronize_session=False)
84+
session.query(TaxonomicAssignment).filter(
85+
TaxonomicAssignment.assembly_id.in_(assembly_ids)
86+
).delete(synchronize_session=False)
87+
session.query(AssemblyQC).filter(
88+
AssemblyQC.assembly_id.in_(assembly_ids)
89+
).delete(synchronize_session=False)
90+
session.query(Assembly).filter(
91+
Assembly.isolate_id == sample_id
92+
).delete(synchronize_session=False)
93+
session.query(Aliquot).filter(Aliquot.isolate_id == sample_id).delete(
94+
synchronize_session=False
95+
)
96+
session.query(Isolate).filter(Isolate.sample_id == sample_id).delete(
97+
synchronize_session=False
98+
)
99+
100+
trans.commit()
101+
except Exception:
102+
trans.rollback()
103+
raise
104+
finally:
105+
if created_session:
106+
session.close()

‎tests/test_remove.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import pandas as pd
2+
from pathlib import Path
3+
from sqlalchemy import create_engine
4+
from sqlalchemy.orm import sessionmaker
5+
6+
from marc_db.ingest import ingest_from_tsvs
7+
from marc_db.models import (
8+
Aliquot,
9+
Antimicrobial,
10+
Assembly,
11+
AssemblyQC,
12+
Base,
13+
Isolate,
14+
TaxonomicAssignment,
15+
)
16+
from marc_db.remove import remove_isolate
17+
18+
19+
data_dir = Path(__file__).parent
20+
21+
22+
def _build_session():
23+
engine = create_engine("sqlite:///:memory:")
24+
Session = sessionmaker(bind=engine)
25+
session = Session()
26+
Base.metadata.create_all(engine)
27+
return session, engine
28+
29+
30+
def test_remove_isolate_deletes_associations():
31+
session, engine = _build_session()
32+
isolates_df = pd.read_csv(data_dir / "test_multi_aliquot.tsv", sep="\t")
33+
assemblies_df = pd.read_csv(data_dir / "test_assembly_data.tsv", sep="\t")
34+
tax_df = pd.read_csv(data_dir / "test_taxonomic_assignment.tsv", sep="\t")
35+
amr_df = pd.read_csv(data_dir / "test_amr_data.tsv", sep="\t")
36+
37+
ingest_from_tsvs(
38+
isolates=isolates_df,
39+
assemblies=assemblies_df,
40+
assembly_qcs=assemblies_df,
41+
taxonomic_assignments=tax_df,
42+
antimicrobials=amr_df,
43+
yes=True,
44+
session=session,
45+
)
46+
47+
remove_isolate(sample_id="sample1", yes=True, session=session)
48+
49+
remaining_sample = "sample2"
50+
expected_aliquots = isolates_df.loc[
51+
isolates_df["SampleID"] == remaining_sample
52+
].shape[0]
53+
expected_assemblies = assemblies_df.loc[
54+
assemblies_df["SampleID"] == remaining_sample
55+
].shape[0]
56+
expected_taxonomic = tax_df.loc[tax_df["SampleID"] == remaining_sample].shape[0]
57+
expected_amr = amr_df.loc[amr_df["SampleID"] == remaining_sample].shape[0]
58+
59+
assert session.query(Isolate).count() == 1
60+
assert session.query(Aliquot).count() == expected_aliquots
61+
assert session.query(Assembly).count() == expected_assemblies
62+
63+
remaining_assembly_ids = [
64+
asm.id
65+
for asm in session.query(Assembly).filter(
66+
Assembly.isolate_id == remaining_sample
67+
)
68+
]
69+
assert session.query(AssemblyQC).filter(
70+
AssemblyQC.assembly_id.in_(remaining_assembly_ids)
71+
).count() == expected_assemblies
72+
assert session.query(TaxonomicAssignment).filter(
73+
TaxonomicAssignment.assembly_id.in_(remaining_assembly_ids)
74+
).count() == expected_taxonomic
75+
assert session.query(Antimicrobial).filter(
76+
Antimicrobial.assembly_id.in_(remaining_assembly_ids)
77+
).count() == expected_amr
78+
79+
session.close()
80+
engine.dispose()
81+
82+
83+
def test_remove_isolate_cancelled_does_not_delete(capsys):
84+
session, engine = _build_session()
85+
isolates_df = pd.read_csv(data_dir / "test_multi_aliquot.tsv", sep="\t")
86+
87+
ingest_from_tsvs(isolates=isolates_df, yes=True, session=session)
88+
89+
remove_isolate(
90+
sample_id="sample1",
91+
yes=False,
92+
session=session,
93+
input_fn=lambda _: "n",
94+
)
95+
96+
captured = capsys.readouterr()
97+
assert "Removal cancelled." in captured.out
98+
assert session.query(Isolate).count() == 2
99+
100+
session.close()
101+
engine.dispose()

0 commit comments

Comments
 (0)