Skip to content

Commit 6c2257d

Browse files
author
Ramisa Alam
committed
feat: add tenant-id to lambda context
1 parent f65d966 commit 6c2257d

10 files changed

+162
-3
lines changed

‎awslambdaric/bootstrap.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,7 @@ def handle_event_request(
158158
cognito_identity_json,
159159
invoked_function_arn,
160160
epoch_deadline_time_in_ms,
161+
tenant_id,
161162
log_sink,
162163
):
163164
error_result = None
@@ -168,6 +169,7 @@ def handle_event_request(
168169
epoch_deadline_time_in_ms,
169170
invoke_id,
170171
invoked_function_arn,
172+
tenant_id,
171173
)
172174
event = lambda_runtime_client.marshaller.unmarshal_request(
173175
event_body, content_type
@@ -229,6 +231,7 @@ def create_lambda_context(
229231
epoch_deadline_time_in_ms,
230232
invoke_id,
231233
invoked_function_arn,
234+
tenant_id,
232235
):
233236
client_context = None
234237
if client_context_json:
@@ -243,6 +246,7 @@ def create_lambda_context(
243246
cognito_identity,
244247
epoch_deadline_time_in_ms,
245248
invoked_function_arn,
249+
tenant_id,
246250
)
247251

248252

@@ -528,5 +532,6 @@ def run(app_root, handler, lambda_runtime_api_addr):
528532
event_request.cognito_identity,
529533
event_request.invoked_function_arn,
530534
event_request.deadline_time_in_ms,
535+
event_request.tenant_id,
531536
log_sink,
532537
)

‎awslambdaric/lambda_context.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ def __init__(
1616
cognito_identity,
1717
epoch_deadline_time_in_ms,
1818
invoked_function_arn=None,
19+
tenant_id=None,
1920
):
2021
self.aws_request_id = invoke_id
2122
self.log_group_name = os.environ.get("AWS_LAMBDA_LOG_GROUP_NAME")
@@ -24,6 +25,7 @@ def __init__(
2425
self.memory_limit_in_mb = os.environ.get("AWS_LAMBDA_FUNCTION_MEMORY_SIZE")
2526
self.function_version = os.environ.get("AWS_LAMBDA_FUNCTION_VERSION")
2627
self.invoked_function_arn = invoked_function_arn
28+
self.tenant_id = tenant_id
2729

2830
self.client_context = make_obj_from_dict(ClientContext, client_context)
2931
if self.client_context is not None:
@@ -65,7 +67,8 @@ def __repr__(self):
6567
f"function_version={self.function_version},"
6668
f"invoked_function_arn={self.invoked_function_arn},"
6769
f"client_context={self.client_context},"
68-
f"identity={self.identity}"
70+
f"identity={self.identity},"
71+
f"tenant_id={self.tenant_id}"
6972
"])"
7073
)
7174

‎awslambdaric/lambda_runtime_client.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ def wait_next_invocation(self):
137137
deadline_time_in_ms=headers.get("Lambda-Runtime-Deadline-Ms"),
138138
client_context=headers.get("Lambda-Runtime-Client-Context"),
139139
cognito_identity=headers.get("Lambda-Runtime-Cognito-Identity"),
140+
tenant_id=headers.get("Lambda-Runtime-Aws-Tenant-Id"),
140141
content_type=headers.get("Content-Type"),
141142
event_body=response_body,
142143
)

