External storage: Test oversized query result storage

features/features.go

@@ -33,6 +33,7 @@ import (
 	eager_activity_non_remote_activities_worker "github.com/temporalio/features/features/eager_activity/non_remote_activities_worker"
 	eager_workflow_successful_start "github.com/temporalio/features/features/eager_workflow/successful_start"
 	nexus_sync_success "github.com/temporalio/features/features/nexus/sync_success"
+	query_oversized_result_external_storage "github.com/temporalio/features/features/query/oversized_result_external_storage"
 	query_successful_query "github.com/temporalio/features/features/query/successful_query"
 	query_timeout_due_to_no_active_workers "github.com/temporalio/features/features/query/timeout_due_to_no_active_workers"
 	query_unexpected_arguments "github.com/temporalio/features/features/query/unexpected_arguments"
@@ -96,6 +97,7 @@ func init() {
 		eager_activity_non_remote_activities_worker.Feature,
 		eager_workflow_successful_start.Feature,
 		nexus_sync_success.Feature,
+		query_oversized_result_external_storage.Feature,
 		query_successful_query.Feature,
 		query_timeout_due_to_no_active_workers.Feature,
 		query_unexpected_arguments.Feature,

features/query/oversized_result_external_storage/README.md

@@ -0,0 +1,9 @@
+# Oversized query result with external storage
+
+An oversized query result succeeds when external storage is configured. The
+SDK must offload the result before applying the server payload-size limit.
+
+The query returns 3 MiB, above the server's default 2 MiB error limit. The
+checker verifies the value and confirms that storage and retrieval occurred.
+
+Go, Python, and TypeScript currently expose functional external storage.

features/query/oversized_result_external_storage/config.json

@@ -0,0 +1,11 @@
+{
+  "go": {
+    "minVersion": "1.48.0"
+  },
+  "py": {
+    "minVersion": "1.24.0"
+  },
+  "ts": {
+    "minVersion": "1.21.0"
+  }
+}

features/query/oversized_result_external_storage/feature.go

@@ -0,0 +1,142 @@
+package oversized_result_external_storage
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+
+	"github.com/temporalio/features/harness/go/harness"
+	commonpb "go.temporal.io/api/common/v1"
+	"go.temporal.io/sdk/client"
+	"go.temporal.io/sdk/converter"
+	"go.temporal.io/sdk/workflow"
+	"google.golang.org/protobuf/proto"
+)
+
+const (
+	queryName    = "oversized-result"
+	finishSignal = "finish"
+	driverName   = "query-result-memory"
+	// Exceed the server limit so the SDK must offload the result.
+	resultSize       = 3 * 1024 * 1024
+	storageThreshold = 1024
+)
+
+var storage = newMemoryDriver()
+
+var Feature = harness.Feature{
+	Workflows: Workflow,
+	ClientOptions: client.Options{
+		ExternalStorage: converter.ExternalStorage{
+			Drivers:              []converter.StorageDriver{storage},
+			PayloadSizeThreshold: storageThreshold,
+		},
+	},
+	CheckResult: checkResult,
+}
+
+func Workflow(ctx workflow.Context) error {
+	result := strings.Repeat("a", resultSize)
+	if err := workflow.SetQueryHandler(ctx, queryName, func() (string, error) {
+		return result, nil
+	}); err != nil {
+		return err
+	}
+
+	workflow.GetSignalChannel(ctx, finishSignal).Receive(ctx, nil)
+	return nil
+}
+
+func checkResult(ctx context.Context, r *harness.Runner, run client.WorkflowRun) error {
+	value, err := r.Client.QueryWorkflow(ctx, run.GetID(), run.GetRunID(), queryName)
+	if err != nil {
+		return err
+	}
+
+	var result string
+	if err := value.Get(&result); err != nil {
+		return err
+	}
+	if result != strings.Repeat("a", resultSize) {
+		return fmt.Errorf("unexpected query result")
+	}
+	if storage.storeCalls() == 0 || storage.retrieveCalls() == 0 {
+		return fmt.Errorf("query result did not use external storage")
+	}
+
+	if err := r.Client.SignalWorkflow(ctx, run.GetID(), run.GetRunID(), finishSignal, nil); err != nil {
+		return err
+	}
+	return r.CheckResultDefault(ctx, run)
+}
+
+type memoryDriver struct {
+	mu        sync.Mutex
+	payloads  map[string]*commonpb.Payload
+	nextID    int
+	stores    int
+	retrieves int
+}
+
+func newMemoryDriver() *memoryDriver {
+	return &memoryDriver{payloads: make(map[string]*commonpb.Payload)}
+}
+
+func (*memoryDriver) Name() string {
+	return driverName
+}
+
+func (*memoryDriver) Type() string {
+	return driverName
+}
+
+func (d *memoryDriver) Store(
+	_ converter.StorageDriverStoreContext,
+	payloads []*commonpb.Payload,
+) ([]converter.StorageDriverClaim, error) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.stores++
+	claims := make([]converter.StorageDriverClaim, len(payloads))
+	for i, payload := range payloads {
+		key := fmt.Sprintf("payload-%d", d.nextID)
+		d.nextID++
+		d.payloads[key] = proto.Clone(payload).(*commonpb.Payload)
+		claims[i] = converter.StorageDriverClaim{ClaimData: map[string]string{"key": key}}
+	}
+	return claims, nil
+}
+
+func (d *memoryDriver) Retrieve(
+	_ converter.StorageDriverRetrieveContext,
+	claims []converter.StorageDriverClaim,
+) ([]*commonpb.Payload, error) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.retrieves++
+	payloads := make([]*commonpb.Payload, len(claims))
+	for i, claim := range claims {
+		key := claim.ClaimData["key"]
+		payload, ok := d.payloads[key]
+		if !ok {
+			return nil, fmt.Errorf("payload %q not found", key)
+		}
+		payloads[i] = proto.Clone(payload).(*commonpb.Payload)
+	}
+	return payloads, nil
+}
+
+func (d *memoryDriver) storeCalls() int {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	return d.stores
+}
+
+func (d *memoryDriver) retrieveCalls() int {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+	return d.retrieves
+}

