Skip to content

Commit

Permalink
Add option to upload result to s3 (#49)
Browse files Browse the repository at this point in the history 
* 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.
  • Loading branch information
@dbernstein
dbernstein authored Nov 2, 2023
1 parent 07eecb7 commit 7e4d6fb
Show file tree
Hide file tree
Showing 3 changed files with 194 additions and 15 deletions.
119 changes: 118 additions & 1 deletion core/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
34 changes: 30 additions & 4 deletions core/operation/publish_dashboard_from_template.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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.
Expand Down
56 changes: 46 additions & 10 deletions tests/core/operation/test_publish_operation.py
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
import json
import os.path
import tempfile

import botocore
from botocore.config import Config
from botocore.stub import Stubber
Expand All @@ -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",
Expand All @@ -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": {
Expand Down Expand Up @@ -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={
Expand All @@ -71,15 +85,15 @@ 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={
"AwsAccountId": account,
"DataSetId": f"{target_namespace}-circulation_view",
},
)
stub.add_response(
qs_stub.add_response(
"describe_data_set",
service_response={"DataSet": {"Arn": ds2_arn}},
expected_params={
Expand All @@ -88,7 +102,7 @@ def test(self):
},
)

stub.add_response(
qs_stub.add_response(
"create_dashboard",
service_response={
"ResponseMetadata": {
Expand All @@ -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",
Expand Down Expand Up @@ -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={
Expand All @@ -142,7 +156,7 @@ def test(self):
},
)

stub.add_response(
qs_stub.add_response(
"update_dashboard_permissions",
service_response={
"ResponseMetadata": {
Expand Down Expand Up @@ -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]
}

0 comments on commit 7e4d6fb

Please sign in to comment.