name: get_vertex_dataset
description: |
  Resolves a dataset artifact to an existing Vertex Dataset.
  Args:
    dataset_resource_name (str):
      Required. Vertex Dataset resource name in the format of projects/{project}/locations/{location}/datasets/{dataset}
  Returns:
    dataset (Optional[google.VertexDataset]):
      Vertex Dataset artifact with a resourceName attribute in the format of projects/{project}/locations/{location}/datasets/{dataset}
inputs:
  - {name: dataset_resource_name, type: String}
outputs:
  - {name: dataset, type: google.VertexDataset}
  - {name: gcp_resources, type: String}
implementation:
  container:
    image: python:3.8
    command:
    - sh
    - -ec
    - |
      program_path=$(mktemp)
      printf "%s" "$0" > "$program_path"
      python3 -u "$program_path" "$@"
    - |2

      import os

      def get_vertex_dataset(executor_input_json, dataset_resource_name: str):
          executor_output = {}
          executor_output['artifacts'] = {}
          for name, artifacts in executor_input_json.get('outputs',
                                                        {}).get('artifacts',
                                                                {}).items():
            artifacts_list = artifacts.get('artifacts')
            if name == 'dataset' and artifacts_list:
              updated_runtime_artifact = artifacts_list[0]
              updated_runtime_artifact['uri'] = dataset_resource_name
              updated_runtime_artifact['metadata'] = {'resourceName': dataset_resource_name}
              artifacts_list = {'artifacts': [updated_runtime_artifact]}

            executor_output['artifacts'][name] = artifacts_list

          # Update the output artifacts.
          os.makedirs(
              os.path.dirname(executor_input_json['outputs']['outputFile']),
              exist_ok=True)
          with open(executor_input_json['outputs']['outputFile'], 'w') as f:
            f.write(json.dumps(executor_output))

      if __name__ == '__main__':
        import argparse
        import json

        parser = argparse.ArgumentParser(description='get resource name')
        parser.add_argument('--dataset_resource_name', type=str)
        parser.add_argument('--executor_input', type=str)
        args, _ = parser.parse_known_args()
        executor_input_json = json.loads(args.executor_input)
        dataset_resource_name = args.dataset_resource_name
        get_vertex_dataset(executor_input_json, dataset_resource_name)
    args:
      - --dataset_resource_name
      - {inputValue: dataset_resource_name}
      - --gcp_resources
      - {outputPath: gcp_resources}
      - --executor_input
      - "{{$}}"