diff --git a/main.py b/main.py index 2dd74b7..369329d 100644 --- a/main.py +++ b/main.py @@ -1,15 +1,28 @@ # https://pysam.readthedocs.io/en/latest/api.html#fasta-files import json import pysam -from fastapi import FastAPI, HTTPException, Request, status +from fastapi import FastAPI, HTTPException, Request, status, Query import urllib import itertools import os +from starlette.middleware.cors import CORSMiddleware ALLOWED_URLS = os.environ.get("ALLOWED_URLS", "").split(",") +ALLOWED_ORIGINS_ENV = os.environ.get("ALLOWED_ORIGINS", "*") app = FastAPI() +# Configure CORS +_origins = [o.strip() for o in ALLOWED_ORIGINS_ENV.split(",") if o.strip()] +_allow_origins = ["*"] if not _origins or _origins == ["*"] else _origins +app.add_middleware( + CORSMiddleware, + allow_origins=_allow_origins, + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + def check_url(url): url = urllib.parse.unquote(url) if not any(url.startswith(allowed_url) for allowed_url in ALLOWED_URLS): @@ -150,8 +163,11 @@ def vcf_contigs(url: str): send_400_resp(f"Unable to open file: {e}") @app.get("/vcf/fetch/{seqid}:{start}-{end}/{url:path}") -def vcf_features(url: str, seqid: str, start: int, end: int): +def vcf_features(url: str, seqid: str, start: int, end: int, samples: list[str] = Query(default=[])): try: + vf = pysam.VariantFile(check_url(url)) + if samples: + vf.subset_samples(samples) return [ {"chrom": feature.chrom, "pos": feature.pos, "id": feature.id, @@ -164,15 +180,18 @@ def vcf_features(url: str, seqid: str, start: int, end: int): "samples": list(feature.samples), "alleles": feature.alleles} for feature - in pysam.VariantFile(check_url(url)).fetch(seqid, start, end) ] + in vf.fetch(seqid, start, end) ] except OSError as e: send_400_resp(f"Unable to open file: {e}") except KeyError as e: send_400_resp(f"Unable to find feature: {e}") @app.get("/vcf/fetch/{seqid}/{url:path}") -def vcf_features(url: str, seqid: str): +def vcf_features(url: str, seqid: str, samples: list[str] = Query(default=[])): try: + vf = pysam.VariantFile(check_url(url)) + if samples: + vf.subset_samples(samples) return [ {"chrom": feature.chrom, "pos": feature.pos, "id": feature.id, @@ -185,7 +204,7 @@ def vcf_features(url: str, seqid: str): "samples": list(feature.samples), "alleles": feature.alleles} for feature - in pysam.VariantFile(check_url(url)).fetch(seqid) ] + in vf.fetch(seqid) ] except OSError as e: send_400_resp(f"Unable to open file: {e}") except KeyError as e: @@ -291,3 +310,11 @@ def alignment_lengths(reference: str , url: str): except OSError as e: send_400_resp(f"Unable to open file: {e}") + +# Return all strains present in VCF file +@app.get("/vcf/samples/{url:path}") +def strains(url: str): + try: + return { "strains": list(pysam.VariantFile(check_url(url)).header.samples) } + except OSError as e: + send_400_resp(f"Unable to open file: {e}") \ No newline at end of file