def lambda_handler(event, context):
x_api_key = event['headers']['x-api-key']
response = requests.get(
CREDENTIALS_API_URL,
headers={'x-api-key': x_api_key}
)
if response.status_code == requests.codes.ok:
metaflow_creds = response.json()
print(metaflow_creds)
for key in metaflow_creds.keys():
os.environ[key] = metaflow_creds[key]
os.environ['USERNAME'] = '******'
subsegment = xray_recorder.begin_subsegment('load_metaflow')
import api
xray_recorder.end_subsegment()
return api.handle_request(event, context)
else:
print("Unable to locate credentials")
return {
'statusCode': 403,
'body': json.dumps("Could not obtain Metaflow credentials")
}
def download_intermediate_results(filename):
"""Download a file from S3
Parameters
----------
filename: string, required
Name of the file in S3 source bucket (OpenAQ)
Returns
-------
processed_file: string
Local path to downloaded file
"""
xray_recorder.begin_subsegment('## download_data_file')
subsegment = xray_recorder.current_subsegment()
subsegment.put_metadata('filename', f's3://{output_bucket}/{filename}')
try:
processed_file = os.path.join('/tmp', os.path.basename(filename))
s3.download_file(output_bucket, filename, processed_file)
subsegment.put_annotation('DATA_DOWNLOAD', 'SUCCESS')
except botocore.exceptions.ClientError as e:
subsegment.put_annotation('DATA_DOWNLOAD', 'FAILURE')
log.error(f'Unable to download rsult file: {filename}')
log.debug(e)
raise
xray_recorder.end_subsegment()
return processed_file
def handler(event, context):
raw_records = event["Records"]
logger.debug(raw_records)
log_dict = dict()
failed_dict = dict()
xray_recorder.begin_subsegment("parse")
for payload in kinesis.parse_records(raw_records):
try:
payload_parsed = json.loads(payload)
except json.JSONDecodeError:
logger.debug(f"Ignoring non-JSON data: {payload}")
continue
baikonur_logging.parse_payload_to_log_dict(
payload_parsed,
log_dict,
failed_dict,
LOG_TYPE_FIELD,
LOG_TIMESTAMP_FIELD,
LOG_ID_FIELD,
LOG_TYPE_UNKNOWN_PREFIX,
LOG_TYPE_FIELD_WHITELIST,
timestamp_required=True,
)
xray_recorder.end_subsegment()
baikonur_logging.save_json_logs_to_s3(s3_client, failed_dict,
"Valid log data")
baikonur_logging.save_json_logs_to_s3(
s3_client, failed_dict, "One or more necessary fields are unavailable")
def wrapper(*args, **kw):
from ..query import XRayQuery, XRaySession
from ...flask_sqlalchemy.query import XRaySignallingSession
class_name = str(cls.__module__)
c = xray_recorder._context
sql = None
subsegment = None
if class_name == "sqlalchemy.orm.session":
for arg in args:
if isinstance(arg, XRaySession):
sql = parse_bind(arg.bind)
if isinstance(arg, XRaySignallingSession):
sql = parse_bind(arg.bind)
if class_name == 'sqlalchemy.orm.query':
for arg in args:
if isinstance(arg, XRayQuery):
try:
sql = parse_bind(arg.session.bind)
# Commented our for later PR
# sql['sanitized_query'] = str(arg)
except:
sql = None
if sql is not None:
if getattr(c._local, 'entities', None) is not None:
subsegment = xray_recorder.begin_subsegment(sql['url'],
namespace='remote')
else:
subsegment = None
res = func(*args, **kw)
if subsegment is not None:
subsegment.set_sql(sql)
subsegment.put_annotation("sqlalchemy",
class_name + '.' + func.__name__)
xray_recorder.end_subsegment()
return res
def handler(event, context):
# Extract out S3 object details
valid_bucket = event['valid_bucket_name']
valid_key = event['valid_key']
nonexistent_bucket = event['nonexistent_bucket']
nonexistent_key = event['nonexistent_key']
# Create a segment
xray_recorder.begin_segment('s3trace')
# Copy valid S3 object locally to /tmp
valid_subsegment = xray_recorder.begin_subsegment('valid')
local_file = local_directory + '/' + valid_key
print 'Copying valid object s3://{}/{} to {}...'.format(valid_bucket, valid_key, local_file)
try:
s3.Bucket(valid_bucket).download_file(valid_key, '{}'.format(local_file))
except Exception as e:
print 'Error: {}'.format(str(e))
xray_recorder.end_subsegment()
# Copy invalid S3 object locally to /tmp
invalid_subsegment = xray_recorder.begin_segment('invalid')
local_file = local_directory + '/' + nonexistent_key
print 'Copying invalid object s3://{}/{} to {}...'.format(nonexistent_bucket, nonexistent_key, local_file)
try:
s3.Bucket(nonexistent_bucket).download_file(nonexistent_key, '{}'.format(local_file))
except Exception as e:
print 'Error: {}'.format(str(e))
xray_recorder.end_subsegment()
# End segment
xray_recorder.end_segment()
return '{"message": "X-Ray worked"}'
def load_prerequisites(ctx, object_list):
for o in object_list:
xray_recorder.begin_subsegment("prereq:%s" % o)
log.debug(f"Loading prerequisite '{o}'...")
ctx[o].get_prerequisites()
xray_recorder.end_subsegment()
log.debug(f"End prerequisite loading...")
def upload_final_results(results):
"""Upload a file to S3
Parameters
----------
results: string, required
Name of the local file with final results
"""
xray_recorder.begin_subsegment('## upload_final_results')
subsegment = xray_recorder.current_subsegment()
subsegment.put_metadata(
'filename',
f's3://{output_bucket}/lambda-etl-refarch/output/{results}')
results_path = os.path.join('/tmp', results)
# upload to target S3 bucket
try:
response = s3.upload_file(
results_path, output_bucket,
'lambda-etl-refarch/output/{}'.format(results))
log.info(
f"Uploaded final results to s3://{output_bucket}/lambda-etl-refarch/output/{results}"
)
subsegment.put_annotation('FINAL_RESULTS_UPLOAD', 'SUCCESS')
except botocore.exceptions.ClientError as e:
subsegment.put_annotation('FINAL_RESULTS_UPLOAD', 'FAILURE')
log.error(f'Unable to upload final results: {results}')
log.debug(e)
raise
xray_recorder.end_subsegment()
async def run_task(task, sqs_msg):
global execution_is_completed_flag
xray_recorder.begin_segment('run_task')
logging.info("Running Task: {}".format(task))
xray_recorder.begin_subsegment('encoding')
bin_protobuf = prepare_arguments_for_execution(task)
tast_str = bin_protobuf.decode("utf-8")
task_def = json.loads(tast_str)
submit_pre_agent_measurements(task)
task_id = task["task_id"]
fname_stdout = "./stdout-{task_id}.log".format(task_id=task_id)
fname_stderr = "./stderr-{task_id}.log".format(task_id=task_id)
f_stdout = open(fname_stdout, "w")
f_stderr = open(fname_stderr, "w")
xray_recorder.end_subsegment()
execution_is_completed_flag = 0
task_execution = asyncio.create_task(
do_task_local_lambda_execution_thread(perf_tracker_post, task, sqs_msg,
task_def))
task_ttl_update = asyncio.create_task(do_ttl_updates_thread(task))
await asyncio.gather(task_execution, task_ttl_update)
f_stdout.close()
f_stderr.close()
xray_recorder.end_segment()
logging.info("Finished Task: {}".format(task))
return True
async def end_subsegment(session, trace_config_ctx, params):
if trace_config_ctx.give_up:
return
subsegment = xray_recorder.current_subsegment()
subsegment.put_http_meta(http.STATUS, params.response.status)
xray_recorder.end_subsegment()
def _get_wiki_url(_url, q):
BASE_URL = _url
payload = {}
resp = {"statusCode": 400}
HOT_TOPICS = [
'cholas', 'cheras', 'pandyas', 'pallavas', 'sangam_era', 'kural'
]
if q:
q = q.split('/')[-1]
if not q:
q = random.choice(HOT_TOPICS)
try:
random_sleep()
if _trigger_exception():
xray_recorder.put_annotation("SIMULATED_ERRORS", "True")
xray_recorder.begin_subsegment("BRITTLE_LEGACY_APP")
d = xray_recorder.current_subsegment()
d.put_annotation("MANUALLY_TRIGGERRED_IN_SUBSEGMENT", "True")
xray_recorder.end_subsegment()
raise Exception("RANDOM_ERROR: Simulate Mystique Failure")
r1 = requests.get(f'{BASE_URL}/{q}', params=payload)
xray_recorder.put_metadata('RESPONSE', resp)
resp["statusCode"] = r1.status_code
z = r1.json()
resp["body"] = json.dumps({"message": z["body"]["message"]})
_ddb_put_item(resp)
except Exception as e:
resp["body"] = json.dumps({"message": str(e)})
return resp
def lambda_handler(event, context):
"""
AWS Lambda handler
This method is invoked by the API Gateway: /Prod/first/{proxy+} endpoint.
"""
xray_subsegment = xray_recorder.current_subsegment()
xray_subsegment.put_annotation(
'application', '{{ cookiecutter.project_name.lower().replace('
',' - ') }}')
xray_subsegment.put_metadata(
'event', event, '{{ cookiecutter.project_name.lower().replace('
',' - ') }}')
try:
subsegment = xray_recorder.begin_subsegment('message')
message = {
'Id': uuid.uuid4().hex,
'Count': random.random() * 100,
}
subsegment.put_metadata(
'message', message, '{{ cookiecutter.project_name.lower().replace('
',' - ') }}')
xray_recorder.end_subsegment()
return {"statusCode": 200, "body": json.dumps(message)}
except Exception as err: # pragma: no cover
logger.error(str(err))
traceback.print_exc()
raise err
def failed(self, event):
subsegment = xray_recorder.current_subsegment()
subsegment.add_fault_flag()
subsegment.put_annotation('mongodb_duration_micros',
event.duration_micros)
subsegment.put_metadata('failure', event.failure)
xray_recorder.end_subsegment()
def hello(event, context):
body = {
"message": "Go Serverless v1.0! Your function executed successfully!",
"input": event
}
# Start an xray subsegment, because hte segment is already crreated automatically
subsegment = xray_recorder.begin_subsegment('numpy_subsegment')
# Create a random matrix using numpy and print it
a = np.arange(50).reshape(10, 5)
print("Your numpy array:")
print(a)
big_array()
# Add metadata and annotations
subsegment.put_annotation('event_key', 'numpy')
subsegment.put_metadata('event-data', event, 'data-received')
# Close the subsegment
xray_recorder.end_subsegment()
response = {"statusCode": 200, "body": json.dumps(body)}
return response
async def do_task_local_execution_thread(perf_tracker, task, sqs_msg, task_def,
f_stdout, f_stderr, fname_stdout):
global execution_is_completed_flag
xray_recorder.begin_subsegment('sub-process-1')
command = [
"./mock_compute_engine", task_def["worker_arguments"][0],
task_def["worker_arguments"][1], task_def["worker_arguments"][2]
]
print(command)
proc = subprocess.Popen(command,
stdout=f_stdout,
stderr=f_stderr,
shell=False)
while True:
retcode = proc.poll()
if retcode is not None:
execution_is_completed_flag = 1 # indicate that this thread is completed
process_subprocess_completion(perf_tracker, task, sqs_msg,
fname_stdout)
xray_recorder.end_subsegment()
return retcode
await asyncio.sleep(work_proc_status_pull_interval_sec)
def handler(event: sqs_event, context: LambdaContext) -> str:
xray: Subsegment = xray_recorder.begin_subsegment('Invocation Data')
xray.put_metadata('event', event)
xray.put_metadata('context', context)
xray_recorder.end_subsegment()
records: List[Dict] = event.get('Records')
return json.dumps({"Result": send_email(records)})
def extract_dd_trace_context(event):
"""
Extract Datadog trace context from the Lambda `event` object.
Write the context to a global `dd_trace_context`, so the trace
can be continued on the outgoing requests with the context injected.
Save the context to an X-Ray subsegment's metadata field, so the X-Ray
trace can be converted to a Datadog trace in the Datadog backend with
the correct context.
"""
global dd_trace_context
headers = event.get('headers', {})
trace_id = headers.get(TraceHeader.TRACE_ID)
parent_id = headers.get(TraceHeader.PARENT_ID)
sampling_priority = headers.get(TraceHeader.SAMPLING_PRIORITY)
if trace_id and parent_id and sampling_priority:
dd_trace_context = {
'trace-id': trace_id,
'parent-id': parent_id,
'sampling-priority': sampling_priority,
}
xray_recorder.begin_subsegment(XraySubsegment.NAME)
subsegment = xray_recorder.current_subsegment()
subsegment.put_metadata(
XraySubsegment.KEY,
dd_trace_context,
XraySubsegment.NAMESPACE
)
xray_recorder.end_subsegment()
else:
# AWS Lambda runtime caches global variables between invocations,
# reset to avoid using the context from the last invocation.
dd_trace_context = {}
def succeeded(self, event):
subsegment = xray_recorder.current_subsegment()
subsegment.put_annotation('mongodb_duration_micros',
event.duration_micros)
if self.record_full_documents:
subsegment.put_metadata('mongodb_reply', event.reply)
xray_recorder.end_subsegment()
def lambda_handler(event, context):
subsegment = xray_recorder.begin_subsegment('active-directory')
print("active directory event:{}".format(event))
subsegment.put_annotation('lambda_id', context.aws_request_id)
subsegment.put_metadata('AD', json, "ad.widgets.com")
xray_recorder.end_subsegment()
return {"statusCode": 200}
def lambda_handler(event, context):
symbol = event.get('queryStringParameters', {}).get('symbol', None)
xray_recorder.begin_subsegment('dynamo-call')
dynamo = boto3.client('dynamodb')
_item = dynamo.get_item(TableName='ChaosTrader',
Key={'symbol': {
'S': symbol
}})
xray_recorder.end_subsegment()
item = _item.get('Item')
if item:
is_good_buy = _item.get('Item', {}).get('is_good_buy',
{}).get('BOOL', None)
return {
"isBase64Encoded": False,
"statusCode": 200,
"body": json.dumps({"is_good_buy": is_good_buy}),
"headers": {
"Access-Control-Allow-Origin": "*"
}
}
else:
return {
"isBase64Encoded": False,
"statusCode": 404,
"body": json.dumps({}),
"headers": {
"Access-Control-Allow-Origin": "*"
}
}
def stuff(key):
if request.method == 'PUT':
xray_recorder.begin_subsegment('encryption')
data = codecs.encode(request.data.decode('utf-8'), 'rot_13')
xray_recorder.end_subsegment()
app.logger.info('Setting \"{}\"=\"{}\"'.format(key, data))
s3_client.Object(bucket, 'cache/' + key).put(Body=data)
body = json.dumps({'status': 'OK'})
status = 201
elif request.method == 'DELETE':
obj = s3_client.Object(bucket, 'cache/' + key)
obj.delete()
body = ''
status = 204
elif request.method == 'GET':
try:
app.logger.info('Getting \"{}\"'.format(key))
obj = s3_client.Object(bucket, 'cache/' + key)
body = obj.get()['Body'].read().decode('utf-8')
xray_recorder.begin_subsegment('decryption')
body = codecs.decode(body, 'rot_13')
xray_recorder.end_subsegment()
status = 200
except ClientError as ex:
app.logger.exception('Error', ex)
if ex.response['Error']['Code'] == 'NoSuchKey':
status = 204
body = ''
else:
raise ex
else:
body = json.dumps({'status': 'Method Not Supported'})
status = 406
app.logger.info('Status {}, Body \"{}\"'.format(status, body))
return make_response(body, status)
def __call__(self, request):
sampling_decision = None
meta = request.META
xray_header = construct_xray_header(meta)
# a segment name is required
name = calculate_segment_name(meta.get(HOST_KEY), xray_recorder)
sampling_req = {
'host': meta.get(HOST_KEY),
'method': request.method,
'path': request.path,
'service': name,
}
sampling_decision = calculate_sampling_decision(
trace_header=xray_header,
recorder=xray_recorder,
sampling_req=sampling_req,
)
if self.in_lambda_ctx:
segment = xray_recorder.begin_subsegment(name)
else:
segment = xray_recorder.begin_segment(
name=name,
traceid=xray_header.root,
parent_id=xray_header.parent,
sampling=sampling_decision,
)
segment.save_origin_trace_header(xray_header)
segment.put_http_meta(http.URL, request.build_absolute_uri())
segment.put_http_meta(http.METHOD, request.method)
if meta.get(USER_AGENT_KEY):
segment.put_http_meta(http.USER_AGENT, meta.get(USER_AGENT_KEY))
if meta.get(X_FORWARDED_KEY):
# X_FORWARDED_FOR may come from untrusted source so we
# need to set the flag to true as additional information
segment.put_http_meta(http.CLIENT_IP, meta.get(X_FORWARDED_KEY))
segment.put_http_meta(http.X_FORWARDED_FOR, True)
elif meta.get(REMOTE_ADDR_KEY):
segment.put_http_meta(http.CLIENT_IP, meta.get(REMOTE_ADDR_KEY))
response = self.get_response(request)
segment.put_http_meta(http.STATUS, response.status_code)
if response.has_header(CONTENT_LENGTH_KEY):
length = int(response[CONTENT_LENGTH_KEY])
segment.put_http_meta(http.CONTENT_LENGTH, length)
response[http.XRAY_HEADER] = prepare_response_header(
xray_header, segment)
if self.in_lambda_ctx:
xray_recorder.end_subsegment()
else:
xray_recorder.end_segment()
return response
def query_by_partition_and_sort_key(self, partition_key, partition_value,
sort_key, sort_value):
xray_recorder.begin_subsegment('query')
response = self.db_table.query(
KeyConditionExpression=Key(partition_key).eq(partition_value)
& Key(sort_key).gte(sort_value))
xray_recorder.end_subsegment('query')
return response.get('Items')
def handler(event, context):
print(json.dumps(event))
subsegment = xray_recorder.begin_subsegment("CustomSubsegment")
subsegment.put_annotation("tbd", "tbd")
response = json.dumps(event)
xray_recorder.end_subsegment()
return response
def index():
xray_recorder.begin_subsegment('SSM')
secret = get_secret()
xray_recorder.end_subsegment()
resp = requests.get('http://example.com')
return f"Secret {secret}, Status: {resp.status_code}"
def get_secret():
secret_name = "starAlliances/dev/seatMap/provider"
region_name = "us-east-1"
try:
xray_recorder.begin_subsegment('init_client')
# Create a Secrets Manager client
session = boto3.session.Session()
client = session.client(
service_name='secretsmanager',
region_name=region_name
)
xray_recorder.current_subsegment().put_annotation('init_client', 'done')
finally:
xray_recorder.end_subsegment()
# In this sample we only handle the specific exceptions for the 'GetSecretValue' API.
# See https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
# We rethrow the exception by default.
try:
get_secret_value_response = client.get_secret_value(
SecretId=secret_name
)
except ClientError as e:
if e.response['Error']['Code'] == 'DecryptionFailureException':
# Secrets Manager can't decrypt the protected secret text using the provided KMS key.
# Deal with the exception here, and/or rethrow at your discretion.
raise e
elif e.response['Error']['Code'] == 'InternalServiceErrorException':
# An error occurred on the server side.
# Deal with the exception here, and/or rethrow at your discretion.
raise e
elif e.response['Error']['Code'] == 'InvalidParameterException':
# You provided an invalid value for a parameter.
# Deal with the exception here, and/or rethrow at your discretion.
raise e
elif e.response['Error']['Code'] == 'InvalidRequestException':
# You provided a parameter value that is not valid for the current state of the resource.
# Deal with the exception here, and/or rethrow at your discretion.
raise e
elif e.response['Error']['Code'] == 'ResourceNotFoundException':
# We can't find the resource that you asked for.
# Deal with the exception here, and/or rethrow at your discretion.
raise e
else:
# Decrypts secret using the associated KMS CMK.
# Depending on whether the secret is a string or binary, one of these fields will be populated.
if 'SecretString' in get_secret_value_response:
secret = get_secret_value_response['SecretString']
return secret
else:
decoded_binary_secret = base64.b64decode(get_secret_value_response['SecretBinary'])
return decoded_binary_secret
# Your code goes here.
return ""
def callCustomSegment():
subsegment = xray_recorder.begin_subsegment('subsegment_name')
subsegment.put_annotation('one', '1')
subsegment.put_metadata('integration-test', 'true')
xray_recorder.end_subsegment()
return 'Ok! annotation-metadata testing'
def main_handler(event, context):
log.debug("Handler start.")
r = RLT(lambda args, kwargs, r: True, main_handler_entrypoint, event, context)
# Persist all aggregated data
xray_recorder.begin_subsegment("main_handler_entrypoint.persist_aggregates")
KVTable.persist_aggregates()
xray_recorder.end_subsegment()
log.log(log.NOTICE, "Normal end.")
return r
def lambda_handler(event, context):
symbol = event.get('symbol', None)
if symbol == 'warm':
time.sleep(10)
xray_recorder.begin_subsegment('compute-duration')
is_good_buy = get_buying_advice(symbol)
xray_recorder.end_subsegment()
store_symbol_details(symbol, is_good_buy)
async def end_subsegment_with_exception(session, trace_config_ctx, params):
subsegment = xray_recorder.current_subsegment()
subsegment.add_exception(
params.exception,
traceback.extract_stack(limit=xray_recorder._max_trace_back))
if isinstance(params.exception, LOCAL_EXCEPTIONS):
subsegment.namespace = LOCAL_NAMESPACE
xray_recorder.end_subsegment()
def subsegment(self, name):
segment = xray_recorder.begin_subsegment(name)
try:
yield segment
except Exception as e:
stack = traceback.extract_stack(limit=xray_recorder.max_trace_back)
segment.add_exception(e, stack)
raise
finally:
xray_recorder.end_subsegment(time.time())
def subsegment(self, name):
segment = xray_recorder.begin_subsegment(name)
try:
yield segment
except Exception as e:
stack = traceback.extract_stack(limit=xray_recorder.max_trace_back)
segment.add_exception(e, stack)
raise
finally:
xray_recorder.end_subsegment(time.time())
def __exit__(self, exc_type=None, exc_value=None, exc_traceback=None):
metadata = self.ctx.get_metadata(('api-stats',))
metadata.update(self.metadata)
xray_recorder.put_metadata('custodian', metadata)
if self.in_lambda:
xray_recorder.end_subsegment()
return
xray_recorder.end_segment()
if not self.use_daemon:
self.emitter.flush()
self.metadata.clear()
