From 113394ef7b80df03363a1337f181c97d4a583671 Mon Sep 17 00:00:00 2001
From: Emanuele De Cupis <ema.decu@gmail.com>
Date: Thu, 28 Mar 2024 15:14:44 +0100
Subject: [PATCH] refactor DocumentSearchKey

---
 src/utils/__tests__/cosmosdb_model.test.ts | 28 +++++++++++++++++++++-
 src/utils/cosmosdb_model.ts                | 20 ++++++++++------
 src/utils/cosmosdb_model_versioned.ts      | 14 ++---------
 3 files changed, 42 insertions(+), 20 deletions(-)

diff --git a/src/utils/__tests__/cosmosdb_model.test.ts b/src/utils/__tests__/cosmosdb_model.test.ts
index e9995cbb..5deba7d5 100644
--- a/src/utils/__tests__/cosmosdb_model.test.ts
+++ b/src/utils/__tests__/cosmosdb_model.test.ts
@@ -11,7 +11,12 @@ import {
 } from "@azure/cosmos";
 
 import { NonEmptyString } from "@pagopa/ts-commons/lib/strings";
-import { BaseModel, CosmosdbModel, CosmosResource } from "../cosmosdb_model";
+import {
+  BaseModel,
+  CosmosdbModel,
+  CosmosResource,
+  DocumentSearchKey
+} from "../cosmosdb_model";
 
 beforeEach(() => {
   jest.resetAllMocks();
@@ -82,6 +87,27 @@ const errorResponse: ErrorResponse = new Error();
 // eslint-disable-next-line functional/immutable-data
 errorResponse.code = 500;
 
+type Equal<X, Y extends X> = X extends Y ? (Y extends X ? X : never) : never;
+
+describe("DocumentSearchKey", () => {
+  type MyModel = { foo: string; bar: number; baz: boolean[] };
+
+  // alway allow id as a search key
+  type _0 = Equal<DocumentSearchKey<MyModel, "id">, [string]>;
+  // allow a string as partition key
+  type _1 = Equal<DocumentSearchKey<MyModel, "id", "foo">, [string, string]>;
+  // same model and partition key
+  type _2 = Equal<DocumentSearchKey<MyModel, "foo", "foo">, [string]>;
+  // @ts-expect-error MyModel["bar"] is not a string
+  type _3 = Equal<DocumentSearchKey<MyModel, "bar">, [string]>;
+  // @ts-expect-error MyModel["baz"] is not a string or number
+  type _4 = Equal<DocumentSearchKey<MyModel, "id", "baz">, [string, string]>;
+  // allow custom fields as search key
+  type _5 = Equal<DocumentSearchKey<MyModel, "foo">, [string]>;
+  // @ts-expect-error "pippo" is not a field of MyModel
+  type _6 = Equal<DocumentSearchKey<MyModel, "pippo">, [string]>;
+});
+
 describe("create", () => {
   it("should create a document", async () => {
     containerMock.items.create.mockResolvedValueOnce(
diff --git a/src/utils/cosmosdb_model.ts b/src/utils/cosmosdb_model.ts
index b61b730f..ed0ca55a 100644
--- a/src/utils/cosmosdb_model.ts
+++ b/src/utils/cosmosdb_model.ts
@@ -57,13 +57,19 @@ export type DocumentSearchKey<
   // Hence we omit "extends BaseModel", but we check keys to be part of "T & BaseModel"
   ModelIdKey extends keyof (T & BaseModel),
   PartitionKey extends keyof (T & BaseModel) = ModelIdKey
-> = (T & BaseModel)[ModelIdKey] extends string // narrow type to the ones that might be an identity
-  ? PartitionKey extends ModelIdKey // eslint-disable-next-line functional/prefer-readonly-type
-    ? [(T & BaseModel)[ModelIdKey]] // eslint-disable-next-line functional/prefer-readonly-type
-    : PartitionKey extends keyof (T & BaseModel) // eslint-disable-next-line functional/prefer-readonly-type
-    ? [(T & BaseModel)[ModelIdKey], (T & BaseModel)[PartitionKey]]
-    : never
-  : never;
+> =
+  // We must be sure the provided keys refer to fields whose value is appropriate
+  // the modelId must be a string
+  // the partitiion key might be a string or a number
+  Pick<BaseModel & T, ModelIdKey | PartitionKey> extends Record<
+    ModelIdKey,
+    string
+  > &
+    Record<PartitionKey, string | number>
+    ? PartitionKey extends ModelIdKey // partition key === model id means no partition key is provided
+      ? readonly [string]
+      : readonly [string, string | number]
+    : never;
 
 export type AzureCosmosResource = t.TypeOf<typeof AzureCosmosResource>;
 export const AzureCosmosResource = t.interface({
diff --git a/src/utils/cosmosdb_model_versioned.ts b/src/utils/cosmosdb_model_versioned.ts
index bd27a17e..4122cbd6 100644
--- a/src/utils/cosmosdb_model_versioned.ts
+++ b/src/utils/cosmosdb_model_versioned.ts
@@ -81,18 +81,8 @@ const getPartitionKeyFromSearchKey = <
   PartitionKey extends keyof T = ModelIdKey
 >(
   searchKey: DocumentSearchKey<T, ModelIdKey, PartitionKey>
-): CosmosPartitionKey => {
-  const value: unknown =
-    typeof searchKey[1] !== "undefined" ? searchKey[1] : searchKey[0];
-  if (
-    typeof value === "string" ||
-    typeof value === "number" ||
-    typeof value === "boolean"
-  ) {
-    return value;
-  }
-  return null;
-};
+): CosmosPartitionKey =>
+  typeof searchKey[1] === "undefined" ? searchKey[0] : searchKey[1];
 
 /**
  * Assumption: the model ID is also the partition key