feat!: store URIs for ingested VCFs

.github/workflows/checks.yaml

@@ -1,7 +1,7 @@
 name: Checks
 on: [push, pull_request]
 jobs:
-  test:
+  python-test:
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -34,6 +34,19 @@ jobs:
 
       - name: Check style
         run: python3 -m ruff check . && ruff format --check .
+  rust-test:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./rust
+    steps:
+      - uses: actions/checkout@v4
+
+      - run: rustup update stable && rustup default stable
+
+      - run: cargo build --verbose
+
+      - run: cargo test --verbose
   rust-lint:
     runs-on: ubuntu-latest
     defaults:

README.md

@@ -8,6 +8,12 @@ From a VCF, ingest a VRS ID and the corresponding VCF-called location (i.e. suff
 % vrsix load chr1.vcf
 ```
 
+By default, this uses the input file's location in the local file system as a URI, but a custom URI can be declared:
+
+```shell
+% vrsix load chr1.vcf gs://my_stuff/chr1.vcf
+```
+
 Given a VRS ID, retrieve VCF-associated data (output format TBD)
 
 ```shell
@@ -62,7 +68,14 @@ cd rust/
 cargo fmt
 ```
 
-Run tests with `pytest`:
+Run tests from the project root with `pytest`:
 ```shell
 pytest
 ```
+
+Some granular tests are written directly into the Rust backend as well:
+
+```shell
+cd rust/
+cargo test
+```

rust/Cargo.toml

@@ -15,3 +15,4 @@ sqlx = { version = "0.8.2", features = ["runtime-tokio", "sqlite"] }
 tokio = { version = "1.40.0", features = ["full"] }
 pyo3 = { version = "0.21.1", features = ["abi3-py310", "experimental-async"]}
 log = "0.4.22"
+tempfile = "3.14.0"

rust/src/lib.rs

@@ -5,9 +5,12 @@ use std::path::PathBuf;
 use tokio::runtime::Runtime;
 
 #[pyfunction]
-pub fn vcf_to_sqlite(vcf_path: PathBuf, db_url: String) -> PyResult<()> {
+#[pyo3(signature = (vcf_path, db_url, vcf_uri = None))]
+pub fn vcf_to_sqlite(vcf_path: PathBuf, db_url: String, vcf_uri: Option<String>) -> PyResult<()> {
+    let uri_value =
+        vcf_uri.unwrap_or_else(|| format!("file://{}", vcf_path.to_string_lossy().into_owned()));
     let rt = Runtime::new().unwrap();
-    rt.block_on(load::load_vcf(vcf_path, &db_url))?;
+    rt.block_on(load::load_vcf(vcf_path, &db_url, uri_value))?;
     Ok(())
 }
 

rust/src/load.rs

@@ -19,12 +19,14 @@ use tokio::{
 
 async fn load_allele(db_row: DbRow, pool: &SqlitePool) -> Result<(), Box<dyn std::error::Error>> {
     let mut conn = pool.acquire().await?;
-    let result = sqlx::query("INSERT INTO vrs_locations (vrs_id, chr, pos) VALUES (?, ?, ?)")
-        .bind(db_row.vrs_id)
-        .bind(db_row.chr)
-        .bind(db_row.pos)
-        .execute(&mut *conn)
-        .await;
+    let result =
+        sqlx::query("INSERT INTO vrs_locations (vrs_id, chr, pos, uri_id) VALUES (?, ?, ?, ?);")
+            .bind(db_row.vrs_id)
+            .bind(db_row.chr)
+            .bind(db_row.pos)
+            .bind(db_row.uri_id)
+            .execute(&mut *conn)
+            .await;
     if let Err(err) = result {
         if let Some(db_error) = err.as_database_error() {
             if let Some(sqlite_error) = db_error.try_downcast_ref::<SqliteError>() {
@@ -93,7 +95,25 @@ async fn get_reader(
     }
 }
 
-pub async fn load_vcf(vcf_path: PathBuf, db_url: &str) -> PyResult<()> {
+async fn load_file_uri(uri: &str, pool: &SqlitePool) -> Result<i64, Box<dyn std::error::Error>> {
+    let mut conn = pool.acquire().await?;
+
+    let insert_result = sqlx::query("INSERT OR IGNORE INTO file_uris (uri) VALUES (?);")
+        .bind(uri)
+        .execute(&mut *conn)
+        .await?;
+    if insert_result.rows_affected() > 0 {
+        Ok(insert_result.last_insert_rowid())
+    } else {
+        let row_id: (i64,) = sqlx::query_as("SELECT id FROM file_uris WHERE uri = ?;")
+            .bind(uri)
+            .fetch_one(&mut *conn)
+            .await?;
+        Ok(row_id.0)
+    }
+}
+
+pub async fn load_vcf(vcf_path: PathBuf, db_url: &str, uri: String) -> PyResult<()> {
     let start = Instant::now();
 
     if !vcf_path.exists() || !vcf_path.is_file() {
@@ -118,6 +138,10 @@ pub async fn load_vcf(vcf_path: PathBuf, db_url: &str) -> PyResult<()> {
         VrsixDbError::new_err(format!("Failed database connection/call: {}", e))
     })?;
 
+    let uri_id = load_file_uri(&uri, &db_pool)
+        .await
+        .map_err(|e| VrsixDbError::new_err(format!("Failed to insert file URI `{uri}`: {e}")))?;
+
     while let Some(record) = records.try_next().await? {
         let vrs_ids = get_vrs_ids(record.info(), &header)?;
         let chrom = record.reference_sequence_name();
@@ -131,6 +155,7 @@ pub async fn load_vcf(vcf_path: PathBuf, db_url: &str) -> PyResult<()> {
                     .to_string(),
                 chr: chrom.strip_prefix("chr").unwrap_or(chrom).to_string(),
                 pos: pos.try_into().unwrap(),
+                uri_id,
             };
             load_allele(row, &db_pool).await.map_err(|e| {
                 error!("Failed to load row {:?}", e);
@@ -143,3 +168,19 @@ pub async fn load_vcf(vcf_path: PathBuf, db_url: &str) -> PyResult<()> {
     info!("Time taken: {:?}", duration);
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::NamedTempFile;
+
+    #[tokio::test]
+    async fn test_load_file_uri() {
+        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
+        let db_url = format!("sqlite://{}", temp_file.path().to_str().unwrap());
+        crate::sqlite::setup_db(&db_url).await.unwrap();
+        let db_pool = get_db_connection(&db_url).await.unwrap();
+        let uri_id = load_file_uri("file:///sdlfkjd", &db_pool).await.unwrap();
+        assert!(uri_id == 1);
+    }
+}

rust/src/sqlite.rs

@@ -20,13 +20,20 @@ pub async fn setup_db(db_url: &str) -> Result<(), Error> {
     let db = get_db_connection(db_url).await?;
     let result = sqlx::query(
         "
+        CREATE TABLE IF NOT EXISTS file_uris (
+            id INTEGER PRIMARY KEY,
+            uri TEXT UNIQUE
+        );
         CREATE TABLE IF NOT EXISTS vrs_locations (
             id INTEGER PRIMARY KEY AUTOINCREMENT,
             vrs_id TEXT NOT NULL,
             chr TEXT NOT NULL,
             pos INTEGER NOT NULL,
-            UNIQUE(vrs_id,chr,pos)
-        );",
+            uri_id INTEGER NOT NULL,
+            FOREIGN KEY (uri_id) REFERENCES file_uris(id),
+            UNIQUE(vrs_id, chr, pos, uri_id)
+        );
+        ",
     )
     .execute(&db)
     .await?;
@@ -39,4 +46,18 @@ pub struct DbRow {
     pub vrs_id: String,
     pub chr: String,
     pub pos: i64,
+    pub uri_id: i64,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::NamedTempFile;
+
+    #[tokio::test]
+    async fn test_setup_db() {
+        let temp_file = NamedTempFile::new().expect("Failed to create temp file");
+        let db_url = format!("sqlite://{}", temp_file.path().to_str().unwrap());
+        setup_db(&db_url).await.expect("Setup DB failed");
+    }
 }

src/vrsix/cli.py

@@ -31,31 +31,40 @@ def cli() -> None:
 
 @cli.command()
 @click.argument(
-    "vcfs",
+    "vcf",
     type=click.Path(
         exists=True, file_okay=True, dir_okay=False, readable=True, path_type=Path
     ),
-    nargs=-1,
+    nargs=1,
+)
+@click.argument(
+    "uri",
+    default=None,
+    required=False,
 )
 @click.option(
     "--db-location",
     type=click.Path(
         file_okay=True, dir_okay=True, readable=True, writable=True, path_type=Path
     ),
 )
-def load(vcfs: tuple[Path], db_location: Path | None) -> None:
+def load(vcf: Path, uri: str, db_location: Path | None) -> None:
     """Index the VRS annotations in a VCF by loading it into the sqlite DB.
 
+    Optionally provide a custom file URI to describe how to retrieve VCF records after
+    index lookup:
+
+        % vrsix load input.vcf gs://my_storage/input.vcf
+
     \f
     :param vcf_path: path to VCF to ingest
     """
     if db_location and db_location.is_dir():
         db_location = db_location / "vrs_vcf_index.db"
-    for vcf in vcfs:
-        start = timer()
-        load_vcf.load_vcf(vcf, db_location)
-        end = timer()
-        _logger.info("Processed `%s` in %s seconds", vcf, end - start)
+    start = timer()
+    load_vcf.load_vcf(vcf, db_location, uri)
+    end = timer()
+    _logger.info("Processed `%s` in %s seconds", vcf, end - start)
 
 
 @cli.command()

src/vrsix/load.py

@@ -26,16 +26,20 @@
 _logger = logging.getLogger(__name__)
 
 
-def load_vcf(vcf_path: Path, db_location: Path | None = None) -> None:
+def load_vcf(
+    vcf_path: Path, db_location: Path | None = None, vcf_uri: str | None = None
+) -> None:
     """Load VRS-annotated VCF into sqlite database.
 
     :param vcf_path: path to VCF (must exist) to ingest
     :param db_location: path to sqlite DB
+    :param vcf_uri: optional URI for original input. Not validated or parsed in any way.
+        Used by fetching lib to re-acquire data.
     """
     sqlite_uri = (
         f"sqlite://{DEFAULT_SQLITE_LOCATION}"
         if db_location is None
         else f"sqlite://{db_location}"
     )
     _logger.debug("Using sqlite file located at %s", sqlite_uri)
-    vcf_to_sqlite(vcf_path, sqlite_uri)
+    vcf_to_sqlite(vcf_path, sqlite_uri, vcf_uri)

tests/test_load.py

@@ -24,16 +24,39 @@ def test_load(fixture_dir: Path, temp_dir: Path, input_filename: str):
     results = conn.execute("SELECT * FROM vrs_locations").fetchall()
     assert len(results) == 10
     assert results == [
-        (1, "dwwiZdvVtfAmomu0OBsiHue1O-bw5SpG", "1", 783006),
-        (2, "MiasxyXMXtOpsZgGelL3c4QgtflCNLHD", "1", 783006),
-        (3, "5cY2k53xdW7WeHw2WG1HA7jl50iH-r9p", "1", 783175),
-        (4, "jHaXepIvlbnapfPtH_62y-Qm81hCrBYn", "1", 783175),
-        (5, "-NGsjBEx0UbPF3uYjStZ_2r-m2LbUtUB", "1", 784860),
-        (6, "HLinVo6Q-i-PryQOiq8QAtOeC9oQ9Q3p", "1", 784860),
-        (7, "qdyeeiC3cLfXeT23zxT9-qlJNN64MKVB", "1", 785417),
-        (8, "cNWXR3OLq9D3L19vQFvbHw-aH0vlA5cN", "1", 785417),
-        (9, "DVMcfA37Llc9QUOA0XfLJbJ-agKyGpGo", "1", 797392),
-        (10, "OTiBHLE2WW93M4-4zGVrWSqP2GBj8-qM", "1", 797392),
+        (1, "dwwiZdvVtfAmomu0OBsiHue1O-bw5SpG", "1", 783006, 1),
+        (2, "MiasxyXMXtOpsZgGelL3c4QgtflCNLHD", "1", 783006, 1),
+        (3, "5cY2k53xdW7WeHw2WG1HA7jl50iH-r9p", "1", 783175, 1),
+        (4, "jHaXepIvlbnapfPtH_62y-Qm81hCrBYn", "1", 783175, 1),
+        (5, "-NGsjBEx0UbPF3uYjStZ_2r-m2LbUtUB", "1", 784860, 1),
+        (6, "HLinVo6Q-i-PryQOiq8QAtOeC9oQ9Q3p", "1", 784860, 1),
+        (7, "qdyeeiC3cLfXeT23zxT9-qlJNN64MKVB", "1", 785417, 1),
+        (8, "cNWXR3OLq9D3L19vQFvbHw-aH0vlA5cN", "1", 785417, 1),
+        (9, "DVMcfA37Llc9QUOA0XfLJbJ-agKyGpGo", "1", 797392, 1),
+        (10, "OTiBHLE2WW93M4-4zGVrWSqP2GBj8-qM", "1", 797392, 1),
+    ]
+    conn.close()
+
+
+def test_load_specify_uri(fixture_dir: Path, temp_dir: Path):
+    input_file = fixture_dir / "input.vcf"
+    temp_db = temp_dir / "tmp.db"
+    load.load_vcf(input_file, temp_db)
+
+    conn = sqlite3.connect(temp_db)
+    results = conn.execute("SELECT * FROM vrs_locations").fetchall()
+    assert len(results) == 10
+    assert results == [
+        (1, "dwwiZdvVtfAmomu0OBsiHue1O-bw5SpG", "1", 783006, 1),
+        (2, "MiasxyXMXtOpsZgGelL3c4QgtflCNLHD", "1", 783006, 1),
+        (3, "5cY2k53xdW7WeHw2WG1HA7jl50iH-r9p", "1", 783175, 1),
+        (4, "jHaXepIvlbnapfPtH_62y-Qm81hCrBYn", "1", 783175, 1),
+        (5, "-NGsjBEx0UbPF3uYjStZ_2r-m2LbUtUB", "1", 784860, 1),
+        (6, "HLinVo6Q-i-PryQOiq8QAtOeC9oQ9Q3p", "1", 784860, 1),
+        (7, "qdyeeiC3cLfXeT23zxT9-qlJNN64MKVB", "1", 785417, 1),
+        (8, "cNWXR3OLq9D3L19vQFvbHw-aH0vlA5cN", "1", 785417, 1),
+        (9, "DVMcfA37Llc9QUOA0XfLJbJ-agKyGpGo", "1", 797392, 1),
+        (10, "OTiBHLE2WW93M4-4zGVrWSqP2GBj8-qM", "1", 797392, 1),
     ]
     conn.close()