refactor: unify storage protocols into single FunctionStorage interface

Merged InitStorage and WorkerStateStorage protocols into a unified
FunctionStorage protocol with consistent naming:

- global_put/get/delete/exists for init data
- worker_put/collect for worker state
- queue_push/pop/clear for work queue

Merged InitStorageSqlite and WorkerStateStorageSqlite into single
FunctionStorageSqlite implementation.

Users now implement one class for custom storage backends.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: rustyconover

.beads/issues.jsonl

@@ -1,3 +1,4 @@
 {"id":"vgi-python-e37","title":"move Invocation from function.py out to own file","description":"The Invocation clas is kind of seperate from functions, so it should be in its own file.  Move it and all of its other associated classes like InvocationType to its own file","status":"closed","priority":2,"issue_type":"task","created_at":"2026-01-04T09:18:46.605941-05:00","created_by":"rusty","updated_at":"2026-01-04T09:24:37.922675-05:00","closed_at":"2026-01-04T09:24:37.922675-05:00","close_reason":"Closed"}
 {"id":"vgi-python-odi","title":"Change max_processes from method to property in Function hierarchy","description":"Refactor max_processes from a method to a property across the Function class hierarchy (Function, ScalarFunction, TableFunctionGenerator, TableInOutFunction, etc.). This makes the API more consistent since max_processes is effectively a constant per function class and properties are more idiomatic for such values.","status":"closed","priority":2,"issue_type":"task","created_at":"2026-01-04T11:25:29.750648-05:00","created_by":"rusty","updated_at":"2026-01-04T11:50:57.566545-05:00","closed_at":"2026-01-04T11:50:57.566545-05:00","close_reason":"Closed"}
 {"id":"vgi-python-p91","title":"Move exception classes from function.py to own file","description":"Move InitIdentifierError and SchemaValidationError from vgi/function.py to a new vgi/exceptions.py file. Update imports in function.py and any other files that reference these exceptions.","status":"closed","priority":2,"issue_type":"task","created_at":"2026-01-04T09:12:28.058227-05:00","created_by":"rusty","updated_at":"2026-01-04T09:17:52.477661-05:00","closed_at":"2026-01-04T09:17:52.477661-05:00","close_reason":"Closed"}
+{"id":"vgi-python-zf7","title":"Unify storage protocols into single FunctionStorage interface","description":"## Problem\n\nCurrently there are two separate storage protocols (InitStorage and WorkerStateStorage) with inconsistent naming:\n- `create` vs `store` (different verbs for similar operations)\n- `collect_and_delete` is verbose and describes implementation\n\nUsers wanting a custom storage backend (Redis, DynamoDB, etc.) must implement two separate classes.\n\n## Solution\n\nUnify into a single `FunctionStorage` protocol with consistent naming using prefixes to group related methods:\n\n```python\nclass FunctionStorage(Protocol):\n    \"\"\"Storage protocol for VGI distributed function execution.\n    \n    Three access patterns:\n    - Global state: Init data shared across all workers (key-value with auto-generated keys)\n    - Worker state: Partial results per worker (collected during finalization)\n    - Work queue: Atomic work distribution across workers (FIFO queue)\n    \"\"\"\n    \n    # Global state (init data)\n    def global_put(self, value: bytes) -\u003e bytes: ...  # Returns auto-generated key\n    def global_get(self, key: bytes) -\u003e bytes: ...\n    def global_delete(self, key: bytes) -\u003e None: ...\n    def global_exists(self, key: bytes) -\u003e bool: ...\n    \n    # Worker state (partial results per worker)\n    def worker_put(self, invocation_id: bytes, worker_id: int, state: bytes) -\u003e None: ...\n    def worker_collect(self, invocation_id: bytes) -\u003e list[bytes]: ...  # Atomic collect+delete\n    \n    # Work queue (distributed work items)\n    def queue_push(self, invocation_id: bytes, items: list[bytes]) -\u003e int: ...\n    def queue_pop(self, invocation_id: bytes) -\u003e bytes | None: ...  # Atomic claim\n    def queue_clear(self, invocation_id: bytes) -\u003e int: ...\n```\n\n## Design Rationale\n\n- **Prefixes** (`global_`, `worker_`, `queue_`): Clear grouping, good autocomplete\n- **Consistent verbs**: `put/get` for storage, `push/pop` for queue\n- **Minimal interface**: 9 methods total (down from 9, but now unified)\n- **Single class variable** in Function: `storage: ClassVar[FunctionStorage]`\n\n## Files to Change\n\n- `vgi/function_storage.py`: New protocol + merged FunctionStorageSqlite\n- `vgi/function.py`: Single `storage` class variable, update all method calls","status":"closed","priority":2,"issue_type":"task","created_at":"2026-01-04T12:34:25.966005-05:00","created_by":"rusty","updated_at":"2026-01-04T12:58:16.913278-05:00","closed_at":"2026-01-04T12:58:16.913278-05:00","close_reason":"Closed"}

vgi/examples/table.py

