text: # Licensed to the Apache Software Foundation (ASF) under one or more
  # contributor license agreements.  See the NOTICE file distributed with
  # this work for additional information regarding copyright ownership.
  # The ASF licenses this file to You under the Apache License, Version 2.0
  # (the "License"); you may not use this file except in compliance with
  # the License.  You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """
    Helper for CD step.
  
    It is used to save beam examples/katas/tests and their output on the GCS.
  """
  import asyncio
  import json
  import logging
  import os
  import shutil
  from pathlib import Path
  from typing import List
  
  from tqdm import tqdm
  from google.cloud import storage
  
  from api.v1.api_pb2 import Sdk, SDK_PYTHON, SDK_JAVA, PrecompiledObjectType
  from config import Config, PrecompiledExample
  from grpc_client import GRPCClient
  from helper import Example, get_statuses
  
  
  class CDHelper:
    """
    Helper for CD step.
  
    It is used to save beam examples/katas/tests and their output on the GCS.
    """
    def store_examples(self, examples: List[Example]):
      """
      Store beam examples and their output in the Google Cloud.
  
      Outputs for multifile examples are left empty.
      """
      single_file_examples = list(filter(
        lambda example: example.tag.multifile is False, examples))
      logging.info("Start of executing only single-file Playground examples ...")
      asyncio.run(self._get_outputs(single_file_examples))
      logging.info("Finish of executing single-file Playground examples")
  
      logging.info("Start of sending Playground examples to the bucket ...")
      self._save_to_cloud_storage(examples)
      logging.info("Finish of sending Playground examples to the bucket")
      self._clear_temp_folder()
  
    async def _get_outputs(self, examples: List[Example]):
      """
      Run beam examples and keep their output.
  
      Call the backend to start code processing for the examples.
      Then receive code output.
  
      Args:
          examples: beam examples that should be run
      """
      await get_statuses(
          examples)  # run examples code and wait until all are executed
      client = GRPCClient()
      tasks = [client.get_run_output(example.pipeline_id) for example in examples]
      outputs = await asyncio.gather(*tasks)
  
      tasks = [client.get_log(example.pipeline_id) for example in examples]
      logs = await asyncio.gather(*tasks)
  
      if len(examples) > 0 and (examples[0].sdk is SDK_PYTHON or
                                examples[0].sdk is SDK_JAVA):
        tasks = [
            client.get_graph(example.pipeline_id, example.filepath)
            for example in examples
        ]
        graphs = await asyncio.gather(*tasks)
  
        for graph, example in zip(graphs, examples):
          example.graph = graph
  
      for output, example in zip(outputs, examples):
        example.output = output
  
      for log, example in zip(logs, examples):
        example.logs = log
  
    def _save_to_cloud_storage(self, examples: List[Example]):
      """
      Save examples, outputs and meta to bucket
  
      Args:
          examples: precompiled examples
      """
      self._storage_client = storage.Client()
      self._bucket = self._storage_client.bucket(Config.BUCKET_NAME)
      for example in tqdm(examples):
        file_names = self._write_to_local_fs(example)
        for cloud_file_name, local_file_name in file_names.items():
          self._upload_blob(
              source_file=local_file_name, destination_blob_name=cloud_file_name)
  
    def _write_to_local_fs(self, example: Example):
      """
      Write code of an example, output and meta info
      to the filesystem (in temp folder)
  
      Args:
          example: example object
  
      Returns: dict {path_at_the_bucket:path_at_the_os}
  
      """
      path_to_object_folder = os.path.join(
          Config.TEMP_FOLDER,
          example.pipeline_id,
          Sdk.Name(example.sdk),
          PrecompiledObjectType.Name(example.type),
          example.tag.name)
      Path(path_to_object_folder).mkdir(parents=True, exist_ok=True)
  
      file_names = {}
      code_path = self._get_gcs_object_name(
          sdk=example.sdk,
          type=example.type,
          base_folder_name=example.tag.name,
          file_name=example.tag.name)
      output_path = self._get_gcs_object_name(
          sdk=example.sdk,
          type=example.type,
          base_folder_name=example.tag.name,
          file_name=example.tag.name,
          extension=PrecompiledExample.OUTPUT_EXTENSION)
      log_path = self._get_gcs_object_name(
          sdk=example.sdk,
          type=example.type,
          base_folder_name=example.tag.name,
          file_name=example.tag.name,
          extension=PrecompiledExample.LOG_EXTENSION)
      graph_path = self._get_gcs_object_name(
          sdk=example.sdk,
          type=example.type,
          base_folder_name=example.tag.name,
          file_name=example.tag.name,
          extension=PrecompiledExample.GRAPH_EXTENSION)
      meta_path = self._get_gcs_object_name(
          sdk=example.sdk,
          type=example.type,
          base_folder_name=example.tag.name,
          file_name=PrecompiledExample.META_NAME,
          extension=PrecompiledExample.META_EXTENSION)
      file_names[code_path] = example.code
      file_names[output_path] = example.output
      meta = example.tag._asdict()
      meta["link"] = example.link
      file_names[meta_path] = json.dumps(meta)
      file_names[log_path] = example.logs
      if example.sdk == SDK_PYTHON or example.sdk == SDK_JAVA:
        file_names[graph_path] = example.graph
      for file_name, file_content in file_names.items():
        local_file_path = os.path.join(
            Config.TEMP_FOLDER, example.pipeline_id, file_name)
        with open(local_file_path, "w", encoding="utf-8") as file:
          file.write(file_content)
        # don't need content anymore, instead save the local path
        file_names[file_name] = local_file_path
      return file_names
  
    def _get_gcs_object_name(
        self,
        sdk: Sdk,
        type: PrecompiledObjectType,
        base_folder_name: str,
        file_name: str,
        extension: str = None):
      """
      Get the path where file will be stored at the bucket.
  
      Args:
        sdk: sdk of the example
        type: type of the example
        file_name: name of the example
        base_folder_name: name of the folder where example is stored
          (eq. to example name)
        extension: extension of the file
  
      Returns: file name
      """
      if extension is None:
        extension = Config.SDK_TO_EXTENSION[sdk]
      return os.path.join(
          Sdk.Name(sdk),
          PrecompiledObjectType.Name(type),
          base_folder_name,
          f"{file_name}.{extension}")
  
    def _upload_blob(self, source_file: str, destination_blob_name: str):
      """
      Upload a file to the bucket.
  
      Args:
          source_file: name of the file to be stored
          destination_blob_name: "storage-object-name"
      """
  
      blob = self._bucket.blob(destination_blob_name)
      blob.upload_from_filename(source_file)
      # change caching to no caching
      blob.cache_control = Config.NO_STORE
      blob.patch()
      logging.info("File uploaded to %s", destination_blob_name)
  
    def _clear_temp_folder(self):
      """
      Remove temporary folder with source files.
      """
      if os.path.exists(Config.TEMP_FOLDER):
        shutil.rmtree(Config.TEMP_FOLDER)
  