Skip to content

Commit

Permalink
Allow custom partitions to be used at runtime (#3610)
Browse files Browse the repository at this point in the history
## Motivation and Context
This PR adds support for custom partition to be loaded at runtime. Use
cases for this feature are mostly private, i.e. using a pre-GA region
from the AWS SDK for Rust (or more broadly, from Smithy SDK for Rust).

## Description
Today, AWS SDK for Rust uses
[sdk-partitions.json](https://github.com/awslabs/aws-sdk-rust/blob/c6f4f53506604b8d52a5fd735b057c6225ed9aae/aws-models/sdk-partitions.json).
Unless a customer specifies a full URL using the `.endpoint_url` method,
the SDK is responsible for constructing the full URL for an endpoint
given a specified region, using a partitions file. It is essentially a
recipe for how to do so, such as providing DNS suffix like
`csp.hci.ic.gov`.

Certain customers need to use the SDK for a pre-GA region that is not
listed in the said public partitions file. When a region is not listed,
it will not be mapped to a known partition, and the SDK defaults to
construct endpoints using the patterns for the public AWS partition,
i.e. `amazonaws.com`. The assumption of "defaulting to the public
partition" may not always work for those customers. This PR, therefore,
adds an environment variable so that customers can specify a full path
to a custom partitions file.

Why look up the environment variable at runtime but not at codegen time?
The primary reason is maximum availability for the feature. Customers
can use a custom partitions file even for SDKs already published on
crates.io. If configuration were at codegen time, it wouldn't be
possible to use a custom partitions file for those SDKs.

## Testing
Since the code change uses an environment variable, no automated tests
have added for this. Instead, manually verified the following. Given a
fake custom paritions file that includes
```
  {
    "id" : "custom-partition",
    "outputs" : {
      "dnsSuffix" : "custom.com",
      "dualStackDnsSuffix" : "custom.com",
      "implicitGlobalRegion" : "us-custom-region-1",
      "name" : "custom-partition",
      "supportsDualStack" : false,
      "supportsFIPS" : true
    },
    "regionRegex" : "^us\\-custom\\-*",
    "regions" : {
      "us-custom-region-1" : {
        "description" : "US CUSTOM REGION 1"
      } 
    } 
  }
```
when we run the following test
```
#[tokio::test]
async fn test_custom_partition() {
    let (http_client, captured_request) = capture_request(None);
    std::env::set_var("SMITHY_CLIENT_SDK_CUSTOM_PARTITIONS", <full path to that custom partition>);
    let sdk_config = aws_config::from_env().region("us-custom-region-1").http_client(http_client).load().await;
    let client = Client::new(&sdk_config);
    let _ = client.list_buckets().send().await;
    dbg!(captured_request.expect_request());
}
```
`dbg!` shows a fully constructed endpoint URL:
```
    uri: Uri {
        as_string: "https://s3.us-custom-region-1.custom.com/?x-id=ListBuckets",
        parsed: H0(
            https://s3.us-custom-region-1.custom.com/?x-id=ListBuckets,
        ),
    },
```
If we do not set the environment variable in the above snippet, then we
get this endpoint URL instead:
```
    uri: Uri {
        as_string: "https://s3.us-custom-region-1.amazonaws.com/?x-id=ListBuckets",
        parsed: H0(
            https://s3.us-custom-region-1.amazonaws.com/?x-id=ListBuckets,
        ),
    },
```

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._
  • Loading branch information
ysaito1001 authored Apr 26, 2024
1 parent f15e7b3 commit 5461e4f
Showing 1 changed file with 14 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ class AwsPartitionResolver(runtimeConfig: RuntimeConfig, private val partitionsD
arrayOf(
"PartitionResolver" to EndpointsLib.partitionResolver(runtimeConfig),
"Lazy" to CargoDependency.OnceCell.toType().resolve("sync::Lazy"),
"tracing" to RuntimeType.Tracing,
)

override fun structFieldInit() =
Expand All @@ -70,7 +71,19 @@ class AwsPartitionResolver(runtimeConfig: RuntimeConfig, private val partitionsD
// Loading the partition JSON is expensive since it involves many regex compilations,
// so cache the result so that it only need to be paid for the first constructed client.
pub(crate) static DEFAULT_PARTITION_RESOLVER: #{Lazy}<#{PartitionResolver}> =
#{Lazy}::new(|| #{PartitionResolver}::new_from_json(b$json).expect("valid JSON"));
#{Lazy}::new(|| {
match std::env::var("SMITHY_CLIENT_SDK_CUSTOM_PARTITION") {
Ok(partitions) => {
#{tracing}::debug!("loading custom partitions located at {partitions}");
let partition_dot_json = std::fs::read_to_string(partitions).expect("should be able to read a custom partition JSON");
#{PartitionResolver}::new_from_json(partition_dot_json.as_bytes()).expect("valid JSON")
},
_ => {
#{tracing}::debug!("loading default partitions");
#{PartitionResolver}::new_from_json(b$json).expect("valid JSON")
}
}
});
""",
*codegenScope,
)
Expand Down

0 comments on commit 5461e4f

Please sign in to comment.