From e3b472bfd19e3f3101fbe606fd2a640c08a20d02 Mon Sep 17 00:00:00 2001 From: NopeCharon <115794254+NopeCharon@users.noreply.github.com> Date: Sat, 31 Oct 2026 22:14:07 +0530 Subject: [PATCH] feat(annotator): let subclasses abort VCF annotation AbstractVcfAnnotator.annotate() catches every exception raised while processing a record, writes it to that record's VRS_Error INFO field and moves on to the next record. This is the right behavior for per-variant problems, but leaves child classes no way to signal an unrecoverable, global failure from hooks like on_vrs_object(). For example, an implementation that batch-inserts alleles into a database keeps iterating (and retrying the insert on every record) after the database goes away. Add VcfAnnotatorAbortError, a VcfAnnotatorError subclass that annotate() re-raises instead of recording, so annotation stops and the error reaches the caller. The input and output VCF files are now closed in a finally block so an aborted run doesn't leave them open, and the collection hook (e.g. PKL/NDJSON dumps) is not run for an aborted annotation. Closes #605 --- src/ga4gh/vrs/extras/annotator/vcf.py | 141 +++++++++++++++----------- tests/extras/test_annotate_vcf.py | 91 ++++++++++++++++- 2 files changed, 172 insertions(+), 60 deletions(-) diff --git a/src/ga4gh/vrs/extras/annotator/vcf.py b/src/ga4gh/vrs/extras/annotator/vcf.py index 3dd70dcf..7d1a3593 100644 --- a/src/ga4gh/vrs/extras/annotator/vcf.py +++ b/src/ga4gh/vrs/extras/annotator/vcf.py @@ -29,6 +29,18 @@ class VcfAnnotatorArgsError(VcfAnnotatorError): """Raise for improper args passed to VCF annotator""" +class VcfAnnotatorAbortError(VcfAnnotatorError): + """Raise to abort annotation of the entire VCF. + + By default, an exception raised while processing a VCF record is logged and + recorded in that record's ``VRS_Error`` INFO field, and annotation continues + with the next record. Child classes of ``AbstractVcfAnnotator`` can raise this + exception (e.g. from ``on_vrs_object``) to signal an unrecoverable error, such as + a lost database connection, that should instead stop annotation and propagate + to the caller of ``annotate()``. + """ + + class FieldName(str, Enum): """Define VCF field names for VRS annotations""" @@ -243,6 +255,8 @@ def annotate( validation checks fail, although all instances of failed validation are logged as warnings regardless. :raise VCFAnnotatorError: if no output formats are selected + :raise VcfAnnotatorAbortError: if a child class aborts annotation while a + record is being processed """ self.raise_for_output_args(output_vcf_path, **kwargs) # This can be pushed up to the click arg parsing too @@ -264,67 +278,72 @@ def annotate( if (self.collect_alleles and self.should_collect_alleles(**kwargs)) else None ) - for record in vcf: - if vcf_out: - additional_info_fields = [FieldName.IDS_FIELD] - if vrs_attributes: - additional_info_fields += [ - FieldName.STARTS_FIELD, - FieldName.ENDS_FIELD, - FieldName.STATES_FIELD, - FieldName.LENGTHS_FIELD, - FieldName.REPEAT_SUBUNIT_LENGTHS_FIELD, - ] - else: - # no INFO field names need to be designated if not producing an annotated VCF - additional_info_fields = [] - try: - vrs_field_data = self._get_vrs_data( - record, - allele_collection, - assembly, - additional_info_fields, - vrs_attributes=vrs_attributes, - compute_for_ref=compute_for_ref, - require_validation=require_validation, - **kwargs, - ) - except Exception as ex: - _logger.exception("VRS error on %s-%s", record.chrom, record.pos) - err_msg = f"{ex}" or f"{type(ex)}" - err_msg = err_msg.translate(VCF_ESCAPE_MAP) - additional_info_fields = [FieldName.ERROR_FIELD] - vrs_field_data = {FieldName.ERROR_FIELD.value: [err_msg]} - - _logger.debug( - "VCF record %s-%s generated vrs_field_data %s", - record.chrom, - record.pos, - vrs_field_data, - ) - - if output_vcf_path and vcf_out: - for k in additional_info_fields: - # pysam correctly converts None → "." for Integer/Float - # INFO fields, but for String fields it converts None → - # "" (empty bytes), violating the VCF spec. Work around - # by using the literal string "." for String-typed fields. - if k in _STRING_FIELDS: - record.info[k.value] = [ - "." if v in ("", None) else v - for v in vrs_field_data[k.value] - ] - else: - record.info[k.value] = [ - None if v in ("", None) else v - for v in vrs_field_data[k.value] + try: + for record in vcf: + if vcf_out: + additional_info_fields = [FieldName.IDS_FIELD] + if vrs_attributes: + additional_info_fields += [ + FieldName.STARTS_FIELD, + FieldName.ENDS_FIELD, + FieldName.STATES_FIELD, + FieldName.LENGTHS_FIELD, + FieldName.REPEAT_SUBUNIT_LENGTHS_FIELD, ] - vcf_out.write(record) - - vcf.close() + else: + # no INFO field names need to be designated if not producing an annotated VCF + additional_info_fields = [] + try: + vrs_field_data = self._get_vrs_data( + record, + allele_collection, + assembly, + additional_info_fields, + vrs_attributes=vrs_attributes, + compute_for_ref=compute_for_ref, + require_validation=require_validation, + **kwargs, + ) + except VcfAnnotatorAbortError: + _logger.exception( + "Aborting VCF annotation at %s-%s", record.chrom, record.pos + ) + raise + except Exception as ex: + _logger.exception("VRS error on %s-%s", record.chrom, record.pos) + err_msg = f"{ex}" or f"{type(ex)}" + err_msg = err_msg.translate(VCF_ESCAPE_MAP) + additional_info_fields = [FieldName.ERROR_FIELD] + vrs_field_data = {FieldName.ERROR_FIELD.value: [err_msg]} + + _logger.debug( + "VCF record %s-%s generated vrs_field_data %s", + record.chrom, + record.pos, + vrs_field_data, + ) - if vcf_out: - vcf_out.close() + if output_vcf_path and vcf_out: + for k in additional_info_fields: + # pysam correctly converts None → "." for Integer/Float + # INFO fields, but for String fields it converts None → + # "" (empty bytes), violating the VCF spec. Work around + # by using the literal string "." for String-typed fields. + if k in _STRING_FIELDS: + record.info[k.value] = [ + "." if v in ("", None) else v + for v in vrs_field_data[k.value] + ] + else: + record.info[k.value] = [ + None if v in ("", None) else v + for v in vrs_field_data[k.value] + ] + vcf_out.write(record) + finally: + vcf.close() + if vcf_out: + vcf_out.close() if allele_collection is not None: self.on_vrs_object_collection(allele_collection, **kwargs) @@ -339,6 +358,10 @@ def on_vrs_object( Reimplement in a child class to add custom logic. Otherwise, simply pass through ``vrs_allele`` without altering it further or storing it. + Exceptions raised here are recorded in the VCF record's ``VRS_Error`` INFO + field and annotation continues with the next record, unless the exception is + a ``VcfAnnotatorAbortError``, which stops annotation altogether. + :param vcf_coords: CHR-POS-REF-ALT from VCF for this allele :param vrs_allele: allele translated from coords :return: final VRS allele diff --git a/tests/extras/test_annotate_vcf.py b/tests/extras/test_annotate_vcf.py index 0d66dbe9..9dd2c39d 100644 --- a/tests/extras/test_annotate_vcf.py +++ b/tests/extras/test_annotate_vcf.py @@ -10,7 +10,12 @@ import pytest from ga4gh.vrs.dataproxy import DataProxyValidationError, SeqRepoRESTDataProxy -from ga4gh.vrs.extras.annotator.vcf import VcfAnnotator, VcfAnnotatorError +from ga4gh.vrs.extras.annotator.vcf import ( + VcfAnnotator, + VcfAnnotatorAbortError, + VcfAnnotatorError, +) +from ga4gh.vrs.models import Allele TEST_DATA_DIR = Path("tests/extras/data") @@ -314,3 +319,87 @@ def test_annotate_vcf_rle(vcf_annotator: VcfAnnotator, vcr_cassette, tmp_path: P assert output_vrs_pkl.exists() if vcr_cassette.write_protected: assert vcr_cassette.all_played + + +class FailingVcfAnnotator(VcfAnnotator): + """Raise ``error`` when ``on_vrs_object`` reaches the given VCF coordinates""" + + def __init__( + self, data_proxy: SeqRepoRESTDataProxy, fail_on: str, error: Exception + ) -> None: + """Initialize the annotator. + + :param data_proxy: GA4GH sequence dataproxy instance + :param fail_on: VCF coordinates (CHR-POS-REF-ALT) at which to raise ``error`` + :param error: exception to raise + """ + super().__init__(data_proxy) + self.fail_on = fail_on + self.error = error + + def on_vrs_object(self, vcf_coords: str, vrs_allele: Allele, **kwargs) -> Allele: + if vcf_coords == self.fail_on: + raise self.error + return super().on_vrs_object(vcf_coords, vrs_allele, **kwargs) + + +def test_annotate_vcf_record_error( + rest_dataproxy: SeqRepoRESTDataProxy, vcr, tmp_path: Path +): + """Test that a non-abort error in a child class hook only fails the current + record, and annotation continues with the rest of the VCF + """ + input_vcf = TEST_DATA_DIR / "test_rle.vcf" + output_vcf = tmp_path / "test_record_error_output.vcf" + output_vrs_pkl = tmp_path / "test_record_error_output.pkl" + + # test_rle.vcf is the input of test_annotate_vcf_rle, so its cassette covers every + # sequence request made here + with vcr.use_cassette("test_annotate_vcf_rle.yaml"): + vcf_annotator = FailingVcfAnnotator( + rest_dataproxy, + fail_on="1-100210778-AA-AA", + error=RuntimeError("Unable to store allele"), + ) + vcf_annotator.annotate(input_vcf, output_vcf, output_pkl_path=output_vrs_pkl) + + with pysam.VariantFile(str(output_vcf)) as vcf: + variants = list(vcf) + assert len(variants) == 2 + assert variants[0].info["VRS_Error"] == ("Unable to store allele",) + assert "VRS_Allele_IDs" not in variants[0].info + assert "VRS_Error" not in variants[1].info + assert len(variants[1].info["VRS_Allele_IDs"]) == 2 + assert output_vrs_pkl.exists() + + +def test_annotate_vcf_abort( + rest_dataproxy: SeqRepoRESTDataProxy, vcr, tmp_path: Path, caplog +): + """Test that raising ``VcfAnnotatorAbortError`` from a child class hook stops + annotation of the entire VCF and propagates to the caller + """ + input_vcf = TEST_DATA_DIR / "test_rle.vcf" + output_vcf = tmp_path / "test_abort_output.vcf.gz" + output_vrs_pkl = tmp_path / "test_abort_output.pkl" + + with vcr.use_cassette("test_annotate_vcf_rle.yaml"): + vcf_annotator = FailingVcfAnnotator( + rest_dataproxy, + fail_on="1-102995989-CTTT-CTTT", + error=VcfAnnotatorAbortError("Lost connection to storage"), + ) + with pytest.raises(VcfAnnotatorAbortError, match="Lost connection to storage"): + vcf_annotator.annotate( + input_vcf, output_vcf, output_pkl_path=output_vrs_pkl + ) + + assert "Aborting VCF annotation at 1-102995989" in caplog.text + # Collection hooks (e.g. PKL dump) are not run for an aborted annotation + assert not output_vrs_pkl.exists() + # Output VCF is closed, containing only the records annotated before aborting + with pysam.VariantFile(str(output_vcf)) as vcf: + variants = list(vcf) + assert len(variants) == 1 + assert variants[0].pos == 100210778 + assert len(variants[0].info["VRS_Allele_IDs"]) == 2