recipe_modules/bqupload/api.py

# Copyright 2019 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

from recipe_engine import recipe_api

BQUPLOAD_CIPD_VERSION = \
  'git_revision:d85fe78f303c3e969f815121e17c8b08868039ef'


class BquploadApi(recipe_api.RecipeApi):
  """BquploadApi provides support for BigQuery upload."""

  def __init__(self, *args, **kwargs):
    super(BquploadApi, self).__init__(*args, **kwargs)
    self._bqupload_tool = None

  def _ensure_bqupload(self):
    """Ensures that the bqupload Go binary is installed."""
    if self._bqupload_tool:
      return  # pragma: no cover

    with self.m.step.nest('ensure bqupload'):
      with self.m.context(infra_steps=True):
        pkgs = self.m.cipd.EnsureFile()
        pkgs.add_package('infra/tools/bqupload/${platform}',
                         BQUPLOAD_CIPD_VERSION)
        cipd_dir = self.m.path['start_dir'].join('cipd', 'bqupload')
        self.m.cipd.ensure(cipd_dir, pkgs)
        self._bqupload_tool = cipd_dir.join('bqupload')

  def _run(self, name, cmd, step_test_data=None):
    """Returns a bqupload command step."""
    self._ensure_bqupload()
    return self.m.step(
        name, [self._bqupload_tool] + list(cmd), step_test_data=step_test_data)

  def insert(self, step_name, project, dataset, table, data_file):
    """Inserts a data file into a BigQuery table.

    Args:
      step_name (str): name of the step.
      project (str): id of the gcloud project.
      dataset (str): name of the BigQuery dataset.
      table (str): name of the BigQuery table.
      data_file (Path or Placeholder): path to the json file that contains
        data to insert into a BigQuery table.
    """
    table_path = '%s.%s.%s' % (project, dataset, table)
    return self._run(step_name, [table_path, data_file])