@@ -473,11 +473,11 @@ def cardinality(self) -> TableCardinality:
         """
         return TableCardinality(estimate=self.count, max=self.count)
 
-    def perform_init(self, init_input: pa.RecordBatch) -> InitResult:
+    def initialize_global_state(self, init_input: pa.RecordBatch) -> InitResult:
         """Populate the work queue with range chunks."""
-        # Parse init data and store in init_storage
-        self.init_data = TableFunctionInitInput.deserialize(init_input)
-        self.init_identifier = self.init_storage.create(self.init_data.serialize())
+        # Parse init data and store in storage
+        self.init_input = TableFunctionInitInput.deserialize(init_input)
+        self.execution_identifier = self.storage.global_put(self.init_input.serialize())
 
         # Create work items for each chunk of the range
         work_items: list[bytes] = []
@@ -489,7 +489,7 @@ def perform_init(self, init_input: pa.RecordBatch) -> InitResult:
         if work_items:
             self.enqueue_work(work_items)
 
-        return InitResult(self.init_identifier)
+        return InitResult(self.execution_identifier)
 
     def process(self) -> OutputGenerator:
         """Generate values by pulling chunks from the work queue."""
@@ -579,7 +579,7 @@ class Meta:
 
     @property
     def output_schema(self) -> pa.Schema:
-        """Return the projected schema based on init_data."""
+        """Return the projected schema based on init_input."""
         return self.apply_projection(self.FULL_SCHEMA)
 
     def cardinality(self) -> TableCardinality:
@@ -591,8 +591,8 @@ def _get_projected_column_indices(self) -> list[int]:
 
         Returns indices from projection_ids if set, otherwise all columns.
         """
-        if self.init_data and self.init_data.projection_ids is not None:
-            return self.init_data.projection_ids
+        if self.init_input and self.init_input.projection_ids is not None:
+            return self.init_input.projection_ids
         return list(range(len(self.FULL_SCHEMA)))
 
     def process(self) -> OutputGenerator:

vgi/function.py

@@ -42,7 +42,7 @@
 
 import vgi.ipc_utils
 from vgi.exceptions import ExecutionIdentifierError, SchemaValidationError
-from vgi.function_storage import SqliteInitStorage, SqliteWorkerStateStorage
+from vgi.function_storage import FunctionStorage, FunctionStorageSqlite
 from vgi.invocation import InitResult, Invocation
 from vgi.metadata import MetadataMixin, ResolvedMetadata
 
@@ -262,10 +262,7 @@ class Meta:
 
     """
 
-    global_state_storage: ClassVar[SqliteInitStorage] = SqliteInitStorage()
-    worker_state_storage: ClassVar[SqliteWorkerStateStorage] = (
-        SqliteWorkerStateStorage()
-    )
+    storage: ClassVar[FunctionStorage] = FunctionStorageSqlite()
 
     # Cache for resolved metadata
     _metadata_cache: ClassVar[ResolvedMetadata | None] = None
@@ -370,7 +367,7 @@ def process(self, batch: pa.RecordBatch) -> OutputGenerator:
                 "Cannot store state: execution_identifier is not set. "
                 "Call initialize_global_state() or load_global_state() first."
             )
-        self.worker_state_storage.store(
+        self.storage.worker_put(
             self.execution_identifier,
             os.getpid(),
             state.serialize(),
@@ -406,9 +403,7 @@ def finalize(self) -> OutputGenerator:
                 "Cannot collect states: execution_identifier is not set. "
                 "Call initialize_global_state() or load_global_state() first."
             )
-        state_bytes_list = self.worker_state_storage.collect_and_delete(
-            self.execution_identifier
-        )
+        state_bytes_list = self.storage.worker_collect(self.execution_identifier)
         return [state_class.deserialize(data) for data in state_bytes_list]
 
     def enqueue_work(self, work_items: list[bytes]) -> int:
@@ -442,9 +437,7 @@ def initialize_global_state(self, init_input: pa.RecordBatch) -> InitResult:
                 "Cannot enqueue work: execution_identifier is not set. "
                 "Call enqueue_work() after initialize_global_state() has completed."
             )
-        return self.worker_state_storage.enqueue_work(
-            self.execution_identifier, work_items
-        )
+        return self.storage.queue_push(self.execution_identifier, work_items)
 
     def dequeue_work(self) -> bytes | None:
         """Claim and return the next work item from the queue.
@@ -477,7 +470,7 @@ def process(self) -> OutputGenerator:
                 "Cannot dequeue work: execution_identifier is not set. "
                 "Call initialize_global_state() or load_global_state() first."
             )
-        return self.worker_state_storage.dequeue_work(self.execution_identifier)
+        return self.storage.queue_pop(self.execution_identifier)
 
     @final
     @cached_property
@@ -534,9 +527,7 @@ def initialize_global_state(self, init_input: pa.RecordBatch) -> InitResult:
         """Perform a new init call and store it in the storage."""
         self.init_input = self.InitInputType.deserialize(init_input)
         assert self.init_input is not None
-        self.execution_identifier = self.global_state_storage.create(
-            self.init_input.serialize()
-        )
+        self.execution_identifier = self.storage.global_put(self.init_input.serialize())
         return InitResult(self.execution_identifier)
 
     def load_global_state(self, init_input: InitResult) -> None:
@@ -548,7 +539,7 @@ def load_global_state(self, init_input: InitResult) -> None:
             )
         self.execution_identifier = init_input.global_execution_identifier
         self.init_input = self.InitInputType.deserialize_bytes(
-            self.global_state_storage.get(self.execution_identifier)
+            self.storage.global_get(self.execution_identifier)
         )
 
     def setup(self) -> None: