vertexai_workflow_clone_bigquery_table.yaml

main:
  params: [args]
  steps:
    - initialize:
        assign:
          - location: ${default(map.get(args,"location"),"us-central1")}
          - trainingPipelineId: ${default(map.get(args,"trainingPipelineId"),"3590189825883373568")}
          - destination_dataset: "vertexai_model_history"
    - getTrainingPipeline:
        call: VertexAI_GetTrainingPipeline
        args:
          location: ${location}
          trainingPipelineId: ${trainingPipelineId}
        result: trainingPipeline
    - getDataset:
        call: VertexAI_GetDataset
        args:
          location: ${location}
          datasetId: ${trainingPipeline.inputDataConfig.datasetId}
        result: dataset
    - bigquerySourceCheck:
        switch:
          - condition: ${not("bigquerySource" in dataset.metadata.inputConfig)}
            next: end
    # proceed further only when we have the bigquerySource in dataset.metadata.inputConfig
    - prepareQueryForClone:
        assign:
          - basetable: ${text.replace_all(dataset.metadata.inputConfig.bigquerySource.uri,"bq://","")}
          - basetable_elements: ${text.split(basetable,".")}
          - pipeline_elements: ${text.split(trainingPipeline.name,"/")}
          - newtable: ${basetable_elements[0]+"."+destination_dataset+"."+basetable_elements[2]+"_"+pipeline_elements[len(pipeline_elements)-1]}
          - query:
              '${"CREATE TABLE IF NOT EXISTS "+newtable+"
              CLONE "+basetable+"
              OPTIONS ( description=\"Clone of "+basetable+" for Vertex AI Training of "+trainingPipeline.displayName+" at time of "+trainingPipeline.updateTime+"\")"}'
    - cloneBigQueryTable:
        call: BQ_Query
        args:
          query: ${query}
    - final:
        return: ${query}
BQ_Query:
  params: [query]
  steps:
    - runBQquery:
        call: googleapis.bigquery.v2.jobs.query
        args:
          projectId: ${sys.get_env("GOOGLE_CLOUD_PROJECT_ID")}
          body:
            useLegacySql: false
            query: ${query}
        result: queryResult
    - documentFound:
        return: ${queryResult}
VertexAI_GetDataset:
  params: [location, datasetId]
  steps:
    - runGetPipeline:
        call: http.get
        args:
          url: ${"https://"+location+"-aiplatform.googleapis.com/v1/projects/"+sys.get_env("GOOGLE_CLOUD_PROJECT_ID")+"/locations/"+location+"/datasets/"+datasetId}
          auth:
            type: OAuth2
        result: dataset
    - documentFound:
        return: ${dataset.body}
VertexAI_GetTrainingPipeline:
  params: [location, trainingPipelineId]
  steps:
    - runGetPipeline:
        call: http.get
        args:
          url: ${"https://"+location+"-aiplatform.googleapis.com/v1/projects/"+sys.get_env("GOOGLE_CLOUD_PROJECT_ID")+"/locations/"+location+"/trainingPipelines/"+trainingPipelineId}
          auth:
            type: OAuth2
        result: trainingPipeline
    - documentFound:
        return: ${trainingPipeline.body}