diff --git a/src/guidellm/data/deserializers/file.py b/src/guidellm/data/deserializers/file.py index f46e0ae52..b58c83712 100644 --- a/src/guidellm/data/deserializers/file.py +++ b/src/guidellm/data/deserializers/file.py @@ -196,7 +196,18 @@ def __call__( f"expected str or Path to a local .db file, got {data}" ) - return Dataset.from_sql(con=str(path), **data_kwargs) + import sqlite3 + + conn = sqlite3.connect(path) + cursor = conn.cursor() + table = cursor.execute( + "SELECT name FROM sqlite_master WHERE type='table'" + ).fetchone()[0] + conn.close() + + query = f"SELECT * FROM {table}" # noqa: S608 ignore table name injection risk + + return Dataset.from_sql(sql=query, con=str(path), **data_kwargs) @DatasetDeserializerFactory.register("tar_file")