chore: refactor generic persistence layer

dreamorosi · dreamorosi · commit d98ca6243664 · 2023-09-14T13:49:47.000Z
diff --git a/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayer.ts b/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayer.ts
@@ -3,115 +3,92 @@ import {
   IdempotencyItemNotFoundError,
   IdempotencyRecordStatus,
 } from '@aws-lambda-powertools/idempotency';
-import {
-  IdempotencyRecordOptions,
-  IdempotencyRecordStatusValue,
-} from '@aws-lambda-powertools/idempotency/types';
+import { IdempotencyRecordOptions } from '@aws-lambda-powertools/idempotency/types';
 import {
   IdempotencyRecord,
   BasePersistenceLayer,
 } from '@aws-lambda-powertools/idempotency/persistence';
 import { getSecret } from '@aws-lambda-powertools/parameters/secrets';
 import { Transform } from '@aws-lambda-powertools/parameters';
 import {
-  CacheClient,
-  CredentialProvider,
-  Configurations,
-  CacheGet,
-  CacheKeyExists,
-  CollectionTtl,
-  CacheDictionarySetFields,
-  CacheDictionaryGetFields,
-} from '@gomomento/sdk';
-import type { MomentoApiSecret, Item } from './types';
-
-class MomentoCachePersistenceLayer extends BasePersistenceLayer {
-  #cacheName: string;
-  #client?: CacheClient;
-
-  public constructor(config: { cacheName: string }) {
+  ProviderClient,
+  ProviderItemAlreadyExists,
+} from './advancedBringYourOwnPersistenceLayerProvider';
+import type { ApiSecret, ProviderItem } from './types';
+
+class CustomPersistenceLayer extends BasePersistenceLayer {
+  #collectionName: string;
+  #client?: ProviderClient;
+
+  public constructor(config: { collectionName: string }) {
     super();
-    this.#cacheName = config.cacheName;
+    this.#collectionName = config.collectionName;
   }
 
   protected async _deleteRecord(record: IdempotencyRecord): Promise<void> {
     await (
       await this.#getClient()
-    ).delete(this.#cacheName, record.idempotencyKey);
+    ).delete(this.#collectionName, record.idempotencyKey);
   }
 
   protected async _getRecord(
     idempotencyKey: string
   ): Promise<IdempotencyRecord> {
-    const response = await (
-      await this.#getClient()
-    ).dictionaryFetch(this.#cacheName, idempotencyKey);
+    try {
+      const item = await (
+        await this.#getClient()
+      ).get(this.#collectionName, idempotencyKey);
 
-    if (
-      response instanceof CacheGet.Error ||
-      response instanceof CacheGet.Miss
-    ) {
+      return new IdempotencyRecord({
+        ...(item as unknown as IdempotencyRecordOptions),
+      });
+    } catch (error) {
       throw new IdempotencyItemNotFoundError();
     }
-    const { data, ...rest } =
-      response.value() as unknown as IdempotencyRecordOptions & {
-        data: string;
-      };
-
-    return new IdempotencyRecord({
-      responseData: JSON.parse(data),
-      ...rest,
-    });
   }
 
   protected async _putRecord(record: IdempotencyRecord): Promise<void> {
-    const item: Partial<Item> = {
+    const item: Partial<ProviderItem> = {
       status: record.getStatus(),
     };
 
     if (record.inProgressExpiryTimestamp !== undefined) {
-      item.in_progress_expiration = record.inProgressExpiryTimestamp.toString();
+      item.in_progress_expiration = record.inProgressExpiryTimestamp;
     }
 
     if (this.isPayloadValidationEnabled() && record.payloadHash !== undefined) {
       item.validation = record.payloadHash;
     }
 
-    try {
-      const lock = await this.#lookupItem(record.idempotencyKey);
-
-      if (
-        lock.getStatus() !== IdempotencyRecordStatus.INPROGRESS &&
-        (lock.inProgressExpiryTimestamp || 0) < Date.now()
-      ) {
-        throw new IdempotencyItemAlreadyExistsError(
-          `Failed to put record for already existing idempotency key: ${record.idempotencyKey}`
-        );
-      }
-    } catch (error) {
-      if (error instanceof IdempotencyItemAlreadyExistsError) {
-        throw error;
-      }
-    }
-
     const ttl = record.expiryTimestamp
       ? Math.floor(new Date(record.expiryTimestamp * 1000).getTime() / 1000) -
         Math.floor(new Date().getTime() / 1000)
       : this.getExpiresAfterSeconds();
 
-    const response = await (
-      await this.#getClient()
-    ).dictionarySetFields(this.#cacheName, record.idempotencyKey, item, {
-      ttl: CollectionTtl.of(ttl).withNoRefreshTtlOnUpdates(),
-    });
-
-    if (response instanceof CacheDictionarySetFields.Error) {
-      throw new Error(`Unable to put item: ${response.errorCode()}`);
+    let existingItem: ProviderItem | undefined;
+    try {
+      existingItem = await (
+        await this.#getClient()
+      ).put(this.#collectionName, record.idempotencyKey, item, {
+        ttl,
+      });
+    } catch (error) {
+      if (error instanceof ProviderItemAlreadyExists) {
+        if (
+          existingItem &&
+          existingItem.status !== IdempotencyRecordStatus.INPROGRESS &&
+          (existingItem.in_progress_expiration || 0) < Date.now()
+        ) {
+          throw new IdempotencyItemAlreadyExistsError(
+            `Failed to put record for already existing idempotency key: ${record.idempotencyKey}`
+          );
+        }
+      }
     }
   }
 
   protected async _updateRecord(record: IdempotencyRecord): Promise<void> {
-    const value: Partial<Item> = {
+    const value: Partial<ProviderItem> = {
       data: JSON.stringify(record.responseData),
       status: record.getStatus(),
     };
@@ -120,82 +97,34 @@ class MomentoCachePersistenceLayer extends BasePersistenceLayer {
       value.validation = record.payloadHash;
     }
 
-    await this.#checkItemExists(record.idempotencyKey);
-
     await (
       await this.#getClient()
-    ).dictionarySetFields(this.#cacheName, record.idempotencyKey, value, {
-      ttl: CollectionTtl.refreshTtlIfProvided().withNoRefreshTtlOnUpdates(),
-    });
+    ).update(this.#collectionName, record.idempotencyKey, value);
   }
 
-  async #getMomentoApiSecret(): Promise<MomentoApiSecret> {
-    const secretName = process.env.MOMENTO_API_SECRET;
+  async #getClient(): Promise<ProviderClient> {
+    if (this.#client) return this.#client;
+
+    const secretName = process.env.API_SECRET;
     if (!secretName) {
-      throw new Error('MOMENTO_API_SECRET environment variable is not set');
+      throw new Error('API_SECRET environment variable is not set');
     }
 
-    const apiSecret = await getSecret<MomentoApiSecret>(secretName, {
+    const apiSecret = await getSecret<ApiSecret>(secretName, {
       transform: Transform.JSON,
     });
 
     if (!apiSecret) {
       throw new Error(`Could not retrieve secret ${secretName}`);
     }
 
-    return apiSecret;
-  }
-
-  async #getClient(): Promise<CacheClient> {
-    if (this.#client) return this.#client;
-
-    const apiSecret = await this.#getMomentoApiSecret();
-    this.#client = await CacheClient.create({
-      configuration: Configurations.InRegion.LowLatency.latest(),
-      credentialProvider: CredentialProvider.fromString({
-        apiKey: apiSecret.apiKey,
-      }),
+    this.#client = new ProviderClient({
+      apiKey: apiSecret.apiKey,
       defaultTtlSeconds: this.getExpiresAfterSeconds(),
     });
 
     return this.#client;
   }
-
-  async #checkItemExists(idempotencyKey: string): Promise<boolean> {
-    const response = await (
-      await this.#getClient()
-    ).keysExist(this.#cacheName, [idempotencyKey]);
-
-    return response instanceof CacheKeyExists.Success;
-  }
-
-  async #lookupItem(idempotencyKey: string): Promise<IdempotencyRecord> {
-    const response = await (
-      await this.#getClient()
-    ).dictionaryGetFields(this.#cacheName, idempotencyKey, [
-      'in_progress_expiration',
-      'status',
-    ]);
-
-    if (response instanceof CacheDictionaryGetFields.Miss) {
-      throw new IdempotencyItemNotFoundError();
-    } else if (response instanceof CacheDictionaryGetFields.Error) {
-      throw new Error('Unable to get item');
-    } else {
-      const { status, in_progress_expiration: inProgressExpiryTimestamp } =
-        response.value() || {};
-
-      if (status !== undefined || inProgressExpiryTimestamp !== undefined) {
-        throw new Error('Unable');
-      }
-
-      return new IdempotencyRecord({
-        idempotencyKey,
-        status: status as IdempotencyRecordStatusValue,
-        inProgressExpiryTimestamp: parseFloat(inProgressExpiryTimestamp),
-      });
-    }
-  }
 }
 
-export { MomentoCachePersistenceLayer };
+export { CustomPersistenceLayer };
diff --git a/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayerProvider.ts b/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayerProvider.ts
@@ -0,0 +1,44 @@
+import type { ProviderItem } from './types';
+
+/**
+ * This is a mock implementation of an SDK client for a generic key-value store.
+ */
+class ProviderClient {
+  public constructor(_config: { apiKey: string; defaultTtlSeconds: number }) {
+    // ...
+  }
+
+  public async delete(_collectionName: string, _key: string): Promise<void> {
+    // ...
+  }
+
+  public async get(
+    _collectionName: string,
+    _key: string
+  ): Promise<ProviderItem> {
+    // ...
+    return {} as ProviderItem;
+  }
+
+  public async put(
+    _collectionName: string,
+    _key: string,
+    _value: Partial<ProviderItem>,
+    _options: { ttl: number }
+  ): Promise<ProviderItem> {
+    // ...
+    return {} as ProviderItem;
+  }
+
+  public async update(
+    _collectionName: string,
+    _key: string,
+    _value: Partial<ProviderItem>
+  ): Promise<void> {
+    // ...
+  }
+}
+
+class ProviderItemAlreadyExists extends Error {}
+
+export { ProviderClient, ProviderItemAlreadyExists };
diff --git a/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayerUsage.ts b/docs/snippets/idempotency/advancedBringYourOwnPersistenceLayerUsage.ts
@@ -1,14 +1,14 @@
 import type { Context } from 'aws-lambda';
 import { randomUUID } from 'node:crypto';
-import { MomentoCachePersistenceLayer } from './advancedBringYourOwnPersistenceLayer';
+import { CustomPersistenceLayer } from './advancedBringYourOwnPersistenceLayer';
 import {
   IdempotencyConfig,
   makeIdempotent,
 } from '@aws-lambda-powertools/idempotency';
 import type { Request, Response, SubscriptionResult } from './types';
 
-const persistenceStore = new MomentoCachePersistenceLayer({
-  cacheName: 'powertools',
+const persistenceStore = new CustomPersistenceLayer({
+  collectionName: 'powertools',
 });
 const config = new IdempotencyConfig({
   expiresAfterSeconds: 60,
diff --git a/docs/snippets/idempotency/templates/tableCdk.ts b/docs/snippets/idempotency/templates/tableCdk.ts
@@ -4,7 +4,7 @@ import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
 import { Runtime } from 'aws-cdk-lib/aws-lambda';
 import { AttributeType, BillingMode, Table } from 'aws-cdk-lib/aws-dynamodb';
 
-export class IdempotencyMomentoStack extends Stack {
+export class IdempotencyStack extends Stack {
   public constructor(scope: Construct, id: string, props?: StackProps) {
     super(scope, id, props);
 
diff --git a/docs/snippets/idempotency/types.ts b/docs/snippets/idempotency/types.ts
@@ -14,16 +14,16 @@ export type SubscriptionResult = {
   productId: string;
 };
 
-export type MomentoApiSecret = {
+export type ApiSecret = {
   apiKey: string;
   refreshToken: string;
   validUntil: number;
   restEndpoint: string;
 };
 
-export type Item = {
+export type ProviderItem = {
   validation?: string;
-  in_progress_expiration?: string;
+  in_progress_expiration?: number;
   status: IdempotencyRecordStatusValue;
   data: string;
 };
diff --git a/docs/snippets/package.json b/docs/snippets/package.json
@@ -33,7 +33,6 @@
     "@aws-sdk/client-secrets-manager": "^3.360.0",
     "@aws-sdk/client-ssm": "^3.360.0",
     "@aws-sdk/util-dynamodb": "^3.360.0",
-    "@gomomento/sdk": "^1.39.3",
     "aws-sdk": "^2.1405.0",
     "aws-sdk-client-mock": "^2.2.0",
     "aws-sdk-client-mock-jest": "^2.2.0",
diff --git a/docs/utilities/idempotency.md b/docs/utilities/idempotency.md
@@ -697,11 +697,11 @@ You can create your own persistent store from scratch by inheriting the `BasePer
 * `_updateRecord()` – Updates an item in the persistence store.
 * `_deleteRecord()` – Removes an item from the persistence store.
 
-Below an example of an alternative persistence layer backed by [Momento Cache](https://www.gomomento.com):
+Below an example implementation of a custom persistence layer backed by a generic key-value store.
 
-=== "MomentoCachePersistenceLayer"
+=== "CustomPersistenceLayer"
 
-    ```typescript hl_lines="12 28 37 43 67 113"
+    ```typescript hl_lines="9 19 28 34 50 90"
     --8<-- "docs/snippets/idempotency/advancedBringYourOwnPersistenceLayer.ts"
     ```
 
