def api_key(key: str = Depends(APIKeyHeader(name=API_KEY_HEADER_NAME))) -> str:
"""Validate an api key string matches the value currently in SecretsManager"""
region = APIServiceConfig().region
current_api_key_secret = get_api_key(region_name=region)
if key == current_api_key_secret:
return key
try:
pending_api_key_secret = get_api_key(version_stage="AWSPENDING", region_name=region)
if key == pending_api_key_secret:
return key
except ClientError as c_e:
response_error = getattr(c_e, "response", {}).get("Error", {})
error_code = response_error.get("Code", "")
if error_code != "ResourceNotFoundException":
raise c_e
raise HTTPException(status_code=HTTP_403_FORBIDDEN, detail="Could not validate credentials")
def test_get_api_key():
with mock.patch.dict(os.environ, {"API_KEY_SECRET_NAME": "test-api-key"}):
client = boto3.client("secretsmanager", region_name="us-west-2")
client.create_secret(
Name=SecurityConfig().api_key_secret_name, SecretString="testvalue123",
)
api_key = get_api_key(region_name="us-west-2")
assert api_key == "testvalue123"
def query(event: Dict[str, Any]) -> None:
"""Run the query portion of a QJ"""
query_config = QueryConfig()
logger = Logger()
logger.info(event=QJLogEvents.InitConfig, config=query_config)
records = event.get("Records", [])
if not records:
raise Exception("No records found")
if len(records) > 1:
raise Exception(
f"More than one record. BatchSize is probably not 1. event: {event}"
)
body = records[0].get("body")
if body is None:
raise Exception(
f"No record body found. BatchSize is probably not 1. event: {event}"
)
body = json.loads(body)
job = schemas.Job(**body)
logger.info(event=QJLogEvents.InitJob, job=job)
logger.info(event=QJLogEvents.RunQueryStart)
query_result = run_query(job=job, config=query_config)
logger.info(event=QJLogEvents.RunQueryEnd,
num_results=query_result.get_length())
results: List[schemas.Result] = []
if query_config.account_id_key not in query_result.query_result_set.fields:
raise Exception(
f"Query results must contain field '{query_config.account_id_key}'"
)
for q_r in query_result.to_list():
account_id = q_r[query_config.account_id_key]
result = schemas.Result(
account_id=account_id,
result={
key: val
for key, val in q_r.items()
if key != query_config.account_id_key
},
)
results.append(result)
graph_spec = schemas.ResultSetGraphSpec(
graph_uris_load_times=query_result.graph_uris_load_times)
result_set = schemas.ResultSetCreate(job=job,
graph_spec=graph_spec,
results=results)
api_key = get_api_key(region_name=query_config.region)
qj_client = QJAPIClient(host=query_config.api_host,
port=query_config.api_port,
api_key=api_key)
logger.info(event=QJLogEvents.CreateResultSetStart)
qj_client.create_result_set(result_set=result_set)
logger.info(event=QJLogEvents.CreateResultSetEnd)
def lambda_handler(_: Dict[str, Any], __: Any) -> None:
"""Lambda entrypoint"""
logger = Logger()
config = PrunerConfig()
logger.info(event=QJLogEvents.InitConfig, config=config)
api_key = get_api_key(region_name=config.region)
qj_client = QJAPIClient(host=config.api_host,
port=config.api_port,
api_key=api_key)
logger.info(event=QJLogEvents.DeleteStart)
result = qj_client.delete_expired_result_sets()
logger.info(event=QJLogEvents.DeleteEnd, result=result)
def pruner() -> None:
"""Prune results according to Job config settings"""
logger = Logger()
pruner_config = PrunerConfig()
logger.info(event=QJLogEvents.InitConfig, config=pruner_config)
api_key = get_api_key(region_name=pruner_config.region)
qj_client = QJAPIClient(
host=pruner_config.api_host, port=pruner_config.api_port, api_key=api_key
)
logger.info(event=QJLogEvents.DeleteStart)
result = qj_client.delete_expired_result_sets()
logger.info(event=QJLogEvents.DeleteEnd, result=result)