feat(sdk): remove runtime dependency for AWS account ID retrieval (#319)

The previous AWS queue resource type implementation required the `AWS_ACCOUNT_ID` environment variable to be set by the runtime handler. The setting only occurred upon receiving a request, causing a panic if the queue resource type was used globally without the `AWS_ACCOUNT_ID` being set. This commit eliminates the need for runtime setting of `AWS_ACCOUNT_ID` by utilizing the `sts.GetCallerIdentity` API to retrieve the account ID.
pluto-lang · Aug 6, 2024 · cdf966f · cdf966f
1 parent 35cba9c
commit cdf966f
Show file tree

Hide file tree

Showing 13 changed files with 99 additions and 69 deletions.
diff --git a/.changeset/early-onions-guess.md b/.changeset/early-onions-guess.md
@@ -0,0 +1,9 @@
+---
+"@plutolang/pluto": patch
+---
+
+feat(sdk): remove runtime dependency for AWS account ID retrieval
+
+The previous AWS queue resource type implementation required the `AWS_ACCOUNT_ID` environment variable to be set by the runtime handler. The setting only occurred upon receiving a request, causing a panic if the queue resource type was used globally without the `AWS_ACCOUNT_ID` being set.
+
+This commit eliminates the need for runtime setting of `AWS_ACCOUNT_ID` by utilizing the `sts.GetCallerIdentity` API to retrieve the account ID.
diff --git a/docs/dev_guide/extend-sdk.en.md b/docs/dev_guide/extend-sdk.en.md
@@ -193,22 +193,19 @@ export class SNSQueue implements IQueueClient {
     };
     await this.client.send(
       new PublishCommand({
-        TopicArn: this.buildARN(this.topicName),
+        TopicArn: await this.buildARN(this.topicName),
         Message: JSON.stringify(evt),
       })
     );
   }
 
-  private buildARN(topicName: string): string {
+  private async buildARN(topicName: string): Promise<string> {
     const region = process.env.AWS_REGION;
     if (!region) {
       throw new Error("Missing AWS Region");
     }
 
-    const accountId = process.env.AWS_ACCOUNT_ID;
-    if (!accountId) {
-      throw new Error("Missing AWS Account ID");
-    }
+    const accountId = await getAwsAccountId();
 
     return `arn:aws:sns:${region}:${accountId}:${topicName}`;
   }

diff --git a/docs/dev_guide/extend-sdk.zh-CN.md b/docs/dev_guide/extend-sdk.zh-CN.md
@@ -193,22 +193,19 @@ export class SNSQueue implements IQueueClient {
     };
     await this.client.send(
       new PublishCommand({
-        TopicArn: this.buildARN(this.topicName),
+        TopicArn: await this.buildARN(this.topicName),
         Message: JSON.stringify(evt),
       })
     );
   }
 
