feat: add write relation tests (#106)

substrait_consumer/functional/queries/sql/relations/write_relations.py

@@ -0,0 +1,28 @@
+from substrait_consumer.producers.duckdb_producer import DuckDBProducer
+from substrait_consumer.producers.datafusion_producer import DataFusionProducer
+from substrait_consumer.producers.isthmus_producer import IsthmusProducer
+
+WRITE_RELATIONS = {
+    "insert": (
+        """
+        INSERT INTO '{}' (r_regionkey, r_name, r_comment)
+        VALUES (99999, 'region_name', 'region comment');
+        """,
+        [DuckDBProducer, DataFusionProducer, IsthmusProducer],
+    ),
+    "update": (
+        """
+        UPDATE '{}'
+        SET c_address = 'Substait Avenue', c_phone = '123-456-7890'
+        WHERE c_custkey = 1;
+        """,
+        [DuckDBProducer, DataFusionProducer, IsthmusProducer],
+    ),
+    "delete": (
+        """
+        DELETE FROM '{}'
+        WHERE c_custkey = 1;
+        """,
+        [DuckDBProducer, DataFusionProducer, IsthmusProducer],
+    ),
+}

substrait_consumer/functional/write_relation_configs.py

@@ -0,0 +1,23 @@
+from substrait_consumer.functional.queries.sql.relations.write_relations import (
+    WRITE_RELATIONS)
+
+WRITE_RELATION_TESTS = (
+    {
+        "test_name": "insert",
+        "file_names": ["region.parquet"],
+        "sql_query": WRITE_RELATIONS["insert"],
+        "ibis_expr": None
+    },
+    {
+        "test_name": "update",
+        "file_names": ["customer.parquet"],
+        "sql_query": WRITE_RELATIONS["update"],
+        "ibis_expr": None
+    },
+    {
+        "test_name": "delete",
+        "file_names": ["customer.parquet"],
+        "sql_query": WRITE_RELATIONS["delete"],
+        "ibis_expr": None
+    },
+)

substrait_consumer/tests/functional/relations/test_write_relation.py

@@ -0,0 +1,103 @@
+from typing import Callable, Iterable
+
+import duckdb
+from ibis.expr.types.relations import Table
+from ibis_substrait.tests.compiler.conftest import *
+
+from substrait_consumer.functional.write_relation_configs import (
+    WRITE_RELATION_TESTS)
+from substrait_consumer.functional.common import (
+    generate_snapshot_results,
+    substrait_consumer_sql_test, substrait_producer_sql_test)
+from substrait_consumer.parametrization import custom_parametrization
+
+
+@pytest.fixture
+def mark_producer_tests_as_xfail(request):
+    """Marks a subset of tests as expected to be fail."""
+    test_case_name = request.node.callspec.id.split('-')[-1]
+    if test_case_name in ["insert", "update", "delete"]:
+        pytest.skip(reason='Creating substrait plans with write relations is not supported')
+
+
+@pytest.fixture
+def mark_consumer_tests_as_xfail(request):
+    """Marks a subset of tests as expected to be fail."""
+    test_case_name = request.node.callspec.id.split('-')[-1]
+    if test_case_name in ["insert", "update", "delete"]:
+        pytest.skip(reason='Creating substrait plans with write relations is not supported')
+
+
+@pytest.mark.usefixtures("prepare_tpch_parquet_data")
+class TestWriteRelation:
+    """
+    Test Class verifying different consumers are able to run substrait plans
+    that include substrait write relations.
+    """
+
+    @staticmethod
+    @pytest.fixture(scope="class", autouse=True)
+    def setup_teardown_class(request):
+        cls = request.cls
+
+        cls.db_connection = duckdb.connect()
+        cls.db_connection.execute("install substrait")
+        cls.db_connection.execute("load substrait")
+        cls.created_tables = set()
+
+        yield
+
+        cls.db_connection.close()
+
+    @custom_parametrization(WRITE_RELATION_TESTS)
+    @pytest.mark.produce_substrait_snapshot
+    @pytest.mark.usefixtures('mark_producer_tests_as_xfail')
+    def test_producer_write_relations(
+        self,
+        snapshot,
+        test_name: str,
+        file_names: Iterable[str],
+        sql_query: tuple,
+        ibis_expr: Callable[[Table], Table],
+        producer,
+        partsupp
+    ) -> None:
+        test_name = f"write_relation_snapshots:{test_name}"
+        substrait_producer_sql_test(
+            test_name,
+            snapshot,
+            self.db_connection,
+            self.created_tables,
+            file_names,
+            sql_query,
+            ibis_expr,
+            producer,
+            partsupp,
+            validate=True
+        )
+
+    @custom_parametrization(WRITE_RELATION_TESTS)
+    @pytest.mark.consume_substrait_snapshot
+    @pytest.mark.usefixtures('mark_consumer_tests_as_xfail')
+    def test_consumer_write_relations(
+        self,
+        snapshot,
+        test_name: str,
+        file_names: Iterable[str],
+        sql_query: tuple,
+        ibis_expr: Callable[[Table], Table],
+        producer,
+        consumer,
+    ) -> None:
+        test_name = f"write_relation_snapshots:{test_name}"
+        substrait_consumer_sql_test(
+            test_name,
+            snapshot,
+            self.db_connection,
+            self.created_tables,
+            file_names,
+            sql_query,
+            ibis_expr,
+            producer,
+            consumer,
+        )