diff --git a/CHANGELOG.md b/CHANGELOG.md index 8965f29..8499883 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,16 @@ All notable changes to `atb-cli` are documented here. Format follows [Keep a Cha ## [Unreleased] +### Fixed + +- `atb amr --species "Enterococcus faecium"` (an unsuffixed NCBI name) returned + 0 results when a genus SQLite index was present. The AMR index path matched the + raw query with an exact `LOWER(species) IN (...)` clause and never applied the + GTDB-suffix stripping used everywhere else, so v0.22.0's GTDB matching reached + the parquet scan but not the SQLite index. The index path now resolves the + requested species against the stored names with the same GTDB-aware matcher + before querying. Explicit clade names still select a single clade. + ## [v0.22.0](https://github.com/allthebacteria/atb-cli/releases/tag/v0.22.0) - 2026-09-02 ### Changed diff --git a/internal/amr/amr_gtdb_test.go b/internal/amr/amr_gtdb_test.go index 0f30ad2..c2eb02e 100644 --- a/internal/amr/amr_gtdb_test.go +++ b/internal/amr/amr_gtdb_test.go @@ -72,6 +72,57 @@ func TestQueryAMRSpeciesGTDBSuffix(t *testing.T) { } } +// TestQueryAMRSpeciesGTDBSuffixIndexed verifies that an NCBI-style --species +// query matches GTDB-split rows through the SQLite index path (not just the +// parquet scan). Building indexes makes Query route via QueryIndex/buildSQL. +func TestQueryAMRSpeciesGTDBSuffixIndexed(t *testing.T) { + dir := t.TempDir() + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_A.parquet"), + []amrFixtureRow{{species: "Enterococcus_A faecium", genus: "Enterococcus_A", count: 2}}) + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_B.parquet"), + []amrFixtureRow{{species: "Enterococcus_B faecium", genus: "Enterococcus_B", count: 3}}) + + if err := amr.BuildIndexes(dir, nil); err != nil { + t.Fatalf("BuildIndexes: %v", err) + } + + results, err := amr.Query(dir, amr.Filters{ + Genera: []string{"Enterococcus"}, + Species: []string{"Enterococcus faecium"}, + }) + if err != nil { + t.Fatalf("Query: %v", err) + } + if len(results) != 5 { + t.Fatalf("expected 5 GTDB-split rows via index, got %d", len(results)) + } +} + +// TestQueryAMRSpeciesCladeSelectIndexed verifies that an explicit clade name +// through the SQLite index path selects only that clade. +func TestQueryAMRSpeciesCladeSelectIndexed(t *testing.T) { + dir := t.TempDir() + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_A.parquet"), + []amrFixtureRow{{species: "Enterococcus_A faecium", genus: "Enterococcus_A", count: 2}}) + writeAMRParquetAt(t, filepath.Join(dir, amr.PartitionDir, "Enterococcus_B.parquet"), + []amrFixtureRow{{species: "Enterococcus_B faecium", genus: "Enterococcus_B", count: 3}}) + + if err := amr.BuildIndexes(dir, nil); err != nil { + t.Fatalf("BuildIndexes: %v", err) + } + + results, err := amr.Query(dir, amr.Filters{ + Genera: []string{"Enterococcus"}, + Species: []string{"Enterococcus_B faecium"}, + }) + if err != nil { + t.Fatalf("Query: %v", err) + } + if len(results) != 3 { + t.Fatalf("expected 3 rows for explicit _B clade, got %d", len(results)) + } +} + // TestQueryAMRGenusGTDBPartitions verifies that an NCBI-style genus query reads // every on-disk GTDB clade partition, even when no monolithic file is present. func TestQueryAMRGenusGTDBPartitions(t *testing.T) { diff --git a/internal/amr/indexer.go b/internal/amr/indexer.go index 295a401..b123934 100644 --- a/internal/amr/indexer.go +++ b/internal/amr/indexer.go @@ -300,6 +300,17 @@ func QueryIndex(dbPath string, filters Filters) ([]Result, error) { } } + if len(filters.Species) > 0 { + resolved, err := resolveIndexSpecies(db, filters.Species) + if err != nil { + return nil, wrapStaleIndexError(dbPath, err) + } + if len(resolved) == 0 { + return nil, nil + } + filters.Species = resolved + } + query, args := buildSQL(filters) rows, err := db.Query(query, args...) if err != nil { @@ -329,6 +340,33 @@ func QueryIndex(dbPath string, filters Filters) ([]Result, error) { return results, rows.Err() } +// resolveIndexSpecies reads the distinct species stored in the index and returns +// the concrete values that match any requested query under GTDB-aware matching. +// SQLite cannot evaluate the GTDB suffix regex, so matching happens in Go and the +// resolved names feed buildSQL's exact IN clause. +func resolveIndexSpecies(db *sql.DB, queries []string) ([]string, error) { + rows, err := db.Query("SELECT DISTINCT species FROM amr") + if err != nil { + return nil, err + } + defer rows.Close() + + var resolved []string + for rows.Next() { + var stored string + if err := rows.Scan(&stored); err != nil { + return nil, err + } + for _, q := range queries { + if match.SpeciesMatches(q, stored) { + resolved = append(resolved, stored) + break + } + } + } + return resolved, rows.Err() +} + func buildSQL(f Filters) (string, []any) { var clauses []string var args []any