‎awslambdaric/runtime_client.cpp

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ static PyObject *method_next(PyObject *self) {
5252
auto client_context = response.client_context.c_str();
5353
auto content_type = response.content_type.c_str();
5454
auto cognito_id = response.cognito_identity.c_str();
55+
auto tenant_id = response.tenant_id.c_str();
5556

5657
PyObject *payload_bytes = PyBytes_FromStringAndSize(payload.c_str(), payload.length());
5758
PyObject *result = Py_BuildValue("(O,{s:s,s:s,s:s,s:l,s:s,s:s,s:s})",
@@ -62,7 +63,8 @@ static PyObject *method_next(PyObject *self) {
6263
"Lambda-Runtime-Deadline-Ms", deadline,
6364
"Lambda-Runtime-Client-Context", NULL_IF_EMPTY(client_context),
6465
"Content-Type", NULL_IF_EMPTY(content_type),
65-
"Lambda-Runtime-Cognito-Identity", NULL_IF_EMPTY(cognito_id)
66+
"Lambda-Runtime-Cognito-Identity", NULL_IF_EMPTY(cognito_id),
67+
"Lambda-Runtime-Aws-Tenant-Id", NULL_IF_EMPTY(tenant_id)
6668
);
6769

6870
Py_XDECREF(payload_bytes);

‎deps/aws-lambda-cpp-0.2.6.tar.gz

-3.29 KB
Binary file not shown.

‎deps/patches/aws-lambda-cpp-add-tenant-id.patch

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
diff --git a/include/aws/lambda-runtime/runtime.h b/include/aws/lambda-runtime/runtime.h
2+
index 7812ff6..96be869 100644
3+
--- a/include/aws/lambda-runtime/runtime.h
4+
+++ b/include/aws/lambda-runtime/runtime.h
5+
@@ -61,6 +61,11 @@ struct invocation_request {
6+
*/
7+
std::string content_type;
8+
9+
+ /**
10+
+ * The Tenant ID of the current invocation.
11+
+ */
12+
+ std::string tenant_id;
13+
+
14+
/**
15+
* Function execution deadline counted in milliseconds since the Unix epoch.
16+
*/
17+
diff --git a/src/runtime.cpp b/src/runtime.cpp
18+
index e53b2b8..9763282 100644
19+
--- a/src/runtime.cpp
20+
+++ b/src/runtime.cpp
21+
@@ -40,6 +40,7 @@ static constexpr auto CLIENT_CONTEXT_HEADER = "lambda-runtime-client-context";
22+
static constexpr auto COGNITO_IDENTITY_HEADER = "lambda-runtime-cognito-identity";
23+
static constexpr auto DEADLINE_MS_HEADER = "lambda-runtime-deadline-ms";
24+
static constexpr auto FUNCTION_ARN_HEADER = "lambda-runtime-invoked-function-arn";
25+
+static constexpr auto TENANT_ID_HEADER = "lambda-runtime-aws-tenant-id";
26+
27+
enum Endpoints {
28+
INIT,
29+
@@ -289,6 +290,10 @@ runtime::next_outcome runtime::get_next()
30+
req.function_arn = resp.get_header(FUNCTION_ARN_HEADER);
31+
}
32+
33+
+ if (resp.has_header(TENANT_ID_HEADER)) {
34+
+ req.tenant_id = resp.get_header(TENANT_ID_HEADER);
35+
+ }
36+
+
37+
if (resp.has_header(DEADLINE_MS_HEADER)) {
38+
auto const& deadline_string = resp.get_header(DEADLINE_MS_HEADER);
39+
constexpr int base = 10;

‎scripts/update_deps.sh

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ wget -c https://github.com/awslabs/aws-lambda-cpp/archive/v$AWS_LAMBDA_CPP_RELEA
3030
patch -p1 < ../patches/aws-lambda-cpp-posting-init-errors.patch && \
3131
patch -p1 < ../patches/aws-lambda-cpp-make-the-runtime-client-user-agent-overrideable.patch && \
3232
patch -p1 < ../patches/aws-lambda-cpp-make-lto-optional.patch && \
33-
patch -p1 < ../patches/aws-lambda-cpp-add-content-type.patch
33+
patch -p1 < ../patches/aws-lambda-cpp-add-content-type.patch && \
34+
patch -p1 < ../patches/aws-lambda-cpp-add-tenant-id.patch
3435
)
3536

3637
## Pack again and remove the folder

‎tests/test_bootstrap.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ def test_handle_event_request_happy_case(self):
8888
{},
8989
"invoked_function_arn",
9090
0,
91+
"tenant_id",
9192
bootstrap.StandardLogSink(),
9293
)
9394
self.lambda_runtime.post_invocation_result.assert_called_once_with(
@@ -111,6 +112,7 @@ def test_handle_event_request_invalid_client_context(self):
111112
{},
112113
"invoked_function_arn",
113114
0,
115+
"tenant_id",
114116
bootstrap.StandardLogSink(),
115117
)
116118
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -152,6 +154,7 @@ def test_handle_event_request_invalid_cognito_idenity(self):
152154
"invalid_cognito_identity",
153155
"invoked_function_arn",
154156
0,
157+
"tenant_id",
155158
bootstrap.StandardLogSink(),
156159
)
157160
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -194,6 +197,7 @@ def test_handle_event_request_invalid_event_body(self):
194197
{},
195198
"invoked_function_arn",
196199
0,
200+
"tenant_id",
197201
bootstrap.StandardLogSink(),
198202
)
199203
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -238,6 +242,7 @@ def invalid_json_response(json_input, lambda_context):
238242
{},
239243
"invoked_function_arn",
240244
0,
245+
"tenant_id",
241246
bootstrap.StandardLogSink(),
242247
)
243248
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -283,6 +288,7 @@ def __init__(self, message):
283288
{},
284289
"invoked_function_arn",
285290
0,
291+
"tenant_id",
286292
bootstrap.StandardLogSink(),
287293
)
288294
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -335,6 +341,7 @@ def __init__(self, message):
335341
{},
336342
"invoked_function_arn",
337343
0,
344+
"tenant_id",
338345
bootstrap.StandardLogSink(),
339346
)
340347
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -386,6 +393,7 @@ def unable_to_import_module(json_input, lambda_context):
386393
{},
387394
"invoked_function_arn",
388395
0,
396+
"tenant_id",
389397
bootstrap.StandardLogSink(),
390398
)
391399
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -425,6 +433,7 @@ def raise_exception_handler(json_input, lambda_context):
425433
{},
426434
"invoked_function_arn",
427435
0,
436+
"tenant_id",
428437
bootstrap.StandardLogSink(),
429438
)
430439
args, _ = self.lambda_runtime.post_invocation_error.call_args
@@ -475,6 +484,7 @@ def raise_exception_handler(json_input, lambda_context):
475484
{},
476485
"invoked_function_arn",
477486
0,
487+
"tenant_id",
478488
bootstrap.StandardLogSink(),
479489
)
480490

