Skip to content

Commit

Permalink
Connecting to DynamoDB from a Lambda Function
Browse files Browse the repository at this point in the history
To connect to DynamoDB we have to set up the code in our Rust function to do so, and also *allow* it to do so from our CDK code. By default our function is not allowed to do much of anything besides run.

```bash
cargo add -p pokemon-api aws_config aws_sdk_dynamodb
```

Ideally we want our DynamoDB connection to get instantiated in the cold start of our lambda function, then get re-used across future invocations.

To do that we’ll set up a `static` called `DDB_CLIENT`. Uppercase letters here is a convention enforced by compiler warnings. A `static` item defined like this is a globally unique constant that lives as long as the program does.

The type of the static is going to be a `OnceCell<Client>` which we’ll instantiate with `OnceCell::const_new()`.

`OnceCell` is a data structure from the `tokio` crate that is only allowed to be set once and `Client` is a DynamoDB client.

This all comes together to mean that `DDB_CLIENT` is a globally unique constant that can only be set once, and contains an initialized DynamoDB Client.

```rust
use tokio::sync::OnceCell;
use aws_sdk_dynamodb::Client;

static DDB_CLIENT: OnceCell<Client> = OnceCell::const_new();

async fn get_global_client() -> &'static Client {
    DDB_CLIENT
        .get_or_init(|| async {
            let config =
                aws_config::load_from_env().await;
            let client = Client::new(&config);
            client
        })
        .await
}
```

We define an extra function here called `get_global_client`. The name of the function is arbitrary, but it’s `async` so that we can wait on it before our function runs.

This function returns a shared reference to a `Client` with a `'static` lifetime. The static lifetime means that this value lives until the end of our program.

This is because we’re sharing the `Client` that we’ve stored in `DDB_CLIENT`.

This function always calls `get_or_init` on the `DDB_CLIENT` constant. This allows us to instantiate the client if one doesn’t exist yet, or return a shared reference to the existing client if it does already.

On the first run of this function, the async block uses `aws_config` to load our AWS credentials from the environment of the lambda function, then constructs a DynamoDB `Client` in the same way we did when uploading data.

We then return that client, storing it in the `OnceCell` and the `get_or_init` will return a shared reference to the value we just constructed.

The first place we use this is in our main function. Note that we don’t actually use the value here, we’re just making sure it’s initialized in the cold start of our lambda function.

```rust
async fn main() -> Result<(), Error> {
    get_global_client().await;
    let handler_fn = service_fn(handler);
    lambda_runtime::run(handler_fn).await?;
    Ok(())
}
```

The second place we use the function is in our handler, where we do actually use the function to talk to DynamoDB.

First we need the `pokemon_table` name, which we’ll get via an environment variable.

Then we can get a shared reference to the global DynamoDB client using `get_global_client`.

```rust
let pokemon_table = env::var("POKEMON_TABLE")?;
let client = get_global_client().await;
```

The client itself offers a `get_item` builder that we can use to set the `pk` field to the requested Pokemon. Note that we have to wrap our string in the `AttributeValue` enum here.

We also pass in the table name and initiate the request with `send`, then immediately await its completion.

```rust
let resp = client
    .get_item()
    .key("pk", AttributeValue::S(pokemon_requested.to_string()))
    .table_name(pokemon_table)
    .send()
    .await?;
```

`resp.item` is an `Option<HashMap<String, AttributeValue>>` which is a type you’ll remember from when we uploaded the data to Dynamo. It’s the key/value pairs we sent up to Dyanmo in the first place.

We can then replace our previous response code, with a match on `resp.item` to handle the `Option`.

```rust
match resp.item {
    Some(item) => {
        Ok(ApiGatewayV2httpResponse {
            status_code: 200,
            headers: HeaderMap::new(),
            multi_value_headers: HeaderMap::new(
            ),
            body: Some(Body::Text(
                serde_json::to_string(
                    &json!({
                        "data": {
                            "id": item.get("pk").unwrap().as_s().unwrap(),
                            "name": item.get("name").unwrap().as_s().unwrap(),
                            "healthPoints": item.get("health_points").unwrap().as_n().unwrap()
                        },
                    }),
                )?,
            )),
            is_base64_encoded: Some(false),
            cookies: vec![],
        })
    }
    None => Ok(ApiGatewayV2httpResponse {
        status_code: 200,
        headers: HeaderMap::new(),
        multi_value_headers: HeaderMap::new(),
        body: Some(Body::Text(
            serde_json::to_string(&json!({
                "data": {}
            }))?,
        )),
        is_base64_encoded: Some(false),
        cookies: vec![],
    }),
}
```

