From 3fe3fc060ee54d182b834d6f7fd5827ddf06345d Mon Sep 17 00:00:00 2001 From: Michael Hall Date: Tue, 12 Dec 2023 12:30:57 -0600 Subject: [PATCH] - Result filename is now settable. --- pylot/plugins/rds_lambda/main.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/pylot/plugins/rds_lambda/main.py b/pylot/plugins/rds_lambda/main.py index 5d3700a..b0e80dc 100644 --- a/pylot/plugins/rds_lambda/main.py +++ b/pylot/plugins/rds_lambda/main.py @@ -46,23 +46,22 @@ def download_file(self, bucket, key, results, s3_client=None): return file -def query_rds(query_data, results='query_results.json', **kwargs): +def query_rds(query, results='query_results.json', **kwargs): rds = QueryRDS() - if isinstance(query_data, str) and os.path.isfile(query_data): - query_data = rds.read_json_file(query_data) + if isinstance(query, str) and os.path.isfile(query): + query = rds.read_json_file(query) else: - query_data = json.loads(query_data) + query = json.loads(query) - query_data = {'rds_config': query_data, 'is_test': True} + query = {'rds_config': query, 'is_test': True} - - rsp = rds.invoke_rds_lambda(query_data) + rsp = rds.invoke_rds_lambda(query) ret_dict = json.loads(rsp.get('Payload').read().decode('utf-8')) # Download results from S3 file = rds.download_file(bucket=ret_dict.get('bucket'), key=ret_dict.get('key'), results=results) print( - f'{ret_dict.get("count")} {query_data.get("rds_config").get("records")} records obtained: ' + f'{ret_dict.get("count")} {query.get("rds_config").get("records")} records obtained: ' f'{os.getcwd()}/{results}' ) return file @@ -95,8 +94,8 @@ def return_parser(subparsers): ) -def main(query=None, records=None, **kwargs): - query_rds(query_data=query, record_type=records) +def main(**kwargs): + query_rds(**kwargs) print('Complete') return 0