feat: add options support to Store for format-specific Spark DataFrameWriter options

Author: james-willis

README.md

@@ -99,6 +99,27 @@ The `Store` class supports the following options:
 Use `Store.for_download()` as a convenient shorthand for storing results
 as a single Parquet file with a presigned URL.
 
+#### Store options
+
+You can pass format-specific Spark write options through the `options`
+parameter. These correspond to the options available in Spark's
+`DataFrameWriter` and are applied after the server's default options,
+allowing you to override them.
+
+```python
+# CSV without headers and a custom delimiter
+store = Store.for_download(
+    format=StorageFormat.CSV,
+    options={"header": "false", "delimiter": "|"},
+)
+
+# GeoJSON preserving null fields
+store = Store.for_download(
+    format=StorageFormat.GEOJSON,
+    options={"ignoreNullFields": "false"},
+)
+```
+
 ### Execution progress
 
 You can monitor the progress of running queries by registering a

tests/test_store_options.py

@@ -0,0 +1,147 @@
+"""Tests for Store options support."""
+
+import json
+
+from wherobots.db.models import Store
+from wherobots.db.types import StorageFormat
+
+
+class TestStoreOptions:
+    """Tests for the options field on Store."""
+
+    def test_default_options_is_none(self):
+        store = Store(format=StorageFormat.PARQUET)
+        assert store.options is None
+
+    def test_options_set(self):
+        store = Store(
+            format=StorageFormat.CSV,
+            options={"header": "false", "delimiter": "|"},
+        )
+        assert store.options == {"header": "false", "delimiter": "|"}
+
+    def test_empty_options_normalized_to_none(self):
+        store = Store(format=StorageFormat.PARQUET, options={})
+        assert store.options is None
+
+    def test_none_options_stays_none(self):
+        store = Store(format=StorageFormat.PARQUET, options=None)
+        assert store.options is None
+
+    def test_options_defensively_copied(self):
+        original = {"header": "false"}
+        store = Store(format=StorageFormat.CSV, options=original)
+        # Mutating the original should not affect the store
+        original["delimiter"] = "|"
+        assert "delimiter" not in store.options
+
+    def test_options_dict_is_mutable(self):
+        """Store is not frozen, so options dict can be mutated after construction."""
+        store = Store(format=StorageFormat.CSV, options={"header": "false"})
+        store.options["delimiter"] = "|"
+        assert store.options == {"header": "false", "delimiter": "|"}
+
+
+class TestStoreForDownloadWithOptions:
+    """Tests for Store.for_download() with options parameter."""
+
+    def test_for_download_default_no_options(self):
+        store = Store.for_download()
+        assert store.options is None
+
+    def test_for_download_with_options(self):
+        store = Store.for_download(options={"header": "false"})
+        assert store.options == {"header": "false"}
+        assert store.single is True
+        assert store.generate_presigned_url is True
+
+    def test_for_download_with_format_and_options(self):
+        store = Store.for_download(
+            format=StorageFormat.CSV,
+            options={"header": "false", "delimiter": "|"},
+        )
+        assert store.format == StorageFormat.CSV
+        assert store.options == {"header": "false", "delimiter": "|"}
+
+    def test_for_download_empty_options_normalized(self):
+        store = Store.for_download(options={})
+        assert store.options is None
+
+
+class TestStoreSerializationWithOptions:
+    """Tests for store dict serialization matching connection.py's format."""
+
+    def _serialize_store(self, store: Store) -> dict:
+        """Replicate the serialization logic from Connection.__execute_sql."""
+        store_dict = {
+            "format": store.format.value,
+            "single": str(store.single).lower(),
+            "generate_presigned_url": str(store.generate_presigned_url).lower(),
+        }
+        if store.options:
+            store_dict["options"] = store.options
+        return store_dict
+
+    def test_serialize_without_options(self):
+        store = Store.for_download(format=StorageFormat.GEOJSON)
+        d = self._serialize_store(store)
+        assert d == {
+            "format": "geojson",
+            "single": "true",
+            "generate_presigned_url": "true",
+        }
+        assert "options" not in d
+
+    def test_serialize_with_options(self):
+        store = Store.for_download(
+            format=StorageFormat.CSV,
+            options={"header": "false", "delimiter": "|"},
+        )
+        d = self._serialize_store(store)
+        assert d == {
+            "format": "csv",
+            "single": "true",
+            "generate_presigned_url": "true",
+            "options": {"header": "false", "delimiter": "|"},
+        }
+
+    def test_serialize_empty_options_omitted(self):
+        store = Store(format=StorageFormat.PARQUET, options={})
+        d = self._serialize_store(store)
+        assert "options" not in d
+
+    def test_json_roundtrip_with_options(self):
+        store = Store.for_download(
+            format=StorageFormat.GEOJSON,
+            options={"ignoreNullFields": "false"},
+        )
+        d = self._serialize_store(store)
+        payload = json.dumps(d)
+        parsed = json.loads(payload)
+        assert parsed["options"] == {"ignoreNullFields": "false"}
+
+    def test_full_request_shape(self):
+        """Verify the full execute_sql request dict shape with store options."""
+        store = Store.for_download(
+            format=StorageFormat.CSV,
+            options={"header": "false"},
+        )
+        request = {
+            "kind": "execute_sql",
+            "execution_id": "test-id",
+            "statement": "SELECT 1",
+        }
+        store_dict = self._serialize_store(store)
+        request["store"] = store_dict
+
+        assert request == {
+            "kind": "execute_sql",
+            "execution_id": "test-id",
+            "statement": "SELECT 1",
+            "store": {
+                "format": "csv",
+                "single": "true",
+                "generate_presigned_url": "true",
+                "options": {"header": "false"},
+            },
+        }

