From f19a98dc0f75aade309ee138ac80b30889e1922c Mon Sep 17 00:00:00 2001 From: Dave Date: Tue, 26 Nov 2024 11:30:26 +0100 Subject: [PATCH] add dataset type selection test and fix bug in tests --- dlt/destinations/dataset/dataset.py | 1 + tests/load/test_read_interfaces.py | 22 ++++++++++++++++++++-- 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/dlt/destinations/dataset/dataset.py b/dlt/destinations/dataset/dataset.py index 0e88da7f3a..0287196cbf 100644 --- a/dlt/destinations/dataset/dataset.py +++ b/dlt/destinations/dataset/dataset.py @@ -115,6 +115,7 @@ def __call__(self, query: Any) -> ReadableDBAPIRelation: def table(self, table_name: str) -> SupportsReadableRelation: # we can create an ibis powered relation if ibis is available + print(table_name in self.schema.tables, self._dataset_type) if table_name in self.schema.tables and self._dataset_type in ("auto", "ibis"): try: from dlt.common.libs.ibis import create_unbound_ibis_table diff --git a/tests/load/test_read_interfaces.py b/tests/load/test_read_interfaces.py index c51d820508..02966155b5 100644 --- a/tests/load/test_read_interfaces.py +++ b/tests/load/test_read_interfaces.py @@ -155,6 +155,24 @@ def double_items(): ) +@pytest.mark.no_load +@pytest.mark.essential +@pytest.mark.parametrize( + "populated_pipeline", + configs, + indirect=True, + ids=lambda x: x.name, +) +def test_explicit_dataset_type_selection(populated_pipeline: Pipeline): + from dlt.destinations.dataset.dataset import ReadableDBAPIRelation + from dlt.destinations.dataset.ibis_relation import ReadableIbisRelation + + assert isinstance( + populated_pipeline._dataset(dataset_type="default").items, ReadableDBAPIRelation + ) + assert isinstance(populated_pipeline._dataset(dataset_type="ibis").items, ReadableIbisRelation) + + @pytest.mark.no_load @pytest.mark.essential @pytest.mark.parametrize( @@ -272,7 +290,7 @@ def test_db_cursor_access(populated_pipeline: Pipeline) -> None: ) def test_hint_preservation(populated_pipeline: Pipeline) -> None: # NOTE: for now hints are only preserved for the default dataset - table_relationship = populated_pipeline._dataset("default").items + table_relationship = populated_pipeline._dataset(dataset_type="default").items # check that hints are carried over to arrow table expected_decimal_precision = 10 expected_decimal_precision_2 = 12 @@ -365,7 +383,7 @@ def test_limit_and_head(populated_pipeline: Pipeline) -> None: ids=lambda x: x.name, ) def test_column_selection(populated_pipeline: Pipeline) -> None: - table_relationship = populated_pipeline._dataset("default").items + table_relationship = populated_pipeline._dataset(dataset_type="default").items columns = ["_dlt_load_id", "other_decimal"] data_frame = table_relationship.select(*columns).head().df() assert [v.lower() for v in data_frame.columns.values] == columns