import boto3 import requests import json import csv from internal.commons import HEADERS, DEFAULT_METHOD from internal.helpers import rds_query, DataProcessorError import logging from internal.pre_request_filters import default_arguments, make_clientid_param, make_current_package from internal.pre_request_filters import decode_alliance_workspace_id_params from product_packages import RIGHT_CHECKS from product_packages.exceptions import QuotaExceededException, InvalidRight log = logging.getLogger() ## # Query executors # handler(endpoint, request, event) -> result ## def exec_endpoint(endpoint, event, internal_request={}): ep = endpoint default_arguments(event, internal_request) # build request parameters if "pre_filters" in ep: for pre_filter in ep["pre_filters"]: internal_request = pre_filter(event, internal_request) # checks current product package validity only if we use clientid for the request. if make_clientid_param in ep["pre_filters"]: internal_request = make_current_package(event, internal_request) # check rights if 'rights_check' in endpoint: for right in endpoint['rights_check']: if right in RIGHT_CHECKS: if not RIGHT_CHECKS[right]['handler'](**internal_request): raise QuotaExceededException(f"Limit on {right} exceeded.") else: InvalidRight(f"Don't know how to check {right}") # call endpoint result = ep['query'](ep, internal_request, event) # finalize output for appsync if 'post_filters' in ep: for post_filter in ep['post_filters']: result = post_filter(result, internal_request) return result def http_query(endpoint, req, event): headers = endpoint['headers'] if 'headers' in endpoint else HEADERS method = endpoint['method'] if 'method' in endpoint else DEFAULT_METHOD if method == "GET": internal_response = requests.request(method, endpoint["url"], headers=headers, params=req) else: internal_response = requests.request(method, endpoint["url"], headers=headers, data=json.dumps(req)) if internal_response.ok: if len(internal_response.text) > 0: return internal_response.json() else: return None else: raise DataProcessorError(internal_response.json()) def decide_query(endpoint, req, event): if 'klaviyo' in req: return table_example_query(collection_id=req['collectionId'], schema=req['workspace_schema'], start=req['start'], count=req['count']) else: return s3csv_query(endpoint, req, event) def table_example_query(collection_id, schema, start, count): collections_query = f""" WITH collections AS ( SELECT DISTINCT ID FROM {schema}.collection WHERE (id = %(collectionId)s OR (parent_id = %(collectionId)s AND source IN ('EnrichKlaviyoEventsData', 'EnrichKlaviyoStaticData'))) AND status IN ('finished', 'labeling edit') ) """ query = f""" {collections_query} ,attribute_names AS ( SELECT DISTINCT file_field_name FROM {schema}.guessed_file_upload_fields WHERE collection_id IN (SELECT * FROM collections) ), attributes AS ( SELECT * FROM {schema}.attribute WHERE name IN (SELECT * FROM attribute_names)) SELECT * FROM attributes ORDER BY id ASC; """ attributes = rds_query(query, {'collectionId': collection_id}) log.info({'query_1': query, 'collection_id': collection_id}) attribute_string = '' select_string = [] for attribute in attributes: string = f',max(case when fa.attribute_id={attribute["id"]} then value end) AS "{attribute["name"]}"' attribute_string = attribute_string + string select_string.append(f'"{attribute["name"]}"') query = f""" {collections_query} SELECT {','.join(select_string)} FROM ( SELECT fan_id, row_id {attribute_string} FROM {schema}.fan_attribute fa WHERE collection_id IN (SELECT * FROM collections) AND row_id >= {start} AND row_id < {start+count} GROUP BY 1, 2 ) temp ; """ #log.info({'query_2': query, # 'collection_id': collection_id}) result = rds_query(query, {'collectionId': collection_id}) #log.info({'result': result}) resp = [] """Logic to support front end separate requests for getting header""" if start == 0 and count == 1: for index, row in enumerate(result): if index == 0: column_names = [] for key in row: column_names.append(key) resp.append(column_names) else: for row in result: values = [] for key in row: values.append(row[key]) resp.append(values) #log.info({'response': resp}) return resp def split_csv_lines(result): sniffer = csv.Sniffer() dialect = sniffer.sniff("\n".join(result[:5])) return [row for row in csv.reader(result, dialect)] def s3csv_query(endpoint, req, event): try: start = int(req['start']) count = int(req['count']) count = max(1, min(count, 1000)) # clip to [1; 1000] s3_client = boto3.client('s3') obj = s3_client.get_object(**req['object_params']) tags = s3_client.get_object_tagging(**req['object_params'])["TagSet"] encoding = next( # get first element of filtered list if it exists, otherwise get default ("utf_8") (tag for tag in tags if tag["Key"] == "encoding"), {"Value": "utf_8"}, # fallback )["Value"] lines_decoded = [] for i, line in enumerate(obj['Body'].iter_lines()): # always adding first line for the sniffer to detect dialect properly. if i == 0 or start <= i < start + count: lines_decoded.append(line.decode(encoding)) if i == start + count: break lines_split = split_csv_lines(lines_decoded) # skipping header line if it was not requested return lines_split[1:] if start > 0 else lines_split except KeyError as e: log.exception(e) raise RuntimeError('Please pass along correct start and count arguments') except Exception as e: log.exception(e) raise RuntimeError('S3 object retrieval failure') def s3_query(endpoint, req, event): try: s3_client = boto3.client('s3') obj = s3_client.get_object(**req['object_params']) stream = obj['Body'] return stream.read().decode() except KeyError as e: log.exception(e) raise RuntimeError('Please pass along correct start and count arguments') except Exception as e: log.exception(e) raise RuntimeError('S3 object retrieval failure') def rds_ep_query(endpoint, req, event): try: return rds_query(req['query'], decode_alliance_workspace_id_params(None, req['params']) if 'params' in req else {}) except Exception as e: log.exception(e) raise RuntimeError("RDS query failed") def channel_id_query(endpoint, req, event): #channel_name = req["channelName"] if "channelName" in req and \ # req["channelName"] in AVAILABLE_CHANNELS else DEFAULT_CHANNEL #return AVAILABLE_CHANNELS[channel_name](event) return req['workspace_schema'] def test_error_query(endpoint, req, event): raise RuntimeError("This is test error message")