features/query/oversized_result_external_storage/feature.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+
+import temporalio.converter
+from temporalio import workflow
+from temporalio.api.common.v1 import Payload
+from temporalio.client import WorkflowHandle
+from temporalio.converter import (
+    ExternalStorage,
+    StorageDriver,
+    StorageDriverClaim,
+    StorageDriverRetrieveContext,
+    StorageDriverStoreContext,
+)
+
+from harness.python.feature import Runner, register_feature
+
+# Exceed the server limit so the SDK must offload the result.
+RESULT_SIZE = 3 * 1024 * 1024
+STORAGE_THRESHOLD = 1024
+
+
+class MemoryDriver(StorageDriver):
+    def __init__(self) -> None:
+        self.payloads: dict[str, bytes] = {}
+        self.stores = 0
+        self.retrieves = 0
+
+    def name(self) -> str:
+        return "query-result-memory"
+
+    async def store(
+        self,
+        context: StorageDriverStoreContext,
+        payloads: Sequence[Payload],
+    ) -> list[StorageDriverClaim]:
+        self.stores += 1
+        claims = []
+        for payload in payloads:
+            key = f"payload-{len(self.payloads)}"
+            self.payloads[key] = payload.SerializeToString()
+            claims.append(StorageDriverClaim(claim_data={"key": key}))
+        return claims
+
+    async def retrieve(
+        self,
+        context: StorageDriverRetrieveContext,
+        claims: Sequence[StorageDriverClaim],
+    ) -> list[Payload]:
+        self.retrieves += 1
+        payloads = []
+        for claim in claims:
+            payload = Payload()
+            payload.ParseFromString(self.payloads[claim.claim_data["key"]])
+            payloads.append(payload)
+        return payloads
+
+
+driver = MemoryDriver()
+data_converter = dataclasses.replace(
+    temporalio.converter.default(),
+    external_storage=ExternalStorage(
+        drivers=[driver],
+        payload_size_threshold=STORAGE_THRESHOLD,
+    ),
+)
+
+
+@workflow.defn
+class Workflow:
+    def __init__(self) -> None:
+        self.finished = False
+
+    @workflow.run
+    async def run(self) -> None:
+        await workflow.wait_condition(lambda: self.finished)
+
+    @workflow.query
+    def oversized_result(self) -> str:
+        return "a" * RESULT_SIZE
+
+    @workflow.signal
+    def finish(self) -> None:
+        self.finished = True
+
+
+async def check_result(_: Runner, handle: WorkflowHandle) -> None:
+    result = await handle.query(Workflow.oversized_result)
+    assert result == "a" * RESULT_SIZE
+    assert driver.stores > 0
+    assert driver.retrieves > 0
+
+    await handle.signal(Workflow.finish)
+    await handle.result()
+
+
+register_feature(
+    workflows=[Workflow],
+    check_result=check_result,
+    data_converter=data_converter,
+)

features/query/oversized_result_external_storage/feature.ts

@@ -0,0 +1,63 @@
+import * as assert from 'assert';
+import { ExternalStorage, type Payload, type StorageDriver, StorageDriverClaim } from '@temporalio/common';
+import { Feature } from '@temporalio/harness';
+import * as wf from '@temporalio/workflow';
+
+// Exceed the server limit so the SDK must offload the result.
+const RESULT_SIZE = 3 * 1024 * 1024;
+const STORAGE_THRESHOLD = 1024;
+
+const query = wf.defineQuery<string>('oversized-result');
+const finishSignal = wf.defineSignal('finish');
+
+const payloads = new Map<string, Payload>();
+let stores = 0;
+let retrieves = 0;
+
+const driver: StorageDriver = {
+  name: 'query-result-memory',
+  type: 'query-result-memory',
+  async store(_context, values) {
+    stores++;
+    return values.map((value) => {
+      const key = `payload-${payloads.size}`;
+      payloads.set(key, value);
+      return new StorageDriverClaim({ key });
+    });
+  },
+  async retrieve(_context, claims) {
+    retrieves++;
+    return claims.map((claim) => {
+      const key = claim.claimData.key;
+      const value = payloads.get(key);
+      if (value === undefined) {
+        throw new Error(`Payload ${key} not found`);
+      }
+      return value;
+    });
+  },
+};
+
+const externalStorage = new ExternalStorage({
+  drivers: [driver],
+  payloadSizeThreshold: STORAGE_THRESHOLD,
+});
+
+export async function workflow(): Promise<void> {
+  wf.setHandler(query, () => 'a'.repeat(RESULT_SIZE));
+  await new Promise<void>((resolve) => wf.setHandler(finishSignal, resolve));
+}
+
+export const feature = new Feature({
+  workflow,
+  dataConverter: { externalStorage },
+  checkResult: async (runner, handle) => {
+    const result = await handle.query(query);
+    assert.equal(result, 'a'.repeat(RESULT_SIZE));
+    assert.ok(stores > 0);
+    assert.ok(retrieves > 0);
+
+    await handle.signal(finishSignal);
+    await runner.waitForRunResult(handle);
+  },
+});