| # 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|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]) |