In the `None` case we return an empty data object. You can choose to do whatever you want here, including returning a 404 or an error in the JSON instead.

In the success case we have a `HashMap<String, AttributeValue>` that we need to turn into JSON. Unfortunately, since the official aws-sdk-dynamodb is fairly new, we don’t have high level APIs to handle this for us.

`AttributeValue`, for example, doesn’t implement any traits that would make it work out of the box with serde.

So we’re left with having to handle it manually ourselves for now.

For each field we want to include in the response, we’ll use `item.get` which returns an `Option`. Since we *know* these fields have to exist, we can `unwrap` them. If we weren’t sure if the field was going to exist, we could handle it another way.

Once we have the underlying `AttributeValue`, we call `as_s`, or the appropriate as* function to “unwrap” the `AttributeValue` back into a `String`. Of course this can also fail, but we’ve stored the data in a way that should never fail, so we unwrap the `Result` here as well.

```rust
item.get("pk").unwrap().as_s().unwrap(),
```

unwrapping may not be the most elegant way to handle this code, but as long as the assumption we’re making is “this should never fail”, then it’s perfectly acceptable.

Now build and copy the binary into our pokemon-api directory once again.

```rust
cargo zigbuild --target x86_64-unknown-linux-gnu.2.26 --release -p pokemon-api
cp target/x86_64-unknown-linux-gnu/release/pokemon-api lambdas/pokemon-api/bootstrap
```

With the Rust code set up, we still need to fixup our CDK code.

Specifically we need to grant access to the DynamoDB table to our lambda. In this case we’ve chosen to give it full access, but we could also restrict it to just being able to run `get_item`.

We also set the `POKEMON_TABLE` environment variable in the lambda’s environment, so that our Rust code has access to the DynamoDB table name.

```rust
pokemonTable.grantFullAccess(pokemonLambda);
pokemonLambda.addEnvironment("POKEMON_TABLE", pokemonTable.tableName);
```

Then we can diff again to see the changes

```rust
❯ npm run cdk diff -- --profile rust-adventure-playground

> infra@0.1.0 cdk
> cdk "diff" "--profile" "rust-adventure-playground"

Stack InfraStack
IAM Statement Changes
┌───┬──────┬──────┬──────┬──────┬──────┐
│   │ Reso │ Effe │ Acti │ Prin │ Cond │
│   │ urce │ ct   │ on   │ cipa │ itio │
│   │      │      │      │ l    │ n    │
├───┼──────┼──────┼──────┼──────┼──────┤
│ + │ ${Po │ Allo │ dyna │ AWS: │      │
│   │ kemo │ w    │ modb │ ${Po │      │
│   │ nTab │      │ :*   │ kemo │      │
│   │ le.A │      │      │ nHan │      │
│   │ rn}  │      │      │ dler │      │
│   │      │      │      │ /Ser │      │
│   │      │      │      │ vice │      │
│   │      │      │      │ Role │      │
│   │      │      │      │ }    │      │
└───┴──────┴──────┴──────┴──────┴──────┘
(NOTE: There may be security-related changes not in this list. See aws/aws-cdk#1299)

Resources
[+] AWS::IAM::Policy PokemonHandler/ServiceRole/DefaultPolicy PokemonHandlerServiceRoleDefaultPolicy09C7DA9D
[~] AWS::Lambda::Function PokemonHandler PokemonHandlerC37D7DE3
 ├─ [+] Environment
 │   └─ {"Variables":{"POKEMON_TABLE":{"Ref":"PokemonTable7DFA0E9C"}}}
 └─ [~] DependsOn
     └─ @@ -1,3 +1,4 @@
        [ ] [
        [+]   "PokemonHandlerServiceRoleDefaultPolicy09C7DA9D",
        [ ]   "PokemonHandlerServiceRoleF58AC6D6"
        [ ] ]
```

