fix(tests): update ScopedStorage tests for type safety

Author: sroussey

packages/knowledge-base/src/knowledge-base/SharedTableSchemas.ts

@@ -78,13 +78,14 @@ export type SharedChunkPrimaryKey = typeof SharedChunkPrimaryKey;
  * Index definitions for efficient KB-scoped queries on shared document table.
  */
 export const SharedDocumentIndexes = [["kb_id"]] as const satisfies readonly (
-  keyof any | readonly (keyof any)[]
+  | keyof any
+  | readonly (keyof any)[]
 )[];
 
 /**
  * Index definitions for efficient KB-scoped queries on shared chunk table.
  */
-export const SharedChunkIndexes = [
-  ["kb_id"],
-  ["kb_id", "doc_id"],
-] as const satisfies readonly (keyof any | readonly (keyof any)[])[]; 
+export const SharedChunkIndexes = [["kb_id"], ["kb_id", "doc_id"]] as const satisfies readonly (
+  | keyof any
+  | readonly (keyof any)[]
+)[];

packages/storage/src/tabular/BaseTabularStorage.ts

@@ -4,8 +4,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { DataPortSchemaObject, FromSchema, TypedArraySchemaOptions } from "@workglow/util/schema";
 import { createServiceToken, EventEmitter, makeFingerprint } from "@workglow/util";
+import { DataPortSchemaObject, FromSchema, TypedArraySchemaOptions } from "@workglow/util/schema";
 import {
   AnyTabularStorage,
   AutoGeneratedKeys,
@@ -70,7 +70,7 @@ export abstract class BaseTabularStorage<
   protected primaryKeySchema: DataPortSchemaObject;
   protected valueSchema: DataPortSchemaObject;
 
-  /** Name of the auto-generated key column (only first primary key column can be auto-generated) */
+  /** Name of the auto-generated key column (at most one primary key column may be auto-generated) */
   protected autoGeneratedKeyName: keyof Entity | null = null;
   /** Strategy for generating the auto-generated key */
   protected autoGeneratedKeyStrategy: KeyGenerationStrategy | null = null;
@@ -165,8 +165,8 @@ export abstract class BaseTabularStorage<
       }
     }
 
-    // Detect and validate auto-generated keys
-    // Only the first primary key column can be auto-generated
+    // Detect and validate auto-generated keys (at most one PK column; any PK position is allowed).
+    // Composite keys often put a scope column first (e.g. kb_id) and auto-generate a second id.
     const autoGeneratedKeys: string[] = [];
     for (const key of primaryKeyNames) {
       const keyStr = String(key);
@@ -181,20 +181,12 @@ export abstract class BaseTabularStorage<
     if (autoGeneratedKeys.length > 1) {
       throw new Error(
         `Multiple auto-generated keys detected: ${autoGeneratedKeys.join(", ")}. ` +
-          `Only the first primary key column can be auto-generated.`
+          `At most one primary key column can be auto-generated.`
       );
     }
 
     if (autoGeneratedKeys.length > 0) {
       const autoGenKeyName = autoGeneratedKeys[0];
-      const firstPrimaryKey = String(primaryKeyNames[0]);
-
-      if (autoGenKeyName !== firstPrimaryKey) {
-        throw new Error(
-          `Auto-generated key "${autoGenKeyName}" must be the first primary key column. ` +
-            `Current first primary key is "${firstPrimaryKey}".`
-        );
-      }
 
       this.autoGeneratedKeyName = autoGenKeyName as keyof Entity;
       this.autoGeneratedKeyStrategy = this.determineGenerationStrategy(

packages/test/src/test/rag/ScopedStorage.test.ts

@@ -4,12 +4,11 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage";
 import {
   createKnowledgeBase,
+  getGlobalKnowledgeBaseRepository,
   isSharedTableMode,
   registerKnowledgeBase,
-  getGlobalKnowledgeBaseRepository,
   ScopedTabularStorage,
   ScopedVectorStorage,
   SharedChunkIndexes,
@@ -19,11 +18,15 @@ import {
   SharedDocumentPrimaryKey,
   SharedDocumentStorageSchema,
 } from "@workglow/knowledge-base";
+import { InMemoryTabularStorage, InMemoryVectorStorage } from "@workglow/storage";
 import { uuid4 } from "@workglow/util";
 import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
 
 describe("ScopedTabularStorage", () => {
-  let sharedStorage: InstanceType<typeof InMemoryTabularStorage>;
+  let sharedStorage: InMemoryTabularStorage<
+    typeof SharedDocumentStorageSchema,
+    typeof SharedDocumentPrimaryKey
+  >;
   let scopeA: ScopedTabularStorage<any, any>;
   let scopeB: ScopedTabularStorage<any, any>;
 
@@ -244,7 +247,10 @@ describe("ScopedTabularStorage", () => {
 });
 
 describe("ScopedVectorStorage", () => {
-  let sharedStorage: InstanceType<typeof InMemoryVectorStorage>;
+  let sharedStorage: InMemoryVectorStorage<
+    typeof SharedChunkVectorStorageSchema,
+    typeof SharedChunkPrimaryKey
+  >;
   let scopeA: ScopedVectorStorage<any, any>;
   let scopeB: ScopedVectorStorage<any, any>;
 
@@ -261,7 +267,7 @@ describe("ScopedVectorStorage", () => {
   });
 
   afterEach(() => {
-    sharedStorage.destroy();
+    sharedStorage?.destroy?.();
   });
 
   test("getVectorDimensions delegates to inner", () => {