@@ -514,6 +524,7 @@ def raise_exception_handler(json_input, lambda_context):
514524
{},
515525
"invoked_function_arn",
516526
0,
527+
"tenant_id",
517528
bootstrap.StandardLogSink(),
518529
)
519530
error_logs = (
@@ -546,6 +557,7 @@ def raise_exception_handler(json_input, lambda_context):
546557
{},
547558
"invoked_function_arn",
548559
0,
560+
"tenant_id",
549561
bootstrap.StandardLogSink(),
550562
)
551563
error_logs = (
@@ -578,6 +590,7 @@ def raise_exception_handler(json_input, lambda_context):
578590
{},
579591
"invoked_function_arn",
580592
0,
593+
"tenant_id",
581594
bootstrap.StandardLogSink(),
582595
)
583596
error_logs = (
@@ -619,6 +632,7 @@ def raise_exception_handler(json_input, lambda_context):
619632
{},
620633
"invoked_function_arn",
621634
0,
635+
"tenant_id",
622636
bootstrap.StandardLogSink(),
623637
)
624638
error_logs = lambda_unhandled_exception_warning_message + "\n[ERROR]\r"
@@ -652,6 +666,7 @@ def raise_exception_handler(json_input, lambda_context):
652666
{},
653667
"invoked_function_arn",
654668
0,
669+
"tenant_id",
655670
bootstrap.StandardLogSink(),
656671
)
657672

@@ -868,6 +883,7 @@ def test_application_json(self):
868883
cognito_identity_json=None,
869884
invoked_function_arn="invocation-arn",
870885
epoch_deadline_time_in_ms=1415836801003,
886+
tenant_id=None,
871887
log_sink=bootstrap.StandardLogSink(),
872888
)
873889

@@ -887,6 +903,7 @@ def test_binary_request_binary_response(self):
887903
cognito_identity_json=None,
888904
invoked_function_arn="invocation-arn",
889905
epoch_deadline_time_in_ms=1415836801003,
906+
tenant_id=None,
890907
log_sink=bootstrap.StandardLogSink(),
891908
)
892909

@@ -906,6 +923,7 @@ def test_json_request_binary_response(self):
906923
cognito_identity_json=None,
907924
invoked_function_arn="invocation-arn",
908925
epoch_deadline_time_in_ms=1415836801003,
926+
tenant_id=None,
909927
log_sink=bootstrap.StandardLogSink(),
910928
)
911929

@@ -924,6 +942,7 @@ def test_binary_with_application_json(self):
924942
cognito_identity_json=None,
925943
invoked_function_arn="invocation-arn",
926944
epoch_deadline_time_in_ms=1415836801003,
945+
tenant_id=None,
927946
log_sink=bootstrap.StandardLogSink(),
928947
)
929948

‎tests/test_lambda_context.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ def test_init(self):
3737
self.assertEqual(context.memory_limit_in_mb, "1234")
3838
self.assertEqual(context.function_version, "version1")
3939
self.assertEqual(context.invoked_function_arn, "arn:test1")
40+
self.assertEqual(context.tenant_id, None)
4041
self.assertEqual(context.identity.cognito_identity_id, None)
4142
self.assertEqual(context.identity.cognito_identity_pool_id, None)
4243
self.assertEqual(context.client_context.client.installation_id, None)
@@ -74,6 +75,21 @@ def test_init_cognito(self):
7475
self.assertEqual(context.identity.cognito_identity_id, "id1")
7576
self.assertEqual(context.identity.cognito_identity_pool_id, "poolid1")
7677

78+
def test_init_tenant_id(self):
79+
client_context = {}
80+
cognito_identity = {}
81+
tenant_id = "blue"
82+
83+
context = LambdaContext(
84+
"invoke-id1",
85+
client_context,
86+
cognito_identity,
87+
1415836801000,
88+
"arn:test",
89+
tenant_id,
90+
)
91+
self.assertEqual(context.tenant_id, "blue")
92+
7793
def test_init_client_context(self):
7894
client_context = {
7995
"client": {

0 commit comments

Comments
 (0)