# Copyright (c) 2020 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Export chromeperf Row data to BigQuery with Beam & Cloud Dataflow."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import datetime
import json
import logging

import apache_beam as beam
from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.metrics import Metrics
from apache_beam.transforms.core import FlatMap

from bq_export.split_by_timestamp import ReadTimestampRangeFromDatastore
from bq_export.export_options import BqExportOptions
from bq_export.utils import (FloatHack, PrintCounters,
                             WriteToPartitionedBigQuery)


def main():
  project = 'chromeperf'
  options = PipelineOptions()
  options.view_as(GoogleCloudOptions).project = project
  bq_export_options = options.view_as(BqExportOptions)

  p = beam.Pipeline(options=options)
  entities_read = Metrics.counter('main', 'entities_read')
  failed_entity_transforms = Metrics.counter('main', 'failed_entity_transforms')

  """
  CREATE TABLE `chromeperf.chromeperf_dashboard_data.rows_test`
  (revision INT64 NOT NULL,
   value FLOAT64 NOT NULL,
   error FLOAT64,
   `timestamp` TIMESTAMP NOT NULL,
   test_path STRING NOT NULL,
   master STRING,
   bot STRING,
   properties STRING)
  PARTITION BY DATE(`timestamp`);
  """  # pylint: disable=pointless-string-statement
  bq_row_schema = {'fields': [
      {'name': 'revision', 'type': 'INT64', 'mode': 'REQUIRED'},
      {'name': 'value', 'type': 'FLOAT', 'mode': 'REQUIRED'},
      {'name': 'error', 'type': 'FLOAT', 'mode': 'NULLABLE'},
      {'name': 'timestamp', 'type': 'TIMESTAMP', 'mode': 'REQUIRED'},
      {'name': 'test_path', 'type': 'STRING', 'mode': 'REQUIRED'},
      {'name': 'master', 'type': 'STRING', 'mode': 'NULLABLE'},
      {'name': 'bot', 'type': 'STRING', 'mode': 'NULLABLE'},
      {'name': 'properties', 'type': 'STRING', 'mode': 'NULLABLE'},
  ]}
  def RowEntityToRowDict(entity):
    entities_read.inc()
    try:
      d = {
          'revision': entity.key.id,
          'value': FloatHack(entity['value']),
          'error': FloatHack(entity.get('error')),
          'timestamp': entity['timestamp'].isoformat(),
          'test_path': entity.key.parent.name,
      }
      # Add the expando properties as a JSON-encoded dict.
      properties = {}
      for key, value in entity.items():
        if key in d or key == 'parent_test':
          # skip properties with dedicated columns.
          continue
        if isinstance(value, float):
          value = FloatHack(value)
        properties[key] = value
      d['properties'] = json.dumps(properties) if properties else None
      # Add columns derived from test_path: master, bot.
      test_path_parts = d['test_path'].split('/', 2)
      if len(test_path_parts) >= 3:
        d['master'] = test_path_parts[0]
        d['bot'] = test_path_parts[1]
      return [d]
    except KeyError:
      logging.getLogger().exception('Failed to convert Row')
      failed_entity_transforms.inc()
      return []

  row_query_params = dict(project=project, kind='Row')
  row_entities = (
      p
      | 'ReadFromDatastore(Row)' >> ReadTimestampRangeFromDatastore(
          row_query_params,
          time_range_provider=bq_export_options.GetTimeRangeProvider(),
          step=datetime.timedelta(minutes=5)))

  row_dicts = (
      row_entities | 'ConvertEntityToRow(Row)' >> FlatMap(RowEntityToRowDict))

  table_name = '{}:chromeperf_dashboard_data.rows{}'.format(
      project, bq_export_options.table_suffix)
  _ = row_dicts | 'WriteToBigQuery(rows)' >> WriteToPartitionedBigQuery(
      table_name, bq_row_schema)

  result = p.run()
  result.wait_until_finish()
  PrintCounters(result)

@@ -1361,16 +1361,20 @@
 ULL,%0A   
+std_
 error FL
@@ -1380,16 +1380,16 @@
 LOAT64,%0A
-
    %60time
@@ -1422,21 +1422,16 @@
 %0A   test
-_path
  STRING 
@@ -1755,16 +1755,20 @@
 name': '
+std_
 error', 
@@ -1895,21 +1895,16 @@
 ': 'test
-_path
 ', 'type
@@ -2309,16 +2309,20 @@
        '
+std_
 error': 
@@ -2424,21 +2424,16 @@
    'test
-_path
 ': entit
@@ -2611,11 +2611,12 @@
 key 
-== 
+in %5B
 'par
@@ -2624,16 +2624,26 @@
 nt_test'
+, 'error'%5D
 :%0A      
@@ -2919,21 +2919,16 @@
 rom test
-_path
 : master
@@ -2965,21 +2965,16 @@
  d%5B'test
-_path
 '%5D.split

Rename error→std_error, test_path→test.