# Copyright 2021 The Kubeflow Authors
#
# Licensed 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.

name: bigquery_export_model_job
description: |
  Launch a BigQuery export model job and waits for it to finish.

    Args:
        project (str):
            Required. Project to run BigQuery model export job.
        location (Optional[str]):
            Location of the job to export the BigQuery model. If not set,
            default to `US` multi-region.

            For more details, see https://cloud.google.com/bigquery/docs/locations#specifying_your_location
        model (google.BQMLModel):
            Required. BigQuery ML model to export.
        model_destination_path(str):
            Required. The gcs bucket to export the model to.
        job_configuration_extract (Optional[dict]):
            A json formatted string describing the rest of the job configuration.

            For more details, see https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#JobConfigurationQuery
        labels (Optional[dict]):
            The labels associated with this job. You can use these to organize and group your jobs.
            Label keys and values can be no longer than 63 characters, can only containlowercase
            letters, numeric characters, underscores and dashes. International characters are
            allowed. Label values are optional. Label keys must start with a letter and each label
            in the list must have a different key.

            Example: { "name": "wrench", "mass": "1.3kg", "count": "3" }.
    Returns:
        gcp_resources (str):
            Serialized gcp_resources proto tracking the BigQuery job.
            For more details, see https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/proto/README.md.
inputs:
- {name: project, type: String}
- {name: location, type: String, default: "us-central1"}
- {name: model, type: google.BQMLModel}
- {name: model_destination_path, type: String}
- {name: job_configuration_extract, type: JsonObject, optional: true, default: '{}'}
- {name: labels, type: JsonObject, optional: true, default: '{}'}
outputs:
- {name: exported_model_path, type: String}
- {name: gcp_resources, type: String}
implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:1.10.13
    command: [python3, -u, -m, google_cloud_pipeline_components.container.v1.gcp_launcher.launcher]
    args: [
      --type, BigqueryExportModelJob,
      --project, {inputValue: project},
      --location, {inputValue: location},
      --model_name,
      concat: [
          "{{$.inputs.artifacts['model'].metadata['projectId']}}",
          '.',
          "{{$.inputs.artifacts['model'].metadata['datasetId']}}",
          '.',
          "{{$.inputs.artifacts['model'].metadata['modelId']}}"
          ],
      --model_destination_path, {inputValue: model_destination_path},
      --payload,
      concat: [
          '{',
            '"configuration": {',
              '"extract": ', {inputValue: job_configuration_extract},
              ', "labels": ', {inputValue: labels},
            '}',
          '}'
      ],
      --exported_model_path, {outputPath: exported_model_path},
      --gcp_resources, {outputPath: gcp_resources},
      --executor_input, "{{$}}",
    ]
