#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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
#
#    http://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.
#

import logging
import random
import sys
import unittest
import uuid

import pytest

import apache_beam as beam
from apache_beam.testing.test_pipeline import TestPipeline
from apache_beam.testing.util import assert_that
from apache_beam.testing.util import equal_to

# Protect against environments where spanner library is not available.
# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
# pylint: disable=unused-import
try:
  from google.cloud import spanner
  from apache_beam.io.gcp import resource_identifiers
  from apache_beam.io.gcp.experimental.spannerio import create_transaction
  from apache_beam.io.gcp.experimental.spannerio import ReadOperation
  from apache_beam.io.gcp.experimental.spannerio import ReadFromSpanner
  from apache_beam.metrics import monitoring_infos
  from apache_beam.metrics.execution import MetricsEnvironment
  from apache_beam.metrics.metricbase import MetricName
except ImportError:
  spanner = None
# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
# pylint: enable=unused-import

_LOGGER = logging.getLogger(__name__)
_TEST_INSTANCE_ID = 'beam-test'


@unittest.skipIf(spanner is None, 'GCP dependencies are not installed.')
class SpannerReadIntegrationTest(unittest.TestCase):
  TEST_DATABASE = None
  _database_prefix = "pybeam-read-{}"
  _data = None
  _SPANNER_CLIENT = None
  _SPANNER_INSTANCE = None

  @classmethod
  def _generate_table_name(cls):
    cls.TEST_DATABASE = cls._database_prefix.format(
        ''.join(random.sample(uuid.uuid4().hex, 15)))
    return cls.TEST_DATABASE

  @classmethod
  def _create_database(cls):
    _LOGGER.info("Creating test database: %s" % cls.TEST_DATABASE)
    instance = cls._SPANNER_INSTANCE
    database = instance.database(
        cls.TEST_DATABASE,
        ddl_statements=[
            """CREATE TABLE Users (
            UserId    INT64 NOT NULL,
            Key       STRING(1024)
        ) PRIMARY KEY (UserId)""",
        ])
    operation = database.create()
    _LOGGER.info("Creating database: Done! %s" % str(operation.result()))

  @classmethod
  def _add_dummy_entries(cls):
    _LOGGER.info("Dummy Data: Adding dummy data...")
    instance = cls._SPANNER_INSTANCE
    database = instance.database(cls.TEST_DATABASE)
    data = cls._data = [[x + 1, uuid.uuid4().hex] for x in range(200)]
    with database.batch() as batch:
      batch.insert(table='Users', columns=('UserId', 'Key'), values=data)

  @classmethod
  def setUpClass(cls):
    _LOGGER.info(".... PyVersion ---> %s" % str(sys.version))
    _LOGGER.info(".... Setting up!")
    cls.test_pipeline = TestPipeline(is_integration_test=True)
    cls.args = cls.test_pipeline.get_full_options_as_args()
    cls.runner_name = type(cls.test_pipeline.runner).__name__
    cls.project = cls.test_pipeline.get_option('project')
    cls.instance = (
        cls.test_pipeline.get_option('instance') or _TEST_INSTANCE_ID)
    _ = cls._generate_table_name()
    spanner_client = cls._SPANNER_CLIENT = spanner.Client()
    _LOGGER.info(".... Spanner Client created!")
    cls._SPANNER_INSTANCE = spanner_client.instance(cls.instance)
    cls._create_database()
    cls._add_dummy_entries()
    _LOGGER.info("Spanner Read IT Setup Complete...")

  @pytest.mark.spannerio_it
  def test_read_via_table(self):
    _LOGGER.info("Spanner Read via table")
    with beam.Pipeline(argv=self.args) as p:
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          table="Users",
          columns=["UserId", "Key"])
      assert_that(r, equal_to(self._data))

  @pytest.mark.spannerio_it
  def test_read_via_sql(self):
    _LOGGER.info("Running Spanner via sql")
    with beam.Pipeline(argv=self.args) as p:
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          sql="select * from Users")
      assert_that(r, equal_to(self._data))

  @pytest.mark.spannerio_it
  def test_transaction_table_metrics_ok_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with beam.Pipeline(argv=self.args) as p:
      transaction = (
          p
          | create_transaction(self.project, self.instance, self.TEST_DATABASE))
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          table="Users",
          columns=["UserId", "Key"],
          transaction=transaction)
      assert_that(r, equal_to(self._data))
    self.verify_table_read_call_metric(
        self.project, self.TEST_DATABASE, 'Users', 'ok', 1)

  @pytest.mark.spannerio_it
  def test_transaction_table_metrics_error_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with self.assertRaises(Exception):
      p = beam.Pipeline(argv=self.args)
      transaction = (
          p
          | create_transaction(self.project, self.instance, self.TEST_DATABASE))
      _ = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          table="INVALID_TABLE",
          columns=["UserId", "Key"],
          transaction=transaction)

      res = p.run()
      res.wait_until_finish()

    self.verify_table_read_call_metric(
        self.project, self.TEST_DATABASE, 'INVALID_TABLE', '404', 1)

  @pytest.mark.spannerio_it
  def test_transaction_sql_metrics_ok_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with beam.Pipeline(argv=self.args) as p:
      transaction = (
          p
          | create_transaction(self.project, self.instance, self.TEST_DATABASE))
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          sql="select * from Users",
          query_name='query-1',
          transaction=transaction)

      assert_that(r, equal_to(self._data))
    self.verify_sql_read_call_metric(
        self.project, self.TEST_DATABASE, 'query-1', 'ok', 1)

  @pytest.mark.spannerio_it
  def test_transaction_sql_metrics_error_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with self.assertRaises(Exception):
      p = beam.Pipeline(argv=self.args)
      transaction = (
          p
          | create_transaction(self.project, self.instance, self.TEST_DATABASE))
      _ = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          sql="select * from NonExistent",
          query_name="query-2",
          transaction=transaction)

      res = p.run()
      res.wait_until_finish()

    self.verify_sql_read_call_metric(
        self.project, self.TEST_DATABASE, 'query-2', '400', 1)

  @pytest.mark.spannerio_it
  def test_table_metrics_ok_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with beam.Pipeline(argv=self.args) as p:
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          table="Users",
          columns=["UserId", "Key"])

      assert_that(r, equal_to(self._data))
    self.verify_table_read_call_metric(
        self.project, self.TEST_DATABASE, 'Users', 'ok', 1)

  @pytest.mark.spannerio_it
  def test_table_metrics_error_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with self.assertRaises(Exception):
      p = beam.Pipeline(argv=self.args)
      _ = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          table="INVALID_TABLE",
          columns=["UserId", "Key"])

      res = p.run()
      res.wait_until_finish()

      self.verify_table_read_call_metric(
          self.project, self.TEST_DATABASE, 'INVALID_TABLE', '404', 1)

  @pytest.mark.spannerio_it
  def test_sql_metrics_ok_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with beam.Pipeline(argv=self.args) as p:
      r = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          sql="select * from Users",
          query_name='query-1')

      assert_that(r, equal_to(self._data))
    self.verify_sql_read_call_metric(
        self.project, self.TEST_DATABASE, 'query-1', 'ok', 1)

  @pytest.mark.spannerio_it
  def test_sql_metrics_error_call(self):
    if 'DirectRunner' not in self.runner_name:
      raise unittest.SkipTest('This test only runs with DirectRunner.')

    MetricsEnvironment.process_wide_container().reset()

    with self.assertRaises(Exception):
      p = beam.Pipeline(argv=self.args)
      _ = p | ReadFromSpanner(
          self.project,
          self.instance,
          self.TEST_DATABASE,
          sql="select * from NonExistent",
          query_name='query-2')

      res = p.run()
      res.wait_until_finish()

      self.verify_sql_read_call_metric(
          self.project, self.TEST_DATABASE, 'query-2', '400', 1)

  def verify_table_read_call_metric(
      self, project, database, table, status, count):
    resource = resource_identifiers.SpannerTable(project, database, table)
    labels = {
        monitoring_infos.SERVICE_LABEL: 'Spanner',
        monitoring_infos.METHOD_LABEL: 'Read',
        monitoring_infos.SPANNER_PROJECT_ID: project,
        monitoring_infos.SPANNER_DATABASE_ID: database,
        monitoring_infos.RESOURCE_LABEL: resource,
        monitoring_infos.SPANNER_TABLE_ID: table,
        monitoring_infos.STATUS_LABEL: status
    }
    metric_name = MetricName(
        None, None, urn=monitoring_infos.API_REQUEST_COUNT_URN, labels=labels)
    metric_value = MetricsEnvironment.process_wide_container().get_counter(
        metric_name).get_cumulative()

    self.assertEqual(metric_value, count)

  def verify_sql_read_call_metric(
      self, project, database, query_name, status, count):
    resource = resource_identifiers.SpannerSqlQuery(project, query_name)
    labels = {
        monitoring_infos.SERVICE_LABEL: 'Spanner',
        monitoring_infos.METHOD_LABEL: 'Read',
        monitoring_infos.SPANNER_PROJECT_ID: project,
        monitoring_infos.SPANNER_DATABASE_ID: database,
        monitoring_infos.RESOURCE_LABEL: resource,
        monitoring_infos.SPANNER_QUERY_NAME: query_name,
        monitoring_infos.STATUS_LABEL: status
    }

    metric_name = MetricName(
        None, None, urn=monitoring_infos.API_REQUEST_COUNT_URN, labels=labels)
    metric_value = MetricsEnvironment.process_wide_container().get_counter(
        metric_name).get_cumulative()

    self.assertEqual(metric_value, count)

  @classmethod
  def tearDownClass(cls):
    # drop the testing database after the tests
    database = cls._SPANNER_INSTANCE.database(cls.TEST_DATABASE)
    database.drop()


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()
