garf-executors 0.1.4__py3-none-any.whl → 1.0.2__py3-none-any.whl

garf/executors/__init__.py

@@ -0,0 +1,60 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Executors to fetch data from various APIs."""
+
+from __future__ import annotations
+
+import importlib
+
+from garf.executors import executor, fetchers
+from garf.executors.api_executor import ApiExecutionContext, ApiQueryExecutor
+from garf.executors.telemetry import tracer
+
+
+@tracer.start_as_current_span('setup_executor')
+def setup_executor(
+  source: str,
+  fetcher_parameters: dict[str, str | int | bool],
+  enable_cache: bool = False,
+  cache_ttl_seconds: int = 3600,
+) -> type[executor.Executor]:
+  """Initializes executors based on a source and parameters."""
+  if source == 'bq':
+    bq_executor = importlib.import_module('garf.executors.bq_executor')
+    query_executor = bq_executor.BigQueryExecutor(**fetcher_parameters)
+  elif source == 'sqldb':
+    sql_executor = importlib.import_module('garf.executors.sql_executor')
+    query_executor = (
+      sql_executor.SqlAlchemyQueryExecutor.from_connection_string(
+        fetcher_parameters.get('connection_string')
+      )
+    )
+  else:
+    concrete_api_fetcher = fetchers.get_report_fetcher(source)
+    query_executor = ApiQueryExecutor(
+      fetcher=concrete_api_fetcher(
+        **fetcher_parameters,
+        enable_cache=enable_cache,
+        cache_ttl_seconds=cache_ttl_seconds,
+      )
+    )
+  return query_executor
+
+
+__all__ = [
+  'ApiQueryExecutor',
+  'ApiExecutionContext',
+]
+
+__version__ = '1.0.2'

garf/executors/api_executor.py

@@ -0,0 +1,143 @@
+# Copyright 2024 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Module for executing Garf queries and writing them to local/remote.
+
+ApiQueryExecutor performs fetching data from API in a form of
+GarfReport and saving it to local/remote storage.
+"""
+# pylint: disable=C0330, g-bad-import-order, g-multiple-import
+
+from __future__ import annotations
+
+import logging
+
+from garf.core import report_fetcher
+from garf.executors import (
+  exceptions,
+  execution_context,
+  executor,
+  fetchers,
+  query_processor,
+)
+from garf.executors.telemetry import tracer
+from opentelemetry import trace
+
+logger = logging.getLogger(__name__)
+
+
+class ApiExecutionContext(execution_context.ExecutionContext):
+  """Common context for executing one or more queries."""
+
+  writer: str | list[str] = 'console'
+
+
+class ApiQueryExecutor(executor.Executor):
+  """Gets data from API and writes them to local/remote storage.
+
+  Attributes:
+      api_client: a client used for connecting to API.
+  """
+
+  def __init__(self, fetcher: report_fetcher.ApiReportFetcher) -> None:
+    """Initializes ApiQueryExecutor.
+
+    Args:
+        fetcher: Instantiated report fetcher.
+    """
+    self.fetcher = fetcher
+    super().__init__(
+      preprocessors=self.fetcher.preprocessors,
+      postprocessors=self.fetcher.postprocessors,
+    )
+
+  @classmethod
+  def from_fetcher_alias(
+    cls,
+    source: str,
+    fetcher_parameters: dict[str, str] | None = None,
+    enable_cache: bool = False,
+    cache_ttl_seconds: int = 3600,
+  ) -> ApiQueryExecutor:
+    if not fetcher_parameters:
+      fetcher_parameters = {}
+    concrete_api_fetcher = fetchers.get_report_fetcher(source)
+    return ApiQueryExecutor(
+      fetcher=concrete_api_fetcher(
+        **fetcher_parameters,
+        enable_cache=enable_cache,
+        cache_ttl_seconds=cache_ttl_seconds,
+      )
+    )
+
+  @tracer.start_as_current_span('api.execute')
+  def execute(
+    self,
+    query: str,
+    title: str,
+    context: ApiExecutionContext,
+  ) -> str:
+    """Reads query, extract results and stores them in a specified location.
+
+    Args:
+      query: Location of the query.
+      title: Name of the query.
+      context: Query execution context.
+
+    Returns:
+      Result of writing the report.
+
+    Raises:
+      GarfExecutorError: When failed to execute query.
+    """
+    context = query_processor.process_gquery(context)
+    span = trace.get_current_span()
+    span.set_attribute('fetcher.class', self.fetcher.__class__.__name__)
+    span.set_attribute(
+      'api.client.class', self.fetcher.api_client.__class__.__name__
+    )
+    try:
+      span.set_attribute('query.title', title)
+      span.set_attribute('query.text', query)
+      logger.debug('starting query %s', query)
+      results = self.fetcher.fetch(
+        query_specification=query,
+        args=context.query_parameters,
+        **context.fetcher_parameters,
+      )
+      writer_clients = context.writer_clients
+      if not writer_clients:
+        logger.warning('No writers configured, skipping write operation')
+        return None
+      writing_results = []
+      for writer_client in writer_clients:
+        logger.debug(
+          'Start writing data for query %s via %s writer',
+          title,
+          type(writer_client),
+        )
+        result = writer_client.write(results, title)
+        logger.debug(
+          'Finish writing data for query %s via %s writer',
+          title,
+          type(writer_client),
+        )
+        writing_results.append(result)
+      logger.info('%s executed successfully', title)
+      # Return the last writer's result for backward compatibility
+      return writing_results[-1] if writing_results else None
+    except Exception as e:
+      logger.error('%s generated an exception: %s', title, str(e))
+      raise exceptions.GarfExecutorError(
+        '%s generated an exception: %s', title, str(e)
+      ) from e

