Update lambda.py

Author: lmangani

lambda.py

@@ -1,16 +1,19 @@
-## import os
 import json
+
 import chdb
 
-def lambda_handler(event, context):
-    query = event['query'] or "SELECT version()"
-    format = event['default_format'] or "JSONCompact"
-    res = chdb.query(query, format)
-    out = json.loads(res.data())
+
+def handler(event, context):
+    if "requestContext" in event:
+        event = json.loads(event["body"])
+    query = event["query"] if "query" in event else "SELECT version()"
+    format = event["default_format"] if "default_format" in event else "JSONCompact"
+
+    res = chdb.query(query, format).data()
     return {
         "statusCode": 200,
         "headers": {
             "Content-Type": "application/json"
         },
-        "body": json.dumps(out)
+        "body": str(res) if not isinstance(res, (dict, list)) else json.dumps(res),
     }