From 7e4d6fbac92adced9b9a9ea5fc057a99a4315690 Mon Sep 17 00:00:00 2001 From: dbernstein Date: Thu, 2 Nov 2023 11:29:56 -0700 Subject: [PATCH] Add option to upload result to s3 (#49) * Add output-json parameter to publish command for use by downstream processes. * Add support for storing publish command results to s3 * Add combined import-and-publish operation to simplify integration with playbook. --- core/cli.py | 119 +++++++++++++++++- .../publish_dashboard_from_template.py | 34 ++++- .../core/operation/test_publish_operation.py | 56 +++++++-- 3 files changed, 194 insertions(+), 15 deletions(-) diff --git a/core/cli.py b/core/cli.py index 114c081..89a7c5f 100644 --- a/core/cli.py +++ b/core/cli.py @@ -18,12 +18,17 @@ log = logging.getLogger("core.cli") +boto3.setup_default_session() + def create_quicksight_client(): - boto3.setup_default_session() return boto3.client("quicksight") +def create_s3_client(): + return boto3.client("s3") + + @click.group() def cli(): pass @@ -121,11 +126,29 @@ def import_template( help="The namespace you wish to target (e.g. tpp-prod, tpp-dev, tpp-staging).", ) @click.option("--group-name", required=True, help="Name of the Quicksight User Group") +@click.option( + "--output-json", + required=False, + help="The file path to which operation output should be written as json", +) +@click.option( + "--result-bucket", + required=False, + help="An S3 bucket to save the results to. If specified, you must also specify a result-key", +) +@click.option( + "--result-key", + required=False, + help="An S3 object key to save the results to. If used, result-bucket must be specified.", +) def publish_dashboard( aws_account_id: str, template_id: str, target_namespace: str, group_name: str, + output_json: str, + result_bucket: str, + result_key: str, ): """ Create/Update a dashboard from a template @@ -137,12 +160,106 @@ def publish_dashboard( log.info(f"group_name = {group_name}") result = PublishDashboardFromTemplateOperation( qs_client=create_quicksight_client(), + s3_client=create_s3_client(), aws_account_id=aws_account_id, template_id=template_id, target_namespace=target_namespace, group_name=group_name, + output_json=output_json, + result_bucket=result_bucket, + result_key=result_key, ).execute() log.info(result) cli.add_command(publish_dashboard) + + +@click.command +@click.option("--aws-account-id", required=True, help="The ID of the AWS account") +@click.option( + "--template-name", required=True, help="The name of the template to be restored" +) +@click.option( + "--data-source-arn", + required=True, + help="The ARN of the data source you want to associate with the data sets", +) +@click.option( + "--target-namespace", + required=True, + help="The namespace you wish to target (e.g. tpp-prod, tpp-dev, tpp-staging).", +) +@click.option( + "--input-dir", + required=True, + help="The path to the input directory from which resources will be imported", +) +@click.option("--group-name", required=True, help="Name of the Quicksight User Group") +@click.option( + "--output-json", + required=False, + help="The file path to which operation output should be written as json", +) +@click.option( + "--result-bucket", + required=False, + help="An S3 bucket to save the results to. If specified, you must also specify a result-key", +) +@click.option( + "--result-key", + required=False, + help="An S3 object key to save the results to. If used, result-bucket must be specified.", +) +def import_and_publish( + aws_account_id: str, + template_name: str, + data_source_arn: str, + target_namespace: str, + input_dir: str, + output_json: str, + group_name: str, + result_bucket: str, + result_key: str, +): + + log.info(f"import_and_publish") + log.info(f"aws_account_id = {aws_account_id}") + log.info(f"template_name = {template_name}") + log.info(f"data_source_arn = {data_source_arn}") + log.info(f"input_dir= {input_dir}") + log.info(f"group_name = {group_name}") + log.info(f"result_bucket = {result_bucket}") + log.info(f"result_key = {result_key}") + log.info(f"output_json = {output_json}") + + log.info(f"Importing {template_name}") + result = ImportFromJsonOperation( + qs_client=create_quicksight_client(), + aws_account_id=aws_account_id, + template_name=template_name, + target_namespace=target_namespace, + data_source_arn=data_source_arn, + input_dir=input_dir, + ).execute() + log.info(f"Import result: {result}") + template_id: str = result["template"]["id"] + log.info( + f"Publishing template {template_id} as dashboard using datasource {data_source_arn}" + ) + + result = PublishDashboardFromTemplateOperation( + qs_client=create_quicksight_client(), + s3_client=create_s3_client(), + aws_account_id=aws_account_id, + template_id=template_id, + target_namespace=target_namespace, + group_name=group_name, + result_bucket=result_bucket, + result_key=result_key, + output_json=output_json, + ).execute() + log.info(f"publish result = {result}") + + +cli.add_command(import_and_publish) diff --git a/core/operation/publish_dashboard_from_template.py b/core/operation/publish_dashboard_from_template.py index 4b62f51..93b7518 100644 --- a/core/operation/publish_dashboard_from_template.py +++ b/core/operation/publish_dashboard_from_template.py @@ -10,11 +10,24 @@ class PublishDashboardFromTemplateOperation(BaseOperation): """ def __init__( - self, template_id: str, target_namespace: str, group_name: str, *args, **kwargs + self, + template_id: str, + target_namespace: str, + group_name: str, + output_json: str, + result_bucket: str, + result_key: str, + s3_client, + *args, + **kwargs, ): self._template_id = template_id self._target_namespace = target_namespace self._group_name = group_name + self._output_json = output_json + self._result_bucket = result_bucket + self._result_key = result_key + self._s3_client = s3_client super().__init__(*args, **kwargs) def execute(self) -> dict: @@ -110,12 +123,25 @@ def execute(self) -> dict: f"Unexpected response from trying to update_dashboard_permissions : {json.dumps(response, indent=4)} " ) - return { + result = { "status": "success", - "dashboard_arn": dashboard_arn, - "dashboard_id": dashboard_id, + "dashboard_info": {self._template_id: [dashboard_arn]}, } + if self._output_json: + with open(self._output_json, "w") as output: + output.write(json.dumps(result)) + self._log.info(f"Output written to {self._output_json}") + + if self._result_bucket and self._result_key: + self._s3_client.put_object( + Bucket=self._result_bucket, + Key=self._result_key, + Body=json.dumps(result["dashboard_info"]), + ) + + return result + def _create_or_update_dashboard(self, dashboard_params: dict) -> tuple[str, str]: """ Creates new or updates existing template. diff --git a/tests/core/operation/test_publish_operation.py b/tests/core/operation/test_publish_operation.py index 316ce8f..9266e97 100644 --- a/tests/core/operation/test_publish_operation.py +++ b/tests/core/operation/test_publish_operation.py @@ -1,3 +1,7 @@ +import json +import os.path +import tempfile + import botocore from botocore.config import Config from botocore.stub import Stubber @@ -13,6 +17,9 @@ def test(self): template_id = f"{target_namespace}-library" account = "012345678910" group_name = "my_group" + result_bucket = "result-bucket" + result_key = "result-key" + output_json = tempfile.NamedTemporaryFile() boto_config = Config( region_name="us-east-1", @@ -22,14 +29,21 @@ def test(self): "quicksight", config=boto_config ) + s3_client = botocore.session.get_session().create_client( + "s3", config=boto_config + ) + + dashboard_arn = ( + "arn:aws:quicksight:us-west-2:128682227026:dashboard/tpp-prod-library" + ) describe_template_definition_params = { "AwsAccountId": account, "TemplateId": template_id, "AliasName": "$LATEST", } - with Stubber(qs_client) as stub: + with Stubber(qs_client) as qs_stub, Stubber(s3_client) as s3_stub: template_arn = f"arn:aws:quicksight:::template/{target_namespace}-library" - stub.add_response( + qs_stub.add_response( "describe_template", service_response={ "Template": { @@ -57,7 +71,7 @@ def test(self): namespace_arn = "arn:quicksight:::namespace/default" - stub.add_response( + qs_stub.add_response( "describe_namespace", service_response={"Namespace": {"Arn": namespace_arn}}, expected_params={ @@ -71,7 +85,7 @@ def test(self): ) ds2_arn = f"arn:aws:quicksight:::dataset/{target_namespace}-patron_events" - stub.add_response( + qs_stub.add_response( "describe_data_set", service_response={"DataSet": {"Arn": ds1_arn}}, expected_params={ @@ -79,7 +93,7 @@ def test(self): "DataSetId": f"{target_namespace}-circulation_view", }, ) - stub.add_response( + qs_stub.add_response( "describe_data_set", service_response={"DataSet": {"Arn": ds2_arn}}, expected_params={ @@ -88,7 +102,7 @@ def test(self): }, ) - stub.add_response( + qs_stub.add_response( "create_dashboard", service_response={ "ResponseMetadata": { @@ -103,7 +117,7 @@ def test(self): }, "RetryAttempts": 0, }, - "Arn": "arn:aws:quicksight:us-west-2:128682227026:dashboard/tpp-prod-library", + "Arn": dashboard_arn, "VersionArn": f"arn:aws:quicksight:::dashboard/{target_namespace}-library/version/6", "DashboardId": f"{target_namespace}-library", "CreationStatus": "CREATION_IN_PROGRESS", @@ -132,7 +146,7 @@ def test(self): }, ) group_arn = f"arn:aws:quicksight:::group/{group_name}" - stub.add_response( + qs_stub.add_response( "describe_group", service_response={"Group": {"Arn": group_arn}}, expected_params={ @@ -142,7 +156,7 @@ def test(self): }, ) - stub.add_response( + qs_stub.add_response( "update_dashboard_permissions", service_response={ "ResponseMetadata": { @@ -182,15 +196,37 @@ def test(self): ], }, ) + + s3_stub.add_response( + "put_object", + service_response={}, + expected_params={ + "Bucket": result_bucket, + "Key": result_key, + "Body": json.dumps({template_id: [dashboard_arn]}), + }, + ) + op = PublishDashboardFromTemplateOperation( qs_client=qs_client, + s3_client=s3_client, template_id=template_id, target_namespace=target_namespace, aws_account_id=account, group_name=group_name, + output_json=output_json.name, + result_bucket=result_bucket, + result_key=result_key, ) result = op.execute() assert result["status"] == "success" - assert result["dashboard_id"] == template_id + assert result["dashboard_info"] == {template_id: [dashboard_arn]} + assert os.path.exists(output_json.name) + + with open(output_json.name) as file: + result_from_file = json.loads(file.read()) + assert result_from_file["dashboard_info"] == { + template_id: [dashboard_arn] + }