Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
51 changes: 51 additions & 0 deletions internal/amr/amr_gtdb_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
38 changes: 38 additions & 0 deletions internal/amr/indexer.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
Expand Down
Loading