feat: add tests for ingestion into specified catalog/schema (#10)

Author: lidavidm

adbc_drivers_validation/model.py

@@ -88,6 +88,9 @@ class DriverFeatures:
     statement_prepare: bool = True
     _current_catalog: str | FromEnv | None = None
     _current_schema: str | FromEnv | None = None
+    _secondary_schema: str | FromEnv | None = None
+    _secondary_catalog: str | FromEnv | None = None
+    _secondary_catalog_schema: str | FromEnv | None = None
     supported_xdbc_fields: list[str] = dataclasses.field(default_factory=list)
 
     def __init__(
@@ -104,6 +107,9 @@ def __init__(
         statement_prepare=True,
         current_catalog=None,
         current_schema=None,
+        secondary_schema=None,
+        secondary_catalog=None,
+        secondary_catalog_schema=None,
         supported_xdbc_fields=None,
     ):
         self.connection_get_table_schema = connection_get_table_schema
@@ -117,6 +123,9 @@ def __init__(
         self.statement_prepare = statement_prepare
         self._current_catalog = current_catalog
         self._current_schema = current_schema
+        self._secondary_schema = secondary_schema
+        self._secondary_catalog = secondary_catalog
+        self._secondary_catalog_schema = secondary_catalog_schema
         self.supported_xdbc_fields = supported_xdbc_fields or []
 
     @property
@@ -131,6 +140,24 @@ def current_schema(self) -> str | None:
             return self._current_schema.get_or_raise()
         return self._current_schema
 
+    @property
+    def secondary_schema(self) -> str | None:
+        if isinstance(self._secondary_schema, FromEnv):
+            return self._secondary_schema.get_or_raise()
+        return self._secondary_schema
+
+    @property
+    def secondary_catalog(self) -> str | None:
+        if isinstance(self._secondary_catalog, FromEnv):
+            return self._secondary_catalog.get_or_raise()
+        return self._secondary_catalog
+
+    @property
+    def secondary_catalog_schema(self) -> str | None:
+        if isinstance(self._secondary_catalog_schema, FromEnv):
+            return self._secondary_catalog_schema.get_or_raise()
+        return self._secondary_catalog_schema
+
 
 class DriverQuirks(abc.ABC):
     @property
@@ -149,7 +176,13 @@ def bind_parameter(self, index: int) -> str:
         return "?"
 
     def drop_table(
-        self, *, table_name: str, if_exists: bool = True, temporary: bool = False
+        self,
+        *,
+        table_name: str,
+        schema_name: str | None = None,
+        catalog_name: str | None = None,
+        if_exists: bool = True,
+        temporary: bool = False,
     ) -> str:
         """
         Drop a table.
@@ -160,17 +193,27 @@ def drop_table(
             The cursor to execute the query.
         table_name : str
             The name of the table to drop.
+        schema_name : str, optional
+            The schema containing the table (if not given, assume current
+            schema).
+        catalog_name : str, optional
+            The catalog containing the table (if not given, assume current
+            catalog).
         if_exists : bool
             If True, do not raise an error if the table does not exist.
         temporary : bool
             If True, the table is a temporary table.
         """
         if temporary:
             raise NotImplementedError
+
+        name = ".".join(
+            [part for part in (catalog_name, schema_name, table_name) if part]
+        )
         if if_exists:
-            return f"DROP TABLE IF EXISTS {table_name}"
+            return f"DROP TABLE IF EXISTS {name}"
         else:
-            return f"DROP TABLE {table_name}"
+            return f"DROP TABLE {name}"
 
     @abc.abstractmethod
     def is_table_not_found(self, table_name: str, error: Exception) -> bool:

adbc_drivers_validation/tests/ingest.py

@@ -34,18 +34,17 @@ def generate_tests(quirks: model.DriverQuirks, metafunc) -> None:
     """Parameterize the tests in this module for the given driver."""
     combinations = []
 
-    if metafunc.definition.name == "test_not_null":
-        metafunc.parametrize(
-            "driver",
-            [pytest.param(quirks.name, id=quirks.name)],
-            scope="module",
-            indirect=["driver"],
-        )
-        return
-
-    if metafunc.definition.name == "test_temporary":
+    if (
+        enabled := {
+            "test_not_null": True,
+            "test_temporary": quirks.features.statement_bulk_ingest_temporary,
+            "test_schema": quirks.features.statement_bulk_ingest_schema,
+            "test_catalog": quirks.features.statement_bulk_ingest_catalog,
+        }.get(metafunc.definition.name, None)
+        is not None
+    ):
         marks = []
-        if not quirks.features.statement_bulk_ingest_temporary:
+        if not enabled:
             marks.append(pytest.mark.skip(reason="not implemented"))
 
         metafunc.parametrize(
@@ -391,3 +390,76 @@ def test_temporary(
 
         compare.compare_tables(data1, result_temporary)
         compare.compare_tables(data2, result_normal)
+
+    def test_schema(
+        self,
+        driver: model.DriverQuirks,
+        conn: adbc_driver_manager.dbapi.Connection,
+    ) -> None:
+        data = pyarrow.Table.from_pydict(
+            {
+                "idx": [1, 2, 3],
+                "value": ["foo", "bar", "baz"],
+            }
+        )
+        table_name = "test_ingest_schema"
+        schema_name = driver.features.secondary_schema
+        with conn.cursor() as cursor:
+            cursor.execute(
+                driver.drop_table(
+                    table_name=table_name,
+                    schema_name=schema_name,
+                )
+            )
+            cursor.adbc_ingest(
+                table_name,
+                data,
+                db_schema_name=schema_name,
+            )
+
+        select = f"SELECT idx, value FROM {schema_name}.{table_name} ORDER BY idx ASC"
+        with conn.cursor() as cursor:
+            cursor.adbc_statement.set_sql_query(select)
+            handle, _ = cursor.adbc_statement.execute_query()
+            with pyarrow.RecordBatchReader._import_from_c(handle.address) as reader:
+                result = reader.read_all()
+
+        compare.compare_tables(data, result)
+
+    def test_catalog(
+        self,
+        driver: model.DriverQuirks,
+        conn: adbc_driver_manager.dbapi.Connection,
+    ) -> None:
+        data = pyarrow.Table.from_pydict(
+            {
+                "idx": [1, 2, 3],
+                "value": ["foo", "bar", "baz"],
+            }
+        )
+        table_name = "test_ingest_catalog"
+        schema_name = driver.features.secondary_catalog_schema
+        catalog_name = driver.features.secondary_catalog
+        with conn.cursor() as cursor:
+            cursor.execute(
+                driver.drop_table(
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    catalog_name=catalog_name,
+                )
+            )
+            cursor.adbc_ingest(
+                table_name,
+                data,
+                db_schema_name=schema_name,
+                catalog_name=catalog_name,
+            )
+
+        select = f"SELECT idx, value FROM {catalog_name}.{schema_name}.{table_name} ORDER BY idx ASC"
+        with conn.cursor() as cursor:
+            cursor.adbc_statement.set_sql_query(select)
+            handle, _ = cursor.adbc_statement.execute_query()
+            with pyarrow.RecordBatchReader._import_from_c(handle.address) as reader:
+                result = reader.read_all()
+
+        compare.compare_tables(data, result)