From 30f7d8ac808b8c53a849baeb60fa2ee83a3375b3 Mon Sep 17 00:00:00 2001 From: Niels Vanspauwen Date: Thu, 15 Jun 2023 20:40:33 +0200 Subject: [PATCH] Add support for uploading to GCS --- .gitignore | 3 ++ pystructurizr/cli.py | 67 +++++++++++++++++++------------------ pystructurizr/cli_helper.py | 30 +++++++++++++++++ requirements.txt | 1 + 4 files changed, 68 insertions(+), 33 deletions(-) diff --git a/.gitignore b/.gitignore index a964bcc..29aa58a 100644 --- a/.gitignore +++ b/.gitignore @@ -13,3 +13,6 @@ build/ # Miscellaneous .DS_Store + +# Credentials +gcs_credentials.json diff --git a/pystructurizr/cli.py b/pystructurizr/cli.py index e80ae61..59426c1 100644 --- a/pystructurizr/cli.py +++ b/pystructurizr/cli.py @@ -6,7 +6,9 @@ import os import shutil import subprocess -from .cli_helper import generate_diagram_code, generate_diagram_code_in_child_process +from google.cloud import storage +from google.cloud.exceptions import GoogleCloudError +from .cli_helper import generate_diagram_code, generate_diagram_code_in_child_process, generate_svg, ensure_tmp_folder_exists from .cli_watcher import observe_modules @@ -15,13 +17,10 @@ @click.option('--view', prompt='Your view file (e.g. example.componentview)', help='The view file to generate.') def dump(view): - initial_modules = set(sys.modules.keys()) diagram_code = generate_diagram_code(view) - current_modules = set(sys.modules.keys()) - # print("Imported modules:") - # print(current_modules - initial_modules) click.echo(diagram_code) - + + @click.command() @click.option('--view', prompt='Your view file (e.g. example.componentview)', @@ -29,29 +28,15 @@ def dump(view): def dev(view): click.echo(f"Setting up live preview of view {view}...") # Prep the /tmp/pystructurizr folder - tmp_folder = "/tmp/pystructurizr" - os.makedirs(tmp_folder, exist_ok=True) + tmp_folder = ensure_tmp_folder_exists() current_script_path = os.path.abspath(__file__) index_html = os.path.join(os.path.dirname(current_script_path), 'index.html') - print(index_html) shutil.copy(index_html, f"{tmp_folder}/index.html") async def async_behavior(): print("Generating diagram...") diagram_code = generate_diagram_code_in_child_process(view) - # print(diagram_code) - url = f"https://kroki.io/structurizr/svg" - - async with httpx.AsyncClient() as client: - resp = await client.post(url, data=diagram_code) - if resp.status_code != 200: - print(resp) - if resp.content: - print(resp.content.decode()) - raise click.ClickException("Failed to create diagram") - async with aiofiles.open(f"{tmp_folder}/diagram.svg", "w") as f: - await f.write(resp.text) - print(f"Updated SVG in {tmp_folder}") + await generate_svg(diagram_code, tmp_folder) async def observe_loop(): await async_behavior() @@ -62,25 +47,41 @@ async def observe_loop(): asyncio.run(observe_loop()) + @click.command() @click.option('--view', prompt='Your view file (e.g. example.componentview)', - help='The view file to generate and upload to S3.') -def build(view): - click.echo(f"Generating view {view} and uploading to S3...") - diagram_code = generate_diagram_code(view) + help='The view file to generate and upload to cloud storage.') +@click.option('--gcs-credentials', prompt='Path to json file containing Google Cloud Storage credentials', type=click.Path(exists=True), + help='Path to the credentials.json file for Google Cloud Storage.') +@click.option('--bucket-name', prompt='Name of the bucket on Google Cloud Storage', + help='The name of the bucket to use on Google Cloud Storage.') +@click.option('--object-name', prompt='Name of the object on Google Cloud Storage', + help='The name of the object to use on Google Cloud Storage.') +def build(view, gcs_credentials, bucket_name, object_name): + click.echo(f"Generating view {view} and uploading to cloud storage...") async def async_behavior(): - async with httpx.AsyncClient() as client: - resp = await client.post(url, data=diagram_code) - if resp.status_code != 200: - print(resp) - raise click.ClickException("Failed to create diagram") - - # TODO: take resp.text and upload it to an AVG file in an S3 bucket + print("Generating diagram...") + diagram_code = generate_diagram_code(view) + tmp_folder = ensure_tmp_folder_exists() + svg_file_path = await generate_svg(diagram_code, tmp_folder) + try: + print(f"Uploading {svg_file_path}...") + gcs_client = storage.Client.from_service_account_json(gcs_credentials) + gcs_bucket = gcs_client.get_bucket(bucket_name) + gcs_blob = gcs_bucket.blob(object_name) + gcs_blob.upload_from_filename(svg_file_path) + svg_file_url = f"https://storage.googleapis.com/{bucket_name}/{object_name}" + print("Done! You can get the SVG file at:") + print(svg_file_url) + except GoogleCloudError as e: + print("An error occurred while uploading the file to Google Cloud Storage:") + print(e) asyncio.run(async_behavior()) + @click.group() def cli(): pass diff --git a/pystructurizr/cli_helper.py b/pystructurizr/cli_helper.py index afa17b6..e47cc3b 100644 --- a/pystructurizr/cli_helper.py +++ b/pystructurizr/cli_helper.py @@ -2,8 +2,13 @@ import importlib import subprocess import sys +import aiofiles +import click +import httpx +import sys import os + # This is useally run within the same process as cli.py # But for the 'dev' command, we run it as a separate process, so it can be run # repeatedly within the same dev session whenever the view changes. @@ -26,3 +31,28 @@ def run_child_process(): # Run the child process and capture its output child_output = run_child_process() return child_output + + +async def generate_svg(diagram_code, tmp_folder): + url = f"https://kroki.io/structurizr/svg" + async with httpx.AsyncClient() as client: + resp = await client.post(url, data=diagram_code) + + if resp.status_code != 200: + print(resp) + if resp.content: + print(resp.content.decode()) + raise click.ClickException("Failed to create diagram") + + svg_file_path = f"{tmp_folder}/diagram.svg" + async with aiofiles.open(svg_file_path, "w") as f: + await f.write(resp.text) + print(f"Updated SVG in {tmp_folder}") + + return svg_file_path + + +def ensure_tmp_folder_exists(): + tmp_folder = "/tmp/pystructurizr" + os.makedirs(tmp_folder, exist_ok=True) + return tmp_folder diff --git a/requirements.txt b/requirements.txt index 8fd32b8..d6d1c5a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ click httpx aiofiles httpwatcher +google-cloud-storage