pick lambda handler based on $_HANDLER environment variable

Author: tonisuter

AWSLambdaSwift/Sources/AWSLambdaSwift/Runtime.swift

@@ -1,6 +1,6 @@
 import Foundation
 
-func log(_ object: Any, flush: Bool = false) {
+public func log(_ object: Any, flush: Bool = false) {
     fputs("\(object)\n", stderr)
     if flush {
         fflush(stderr)
@@ -13,7 +13,7 @@ public class Runtime {
     var counter = 0
     let urlSession: URLSession
     let awsLambdaRuntimeAPI: String
-    let handler: String
+    let lambdaName: String
     var lambdas: [String: (JSONDictionary) -> JSONDictionary]
 
     public init() throws {
@@ -25,9 +25,13 @@ public class Runtime {
            let handler = environment["_HANDLER"] else {
               throw RuntimeError.missingEnvironmentVariables
         }
-        
+
+        guard let periodIndex = handler.index(of: ".") else {
+            throw RuntimeError.invalidHandlerName
+        }
+
         self.awsLambdaRuntimeAPI = awsLambdaRuntimeAPI
-        self.handler = handler
+        self.lambdaName = String(handler[handler.index(after: periodIndex)...])
     }
 
     func getNextInvocation() throws -> (input: JSONDictionary, requestId: String) {
@@ -75,7 +79,10 @@ public class Runtime {
             counter += 1
             log("Invocation-Counter: \(counter)")
 
-            let lambda = lambdas["lambda"]!
+            guard let lambda = lambdas[lambdaName] else {
+                throw RuntimeError.unknownLambdaHandler
+            }
+
             let output = lambda(input)
             try postInvocationResponse(for: requestId, response: output)
         }

AWSLambdaSwift/Sources/AWSLambdaSwift/RuntimeError.swift

@@ -1,6 +1,8 @@
 enum RuntimeError: Error {
     case missingEnvironmentVariables
+    case invalidHandlerName
     case endpointError(String)
     case missingData
     case invalidData
+    case unknownLambdaHandler
 }

ExampleLambda/Sources/ExampleLambda/main.swift

@@ -1,10 +1,16 @@
 import AWSLambdaSwift
 
-func lambda(input: JSONDictionary) -> JSONDictionary {
-    print(input)
+func lambda1(input: JSONDictionary) -> JSONDictionary {
+    log(input)
     return ["hello": "world", "number": 42, "strings": ["one", "two", "three"]]
 }
 
+func lambda2(input: JSONDictionary) -> JSONDictionary {
+    log(input)
+    return ["result": "success"]
+}
+
 let runtime = try Runtime()
-runtime.registerLambda("lambda", handler: lambda)
+runtime.registerLambda("lambda1", handler: lambda1)
+runtime.registerLambda("lambda2", handler: lambda2)
 try runtime.start()

bootstrap

@@ -1,2 +1,3 @@
 #!/bin/sh
-/opt/swift-shared-libs/ld-linux-x86-64.so.2 --library-path /opt/swift-shared-libs/lib ./ExampleLambda
+EXECUTABLE=$LAMBDA_TASK_ROOT/"$(echo $_HANDLER | cut -d. -f1)"
+/opt/swift-shared-libs/ld-linux-x86-64.so.2 --library-path /opt/swift-shared-libs/lib $EXECUTABLE