garf/executors/bq_executor.py

@@ -0,0 +1,177 @@
+# Copyright 2024 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Executes queries in BigQuery."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+
+try:
+  from google.cloud import bigquery  # type: ignore
+except ImportError as e:
+  raise ImportError(
+    'Please install garf-executors with BigQuery support '
+    '- `pip install garf-executors[bq]`'
+  ) from e
+
+import logging
+
+from garf.core import query_editor, report
+from garf.executors import exceptions, execution_context, executor
+from garf.executors.telemetry import tracer
+from google.cloud import exceptions as google_cloud_exceptions
+from opentelemetry import trace
+
+logger = logging.getLogger(__name__)
+
+
+class BigQueryExecutorError(exceptions.GarfExecutorError):
+  """Error when BigQueryExecutor fails to run query."""
+
+
+class BigQueryExecutor(executor.Executor, query_editor.TemplateProcessorMixin):
+  """Handles query execution in BigQuery.
+
+  Attributes:
+      project_id: Google Cloud project id.
+      location: BigQuery dataset location.
+      client: BigQuery client.
+  """
+
+  def __init__(
+    self,
+    project_id: str | None = os.getenv('GOOGLE_CLOUD_PROJECT'),
+    location: str | None = None,
+    **kwargs: str,
+  ) -> None:
+    """Initializes BigQueryExecutor.
+
+    Args:
+        project_id: Google Cloud project id.
+        location: BigQuery dataset location.
+    """
+    if not project_id:
+      raise BigQueryExecutorError(
+        'project_id is required. Either provide it as project_id parameter '
+        'or GOOGLE_CLOUD_PROJECT env variable.'
+      )
+    self.project_id = project_id
+    self.location = location
+    super().__init__()
+
+  @property
+  def client(self) -> bigquery.Client:
+    """Instantiates bigquery client."""
+    return bigquery.Client(self.project_id)
+
+  @tracer.start_as_current_span('bq.execute')
+  def execute(
+    self,
+    query: str,
+    title: str,
+    context: execution_context.ExecutionContext = (
+      execution_context.ExecutionContext()
+    ),
+  ) -> report.GarfReport:
+    """Executes query in BigQuery.
+
+    Args:
+      query: Location of the query.
+      title: Name of the query.
+      context: Query execution context.
+
+    Returns:
+      Report with data if query returns some data otherwise empty Report.
+    """
+    span = trace.get_current_span()
+    logger.info('Executing script: %s', title)
+    query_text = self.replace_params_template(query, context.query_parameters)
+    self.create_datasets(context.query_parameters.macro)
+    job = self.client.query(query_text)
+    try:
+      result = job.result()
+    except google_cloud_exceptions.GoogleCloudError as e:
+      raise BigQueryExecutorError(
+        f'Failed to execute query {title}: Reason: {e}'
+      ) from e
+      logger.debug('%s launched successfully', title)
+    if result.total_rows:
+      results = report.GarfReport.from_pandas(result.to_dataframe())
+    else:
+      results = report.GarfReport()
+    if context.writer and results:
+      writer_clients = context.writer_clients
+      if not writer_clients:
+        logger.warning('No writers configured, skipping write operation')
+      else:
+        writing_results = []
+        for writer_client in writer_clients:
+          logger.debug(
+            'Start writing data for query %s via %s writer',
+            title,
+            type(writer_client),
+          )
+          writing_result = writer_client.write(results, title)
+          logger.debug(
+            'Finish writing data for query %s via %s writer',
+            title,
+            type(writer_client),
+          )
+          writing_results.append(writing_result)
+        # Return the last writer's result for backward compatibility
+        logger.info('%s executed successfully', title)
+        return writing_results[-1] if writing_results else None
+    logger.info('%s executed successfully', title)
+    span.set_attribute('execute.num_results', len(results))
+    return results
+
+  @tracer.start_as_current_span('bq.create_datasets')
+  def create_datasets(self, macros: dict | None) -> None:
+    """Creates datasets in BQ based on values in a dict.
+
+    If dict contains keys with 'dataset' in them, then values for such keys
+    are treated as dataset names.
+
+    Args:
+      macros: Mapping containing data for query execution.
+    """
+    if macros and (datasets := extract_datasets(macros)):
+      for dataset in datasets:
+        dataset_id = f'{self.project_id}.{dataset}'
+        try:
+          self.client.get_dataset(dataset_id)
+        except google_cloud_exceptions.NotFound:
+          bq_dataset = bigquery.Dataset(dataset_id)
+          bq_dataset.location = self.location
+          with contextlib.suppress(google_cloud_exceptions.Conflict):
+            self.client.create_dataset(bq_dataset, timeout=30)
+            logger.info('Created new dataset %s', dataset_id)
+
+
+def extract_datasets(macros: dict | None) -> list[str]:
+  """Finds dataset-related keys based on values in a dict.
+
+  If dict contains keys with 'dataset' in them, then values for such keys
+  are treated as dataset names.
+
+  Args:
+      macros: Mapping containing data for query execution.
+
+  Returns:
+      Possible names of datasets.
+  """
+  if not macros:
+    return []
+  return [value for macro, value in macros.items() if 'dataset' in macro]