wherobots/db/connection.py

@@ -279,11 +279,14 @@ def __execute_sql(
             request["enable_progress_events"] = True
 
         if store:
-            request["store"] = {
+            store_dict: dict[str, Any] = {
                 "format": store.format.value,
                 "single": str(store.single).lower(),
                 "generate_presigned_url": str(store.generate_presigned_url).lower(),
             }
+            if store.options:
+                store_dict["options"] = store.options
+            request["store"] = store_dict
 
         self.__queries[execution_id] = Query(
             sql=sql,

wherobots/db/models.py

@@ -31,25 +31,40 @@ class Store:
         single: If True, store as a single file. If False, store as multiple files.
         generate_presigned_url: If True, generate a presigned URL for the result.
             Requires single=True.
+        options: Optional dict of format-specific Spark DataFrameWriter options
+            (e.g. ``{"header": "false", "delimiter": "|"}`` for CSV). These are
+            applied after the server's default options, so they can override them.
+            An empty dict is normalized to None.
     """
 
     format: StorageFormat
     single: bool = False
     generate_presigned_url: bool = False
+    options: dict[str, str] | None = None
 
     def __post_init__(self) -> None:
         if self.generate_presigned_url and not self.single:
             raise ValueError("Presigned URL can only be generated when single=True")
+        # Normalize empty options to None and defensively copy.
+        if self.options:
+            self.options = dict(self.options)
+        else:
+            self.options = None
 
     @classmethod
-    def for_download(cls, format: StorageFormat | None = None) -> "Store":
+    def for_download(
+        cls,
+        format: StorageFormat | None = None,
+        options: dict[str, str] | None = None,
+    ) -> "Store":
         """Create a configuration for downloading results via a presigned URL.
 
         This is a convenience method that creates a configuration with
         single file mode and presigned URL generation enabled.
 
         Args:
             format: The storage format.
+            options: Optional format-specific Spark DataFrameWriter options.
 
         Returns:
             A Store configured for single-file download with presigned URL.
@@ -58,6 +73,7 @@ def for_download(cls, format: StorageFormat | None = None) -> "Store":
             format=format or DEFAULT_STORAGE_FORMAT,
             single=True,
             generate_presigned_url=True,
+            options=options,
         )