Spaces:
Sleeping
Sleeping
| """Postgres + Mongo materialization into the session DuckDB. | |
| Self-contained: builds a tiny pg_dump-style .sql and a tiny mongodump .bson | |
| in tmp_path. Postgres needs the embedded `pgserver`; Mongo needs `bson` | |
| (pymongo) — both in the `connectors` extra. Each test skips cleanly if its | |
| dependency isn't installed. | |
| """ | |
| from __future__ import annotations | |
| import sqlite3 | |
| import duckdb | |
| import pytest | |
| from lexsi_ds.agent.datasource import AttachSpec, load_attached_handle | |
| bson = pytest.importorskip("bson", reason="needs pymongo/bson (connectors extra)") | |
| def _mongo_dump(folder, collection, docs): | |
| folder.mkdir(parents=True, exist_ok=True) | |
| with open(folder / f"{collection}.bson", "wb") as f: | |
| for d in docs: | |
| f.write(bson.encode(d)) | |
| # mongodump also drops a metadata.json — ensure it's ignored | |
| (folder / f"{collection}.metadata.json").write_text("{}") | |
| def test_mongo_dump_materializes_with_nested_as_json(tmp_path): | |
| dump = tmp_path / "reviews_dump" | |
| _mongo_dump(dump, "reviews", [ | |
| {"_id": bson.ObjectId(), "book_id": 1, "stars": 5, "tags": ["scifi", "classic"]}, | |
| {"_id": bson.ObjectId(), "book_id": 2, "stars": 3, "tags": ["horror"]}, | |
| ]) | |
| handle = load_attached_handle( | |
| [AttachSpec(alias="rev", db_type="mongo", dump_folder=dump, db_name="reviews_db")], | |
| dataset_id="t:mongo", | |
| ) | |
| names = {t.name for t in handle.tables} | |
| assert "rev_reviews" in names | |
| con = duckdb.connect(str(handle.duckdb_path), read_only=True) | |
| rows = con.execute("SELECT book_id, stars, tags FROM rev_reviews ORDER BY book_id").fetchall() | |
| con.close() | |
| assert rows[0][:2] == (1, 5) | |
| assert rows[0][2] == '["scifi", "classic"]' # nested array → JSON text | |
| # queryable as JSON | |
| con = duckdb.connect(str(handle.duckdb_path), read_only=True) | |
| n = con.execute("SELECT count(*) FROM rev_reviews " | |
| "WHERE list_contains(json_transform(tags, '[\"VARCHAR\"]'), 'horror')").fetchone()[0] | |
| con.close() | |
| assert n == 1 | |
| def test_sqlite_mixed_type_blob_column_falls_back(tmp_path): | |
| """Reproduces the krama bug: a dynamically-typed sqlite column holding a | |
| float in one row and an invalid-unicode blob in another. DuckDB's typed | |
| scanner rejects it; the Python sqlite3 fallback must still load both rows.""" | |
| import sqlite3 | |
| p = tmp_path / "dyn.db" | |
| c = sqlite3.connect(p) | |
| c.execute("CREATE TABLE t (id INTEGER, shape)") # no affinity → dynamic typing | |
| c.execute("INSERT INTO t VALUES (1, 3.14)") | |
| c.execute("INSERT INTO t VALUES (2, ?)", (b"\xff\xfeGP",)) # non-utf8 blob | |
| c.commit(); c.close() | |
| handle = load_attached_handle( | |
| [AttachSpec(alias="d", db_type="sqlite", path=p)], dataset_id="t:dyn") | |
| assert "d_t" in {t.name for t in handle.tables} | |
| con = duckdb.connect(str(handle.duckdb_path), read_only=True) | |
| n = con.execute("SELECT count(*) FROM d_t").fetchone()[0] | |
| con.close() | |
| assert n == 2 | |
| def test_postgres_dump_materializes(tmp_path): | |
| pytest.importorskip("pgserver", reason="needs pgserver (connectors extra)") | |
| sql = tmp_path / "books.sql" | |
| sql.write_text( | |
| "CREATE TABLE books (id integer, title text, author text);\n" | |
| "COPY books (id, title, author) FROM stdin;\n" | |
| "1\tDune\tHerbert\n" | |
| "2\tIt\tKing\n" | |
| "\\.\n" | |
| ) | |
| handle = load_attached_handle( | |
| [AttachSpec(alias="books_database", db_type="postgres", sql_file=sql, db_name="bookreview_db")], | |
| dataset_id="t:pg", | |
| ) | |
| assert "books_database_books" in {t.name for t in handle.tables} | |
| con = duckdb.connect(str(handle.duckdb_path), read_only=True) | |
| rows = con.execute("SELECT title, author FROM books_database_books ORDER BY id").fetchall() | |
| con.close() | |
| assert rows == [("Dune", "Herbert"), ("It", "King")] | |
| def test_cross_db_join_pg_mongo_sqlite(tmp_path): | |
| pytest.importorskip("pgserver", reason="needs pgserver (connectors extra)") | |
| sql = tmp_path / "books.sql" | |
| sql.write_text( | |
| "CREATE TABLE books (id integer, title text);\n" | |
| "COPY books (id, title) FROM stdin;\n1\tDune\n2\tIt\n\\.\n" | |
| ) | |
| dump = tmp_path / "rev_dump" | |
| _mongo_dump(dump, "reviews", [ | |
| {"_id": bson.ObjectId(), "book_id": 1, "stars": 5}, | |
| {"_id": bson.ObjectId(), "book_id": 2, "stars": 3}, | |
| ]) | |
| sq = tmp_path / "meta.db" | |
| c = sqlite3.connect(sq) | |
| c.execute("CREATE TABLE genre (book_id int, name text)") | |
| c.executemany("INSERT INTO genre VALUES (?,?)", [(1, "scifi"), (2, "horror")]) | |
| c.commit(); c.close() | |
| handle = load_attached_handle([ | |
| AttachSpec(alias="b", db_type="postgres", sql_file=sql), | |
| AttachSpec(alias="r", db_type="mongo", dump_folder=dump), | |
| AttachSpec(alias="m", db_type="sqlite", path=sq), | |
| ], dataset_id="t:cross") | |
| con = duckdb.connect(str(handle.duckdb_path), read_only=True) | |
| rows = con.execute( | |
| "SELECT b.title, r.stars, m.name FROM b_books b " | |
| "JOIN r_reviews r ON r.book_id = b.id " | |
| "JOIN m_genre m ON m.book_id = b.id ORDER BY b.id" | |
| ).fetchall() | |
| con.close() | |
| assert rows == [("Dune", 5, "scifi"), ("It", 3, "horror")] | |