garf/executors/config.py

@@ -0,0 +1,52 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+#     https://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.
+
+# pylint: disable=C0330, g-bad-import-order, g-multiple-import
+
+"""Stores mapping between API aliases and their execution context."""
+
+from __future__ import annotations
+
+import os
+import pathlib
+
+import pydantic
+import smart_open
+import yaml
+from garf.executors.execution_context import ExecutionContext
+
+
+class Config(pydantic.BaseModel):
+  """Stores necessary parameters for one or multiple API sources.
+
+  Attributes:
+    source: Mapping between API source alias and execution parameters.
+  """
+
+  sources: dict[str, ExecutionContext]
+
+  @classmethod
+  def from_file(cls, path: str | pathlib.Path | os.PathLike[str]) -> Config:
+    """Builds config from local or remote yaml file."""
+    with smart_open.open(path, 'r', encoding='utf-8') as f:
+      data = yaml.safe_load(f)
+    return Config(sources=data)
+
+  def save(self, path: str | pathlib.Path | os.PathLike[str]) -> str:
+    """Saves config to local or remote yaml file."""
+    with smart_open.open(path, 'w', encoding='utf-8') as f:
+      yaml.dump(
+        self.model_dump(exclude_none=True).get('sources'), f, encoding='utf-8'
+      )
+    return f'Config is saved to {str(path)}'