and after deploying, we can curl the url again for `bulbasaur`, `charmander`, `bidoof`, or any other pokemon.

```shell
❯ curl https://72i5uisgr5.execute-api.us-east-1.amazonaws.com/pokemon/bulbasaur
{"data":{"healthPoints":"45","id":"bulbasaur","name":"Bulbasaur"}}
❯ curl https://72i5uisgr5.execute-api.us-east-1.amazonaws.com/pokemon/charmander
{"data":{"healthPoints":"39","id":"charmander","name":"Charmander"}}
❯ curl https://72i5uisgr5.execute-api.us-east-1.amazonaws.com/pokemon/bidoof
{"data":{"healthPoints":"59","id":"bidoof","name":"Bidoof"}}
```
  • Loading branch information
ChristopherBiscardi committed Apr 5, 2022
1 parent b51bc86 commit 7a158dc
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 14 deletions.
2 changes: 2 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions crates/pokemon-api/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
aws-config = "0.9.0"
aws-sdk-dynamodb = "0.9.0"
aws_lambda_events = "0.6.1"
http = "0.2.6"
lambda_runtime = "0.5.1"
Expand Down
81 changes: 67 additions & 14 deletions crates/pokemon-api/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,29 @@ use aws_lambda_events::{
ApiGatewayV2httpRequest, ApiGatewayV2httpResponse,
},
};
use aws_sdk_dynamodb::{model::AttributeValue, Client};
use http::HeaderMap;
use lambda_runtime::{service_fn, Error, LambdaEvent};
use serde_json::json;
use std::env;
use tokio::sync::OnceCell;

static DDB_CLIENT: OnceCell<Client> = OnceCell::const_new();

async fn get_global_client() -> &'static Client {
DDB_CLIENT
.get_or_init(|| async {
let shared_config =
aws_config::load_from_env().await;
let client = Client::new(&shared_config);
client
})
.await
}

#[tokio::main]
async fn main() -> Result<(), Error> {
get_global_client().await;
let handler_fn = service_fn(handler);
lambda_runtime::run(handler_fn).await?;
Ok(())
Expand All @@ -35,20 +52,56 @@ async fn handler(
cookies: vec![],
}),
Some(pokemon_requested) => {
Ok(ApiGatewayV2httpResponse {
status_code: 200,
headers: HeaderMap::new(),
multi_value_headers: HeaderMap::new(),
body: Some(Body::Text(
serde_json::to_string(&json!({
"data": {
"requested": pokemon_requested
}
}))?,
)),
is_base64_encoded: Some(false),
cookies: vec![],
})
let pokemon_table = env::var("POKEMON_TABLE")?;
let client = get_global_client().await;

let resp = client
.get_item()
.key(
"pk",
AttributeValue::S(
pokemon_requested.to_string(),
),
)
.table_name(pokemon_table)
.send()
.await?;

match resp.item {
Some(item) => {
Ok(ApiGatewayV2httpResponse {
status_code: 200,
headers: HeaderMap::new(),
multi_value_headers: HeaderMap::new(
),
body: Some(Body::Text(
serde_json::to_string(
&json!({
"data": {
"id": item.get("pk").unwrap().as_s().unwrap(),
"name": item.get("name").unwrap().as_s().unwrap(),
"healthPoints": item.get("health_points").unwrap().as_n().unwrap()
},
}),
)?,
)),
is_base64_encoded: Some(false),
cookies: vec![],
})
}
None => Ok(ApiGatewayV2httpResponse {
status_code: 200,
headers: HeaderMap::new(),
multi_value_headers: HeaderMap::new(),
body: Some(Body::Text(
serde_json::to_string(&json!({
"data": {}
}))?,
)),
is_base64_encoded: Some(false),
cookies: vec![],
}),
}
}
}
}
3 changes: 3 additions & 0 deletions infra/lib/infra-stack.js
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ class InfraStack extends Stack {
memorySize: 1024,
});

pokemonTable.grantFullAccess(pokemonLambda);
pokemonLambda.addEnvironment("POKEMON_TABLE", pokemonTable.tableName);

const pokemonIntegration = new integrations.HttpLambdaIntegration(
"PokemonIntegration",
pokemonLambda
Expand Down

0 comments on commit 7a158dc

Please sign in to comment.