recipe_modules/bigquery/api.py

# Copyright 2021 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

from RECIPE_MODULES.fuchsia.utils import cached_property


class BigQueryApi(recipe_api.RecipeApi):
    """APIs for BigQuery."""

    def __init__(self, *args, **kwargs):
        super(BigQueryApi, self).__init__(*args, **kwargs)
        self.project = None

    def query(self, step_name, sql, json_output):
        """Run a SQL query.

        Args:
            step_name (str): Name of the step.
            sql (str): The SQL query to run.
            json_output (Path): Path to JSON output.
        """
        assert self.project
        args = [
            "query",
            "-project",
            self.project,
            "-input",
            self.m.raw_io.input(sql),
            "-json-output",
            json_output,
        ]
        step = self._run(step_name, args)
        step.presentation.logs["query"] = sql
        return step

    @cached_property
    def _bigquery_tool(self):
        return self.m.ensure_tool("bigquery", self.resource("tool_manifest.json"))

    def _run(self, step_name, args):
        return self.m.step(step_name, [self._bigquery_tool] + args)