-  private buildARN(topicName: string): string {
+  private async buildARN(topicName: string): Promise<string> {
     const region = process.env.AWS_REGION;
     if (!region) {
       throw new Error("Missing AWS Region");
     }
 
-    const accountId = process.env.AWS_ACCOUNT_ID;
-    if (!accountId) {
-      throw new Error("Missing AWS Account ID");
-    }
+    const accountId = await getAwsAccountId();
 
     return `arn:aws:sns:${region}:${accountId}:${topicName}`;
   }

diff --git a/examples/quickstart-python/package.json b/examples/quickstart-python/package.json
@@ -12,8 +12,8 @@
   "devDependencies": {
     "@types/node": "^20",
     "typescript": "^5.2.2",
-    "@plutolang/base": "latest",
-    "@plutolang/pluto-infra": "latest",
+    "@plutolang/base": "workspace:*",
+    "@plutolang/pluto-infra": "workspace:*",
     "@pulumi/pulumi": "^3.88.0"
   },
   "main": "dist/index.js"

diff --git a/packages/pluto-py/CHANGELOG.md b/packages/pluto-py/CHANGELOG.md
@@ -1,5 +1,13 @@
 # pluto-client
 
+## 0.0.19
+
+feat: remove runtime dependency for AWS account ID retrieval
+
+The previous AWS queue resource type implementation required the `AWS_ACCOUNT_ID` environment variable to be set by the runtime handler. The setting only occurred upon receiving a request, causing a panic if the queue resource type was used globally without the `AWS_ACCOUNT_ID` being set.
+
+This commit eliminates the need for runtime setting of `AWS_ACCOUNT_ID` by utilizing the `sts.GetCallerIdentity` API to retrieve the account ID.
+
 ## 0.0.18
 
 chore: lint and format code

diff --git a/packages/pluto-py/pluto_client/clients/aws/queue_sns.py b/packages/pluto-py/pluto_client/clients/aws/queue_sns.py
@@ -4,7 +4,7 @@
 import boto3
 from typing import Optional
 from pluto_base.utils.resource_id import gen_resource_id
-from pluto_client.clients.aws.utils import gen_aws_resource_name
+from pluto_client.clients.aws.utils import gen_aws_resource_name, get_aws_account_id
 from ...queue import CloudEvent, IQueueClient, Queue, QueueOptions
 
 
@@ -26,8 +26,6 @@ def __build_arn(self, topic_name: str) -> str:
         if not region:
             raise EnvironmentError("Missing AWS Region")
 
-        account_id = os.environ.get("AWS_ACCOUNT_ID")
-        if not account_id:
-            raise EnvironmentError("Missing AWS Account ID")
+        account_id = get_aws_account_id()
 
         return f"arn:aws:sns:{region}:{account_id}:{topic_name}"
diff --git a/packages/pluto-py/pluto_client/clients/aws/utils.py b/packages/pluto-py/pluto_client/clients/aws/utils.py
@@ -1,6 +1,7 @@
 import re
-import hashlib
 import json
+import boto3
+import hashlib
 
 RESOURCE_NAME_MAX_LENGTH = 50
 
@@ -15,3 +16,8 @@ def gen_aws_resource_name(*parts: str) -> str:
         start = len(resource_full_id) - (RESOURCE_NAME_MAX_LENGTH - len(hash))
         end = len(resource_full_id)
         return (resource_full_id[start:end] + hash).strip("-")
+
+
+def get_aws_account_id() -> str:
+    client = boto3.client("sts")  # type: ignore
+    return client.get_caller_identity().get("Account")  # type: ignore
diff --git a/packages/pluto-py/poetry.lock b/packages/pluto-py/poetry.lock
diff --git a/packages/pluto-py/pyproject.toml b/packages/pluto-py/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "pluto-client"
-version = "0.0.18"
+version = "0.0.19"
 description = "The Client Library for Pluto Programming Language."
 authors = ["Jade Zheng <zheng.shoujian@outlook.com>"]
 license = "Apache-2.0"
@@ -25,7 +25,7 @@ Issues = "https://github.com/pluto-lang/pluto/issues"
 
 [tool.poetry.dependencies]
 python = "^3.10"
-pluto-base = "^0.0.6"
+pluto-base = "^0.0.7"
 
 [[tool.poetry.source]]
 name = "testpypi"

diff --git a/packages/pluto/package.json b/packages/pluto/package.json
@@ -23,6 +23,7 @@
     "@aws-sdk/client-sagemaker-runtime": "^3.614.0",
     "@aws-sdk/client-secrets-manager": "^3.614.0",
     "@aws-sdk/client-sns": "^3.614.0",
+    "@aws-sdk/client-sts": "^3.614.0",
     "@aws-sdk/lib-dynamodb": "^3.614.0",
     "@plutolang/base": "workspace:^",
     "redis": "^4.6.10"

diff --git a/packages/pluto/src/clients/aws/queue.sns.ts b/packages/pluto/src/clients/aws/queue.sns.ts
@@ -1,14 +1,14 @@
 import { SNSClient, PublishCommand } from "@aws-sdk/client-sns";
 import { genResourceId } from "@plutolang/base/utils";
 import { CloudEvent, IQueueClient, Queue, QueueOptions } from "../../queue";
-import { genAwsResourceName } from "./utils";
+import { genAwsResourceName, getAwsAccountId } from "./utils";
 
 /**
  * Implementation of Queue using AWS SNS.
  */
 export class SNSQueue implements IQueueClient {
   private readonly id: string;
-  private readonly topicArn: string;
+  private readonly topicArn: Promise<string>;
 
   private client: SNSClient;
 
@@ -28,22 +28,19 @@ export class SNSQueue implements IQueueClient {
     };
     await this.client.send(
       new PublishCommand({
-        TopicArn: this.topicArn,
+        TopicArn: await this.topicArn,
         Message: JSON.stringify(evt),
       })
     );
   }
 
-  private buildARN(topicName: string): string {
+  private async buildARN(topicName: string): Promise<string> {
     const region = process.env.AWS_REGION;
     if (!region) {
       throw new Error("Missing AWS Region");
     }
 
-    const accountId = process.env.AWS_ACCOUNT_ID;
-    if (!accountId) {
-      throw new Error("Missing AWS Account ID");
-    }
+    const accountId = await getAwsAccountId();
 
     return `arn:aws:sns:${region}:${accountId}:${topicName}`;
   }

diff --git a/packages/pluto/src/clients/aws/utils.ts b/packages/pluto/src/clients/aws/utils.ts
@@ -1,3 +1,4 @@
+import { STSClient, GetCallerIdentityCommand } from "@aws-sdk/client-sts";
 import { createHash } from "crypto";
 
 const RESOURCE_NAME_MAX_LENGTH = 50;
@@ -20,3 +21,16 @@ export function genAwsResourceName(...parts: string[]): string {
     return (resourceFullId.substring(start, end) + hash).replace(/^-+|-+$/g, "");
   }
 }
+
+export async function getAwsAccountId(): Promise<string> {
+  const client = new STSClient({});
+
+  const command = new GetCallerIdentityCommand({});
+  const response = await client.send(command);
+
+  const accountId = response.Account;
+  if (!accountId) {
+    throw new Error("Missing AWS Account ID");
+  }
+  return accountId;
+}