garf/executors/entrypoints/cli.py

@@ -0,0 +1,177 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+#     https://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.
+"""Module for defining `garf` CLI utility.
+
+`garf` allows to execute queries and store results in local/remote
+storage.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import pathlib
+import sys
+
+import garf.executors
+from garf.executors import config, exceptions, workflow
+from garf.executors.entrypoints import utils
+from garf.executors.entrypoints.tracer import initialize_tracer
+from garf.executors.telemetry import tracer
+from garf.io import reader
+from opentelemetry import trace
+
+initialize_tracer()
+
+
+@tracer.start_as_current_span('garf.entrypoints.cli')
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument('query', nargs='*')
+  parser.add_argument('-c', '--config', dest='config', default=None)
+  parser.add_argument('-w', '--workflow', dest='workflow', default=None)
+  parser.add_argument('--source', dest='source', default=None)
+  parser.add_argument('--output', dest='output', default='console')
+  parser.add_argument('--input', dest='input', default='file')
+  parser.add_argument('--log', '--loglevel', dest='loglevel', default='info')
+  parser.add_argument('--logger', dest='logger', default='local')
+  parser.add_argument('--log-name', dest='log_name', default='garf')
+  parser.add_argument(
+    '--parallel-queries', dest='parallel_queries', action='store_true'
+  )
+  parser.add_argument(
+    '--no-parallel-queries', dest='parallel_queries', action='store_false'
+  )
+  parser.add_argument('--dry-run', dest='dry_run', action='store_true')
+  parser.add_argument('-v', '--version', dest='version', action='store_true')
+  parser.add_argument(
+    '--parallel-threshold', dest='parallel_threshold', default=10, type=int
+  )
+  parser.add_argument(
+    '--enable-cache', dest='enable_cache', action='store_true'
+  )
+  parser.add_argument(
+    '--cache-ttl-seconds',
+    dest='cache_ttl_seconds',
+    default=3600,
+    type=int,
+  )
+  parser.set_defaults(parallel_queries=True)
+  parser.set_defaults(enable_cache=False)
+  parser.set_defaults(dry_run=False)
+  args, kwargs = parser.parse_known_args()
+
+  span = trace.get_current_span()
+  command_args = ' '.join(sys.argv[1:])
+  span.set_attribute('cli.command', f'garf {command_args}')
+  if args.version:
+    print(garf.executors.__version__)
+    sys.exit()
+  logger = utils.init_logging(
+    loglevel=args.loglevel.upper(), logger_type=args.logger, name=args.log_name
+  )
+  reader_client = reader.create_reader(args.input)
+  if workflow_file := args.workflow:
+    wf_parent = pathlib.Path.cwd() / pathlib.Path(workflow_file).parent
+    execution_workflow = workflow.Workflow.from_file(workflow_file)
+    for i, step in enumerate(execution_workflow.steps, 1):
+      with tracer.start_as_current_span(f'{i}-{step.fetcher}'):
+        query_executor = garf.executors.setup_executor(
+          source=step.fetcher,
+          fetcher_parameters=step.fetcher_parameters,
+          enable_cache=args.enable_cache,
+          cache_ttl_seconds=args.cache_ttl_seconds,
+        )
+        batch = {}
+        if not (queries := step.queries):
+          logger.error('Please provide one or more queries to run')
+          raise exceptions.GarfExecutorError(
+            'Please provide one or more queries to run'
+          )
+        for query in queries:
+          if isinstance(query, garf.executors.workflow.QueryPath):
+            query_path = wf_parent / pathlib.Path(query.path)
+            if not query_path.exists():
+              raise workflow.GarfWorkflowError(f'Query: {query_path} not found')
+            batch[query.path] = reader_client.read(query_path)
+          elif isinstance(query, garf.executors.workflow.QueryFolder):
+            query_path = wf_parent / pathlib.Path(query.folder)
+            if not query_path.exists():
+              raise workflow.GarfWorkflowError(
+                f'Folder: {query_path} not found'
+              )
+            for p in query_path.rglob('*'):
+              if p.suffix == '.sql':
+                batch[p.stem] = reader_client.read(p)
+          else:
+            batch[query.query.title] = query.query.text
+        query_executor.execute_batch(
+          batch, step.context, args.parallel_threshold
+        )
+    sys.exit()
+
+  if not args.query:
+    logger.error('Please provide one or more queries to run')
+    raise exceptions.GarfExecutorError(
+      'Please provide one or more queries to run'
+    )
+  if config_file := args.config:
+    execution_config = config.Config.from_file(config_file)
+    if not (context := execution_config.sources.get(args.source)):
+      raise exceptions.GarfExecutorError(
+        f'No execution context found for source {args.source} in {config_file}'
+      )
+  else:
+    param_types = ['source', 'macro', 'template']
+    outputs = args.output.split(',')
+    extra_parameters = utils.ParamsParser([*param_types, *outputs]).parse(
+      kwargs
+    )
+    source_parameters = extra_parameters.get('source', {})
+    writer_parameters = {}
+    for output in outputs:
+      writer_parameters.update(extra_parameters.get(output))
+
+    context = garf.executors.api_executor.ApiExecutionContext(
+      query_parameters={
+        'macro': extra_parameters.get('macro'),
+        'template': extra_parameters.get('template'),
+      },
+      writer=outputs,
+      writer_parameters=writer_parameters,
+      fetcher_parameters=source_parameters,
+    )
+  query_executor = garf.executors.setup_executor(
+    source=args.source,
+    fetcher_parameters=context.fetcher_parameters,
+    enable_cache=args.enable_cache,
+    cache_ttl_seconds=args.cache_ttl_seconds,
+  )
+  batch = {query: reader_client.read(query) for query in args.query}
+  if args.parallel_queries and len(args.query) > 1:
+    logger.info('Running queries in parallel')
+    batch = {query: reader_client.read(query) for query in args.query}
+    query_executor.execute_batch(batch, context, args.parallel_threshold)
+  else:
+    if len(args.query) > 1:
+      logger.info('Running queries sequentially')
+    for query in args.query:
+      query_executor.execute(
+        query=reader_client.read(query), title=query, context=context
+      )
+  logging.shutdown()
+
+
+if __name__ == '__main__':
+  main()

garf/executors/entrypoints/grpc_server.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+#     https://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.
+
+"""gRPC endpoint for garf."""
+
+import argparse
+import logging
+from concurrent import futures
+
+import garf.executors
+import grpc
+from garf.executors import garf_pb2, garf_pb2_grpc
+from garf.executors.entrypoints.tracer import initialize_tracer
+from google.protobuf.json_format import MessageToDict
+from grpc_reflection.v1alpha import reflection
+
+
+class GarfService(garf_pb2_grpc.GarfService):
+  def Execute(self, request, context):
+    query_executor = garf.executors.setup_executor(
+      request.source, request.context.fetcher_parameters
+    )
+    execution_context = garf.executors.execution_context.ExecutionContext(
+      **MessageToDict(request.context, preserving_proto_field_name=True)
+    )
+    result = query_executor.execute(
+      query=request.query,
+      title=request.title,
+      context=execution_context,
+    )
+    return garf_pb2.ExecuteResponse(results=[result])
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--port', dest='port', default=50051, type=int)
+  parser.add_argument(
+    '--parallel-threshold', dest='parallel_threshold', default=10, type=int
+  )
+  args, _ = parser.parse_known_args()
+  initialize_tracer()
+  server = grpc.server(
+    futures.ThreadPoolExecutor(max_workers=args.parallel_threshold)
+  )
+
+  service = GarfService()
+  garf_pb2_grpc.add_GarfServiceServicer_to_server(service, server)
+  SERVICE_NAMES = (
+    garf_pb2.DESCRIPTOR.services_by_name['GarfService'].full_name,
+    reflection.SERVICE_NAME,
+  )
+  reflection.enable_server_reflection(SERVICE_NAMES, server)
+  server.add_insecure_port(f'[::]:{args.port}')
+  server.start()
+  logging.info('Garf service started, listening on port %d', 50051)
+  server.wait_for_termination()