/*
 * 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.
 */
package com.datastax.oss.driver.internal.core.metadata.schema.queries;

import static com.datastax.oss.driver.Assertions.assertThat;
import static com.datastax.oss.driver.Assertions.assertThatStage;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.api.core.Version;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.internal.core.adminrequest.AdminResult;
import com.datastax.oss.driver.internal.core.channel.DriverChannel;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList;
import java.util.Collections;
import java.util.Queue;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.LinkedBlockingDeque;
import org.junit.Before;
import org.junit.Test;

public class Cassandra3SchemaQueriesTest extends SchemaQueriesTest {

  @Before
  @Override
  public void setup() {
    super.setup();

    // By default, no keyspace filter
    when(config.getStringList(
            DefaultDriverOption.METADATA_SCHEMA_REFRESHED_KEYSPACES, Collections.emptyList()))
        .thenReturn(Collections.emptyList());
    when(node.getCassandraVersion()).thenReturn(Version.V3_0_0);
  }

  @Test
  public void should_query_without_keyspace_filter() {
    should_query_with_where_clause("");
  }

  @Test
  public void should_query_with_keyspace_filter() {
    when(config.getStringList(
            DefaultDriverOption.METADATA_SCHEMA_REFRESHED_KEYSPACES, Collections.emptyList()))
        .thenReturn(ImmutableList.of("ks1", "ks2"));

    should_query_with_where_clause(" WHERE keyspace_name IN ('ks1','ks2')");
  }

  private void should_query_with_where_clause(String whereClause) {
    SchemaQueriesWithMockedChannel queries =
        new SchemaQueriesWithMockedChannel(driverChannel, node, config, "test");
    CompletionStage<SchemaRows> result = queries.execute();

    // Keyspace
    Call call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.keyspaces" + whereClause);
    call.result.complete(
        mockResult(mockRow("keyspace_name", "ks1"), mockRow("keyspace_name", "ks2")));

    // Types
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.types" + whereClause);
    call.result.complete(mockResult(mockRow("keyspace_name", "ks1", "type_name", "type")));

    // Tables
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.tables" + whereClause);
    call.result.complete(mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo")));

    // Columns
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.columns" + whereClause);
    call.result.complete(
        mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "k")));

    // Indexes
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.indexes" + whereClause);
    call.result.complete(
        mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "index_name", "index")));

    // Views
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.views" + whereClause);
    call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "view_name", "foo")));

    // Functions
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.functions" + whereClause);
    call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "function_name", "add")));

    // Aggregates
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.aggregates" + whereClause);
    call.result.complete(mockResult(mockRow("keyspace_name", "ks2", "aggregate_name", "add")));

    channel.runPendingTasks();

    assertThatStage(result)
        .isSuccess(
            rows -> {
              assertThat(rows.getNode()).isEqualTo(node);

              // Keyspace
              assertThat(rows.keyspaces()).hasSize(2);
              assertThat(rows.keyspaces().get(0).getString("keyspace_name")).isEqualTo("ks1");
              assertThat(rows.keyspaces().get(1).getString("keyspace_name")).isEqualTo("ks2");

              // Types
              assertThat(rows.types().keySet()).containsOnly(KS1_ID);
              assertThat(rows.types().get(KS1_ID)).hasSize(1);
              assertThat(rows.types().get(KS1_ID).iterator().next().getString("type_name"))
                  .isEqualTo("type");

              // Tables
              assertThat(rows.tables().keySet()).containsOnly(KS1_ID);
              assertThat(rows.tables().get(KS1_ID)).hasSize(1);
              assertThat(rows.tables().get(KS1_ID).iterator().next().getString("table_name"))
                  .isEqualTo("foo");

              // Columns
              assertThat(rows.columns().keySet()).containsOnly(KS1_ID);
              assertThat(rows.columns().get(KS1_ID).keySet()).containsOnly(FOO_ID);
              assertThat(
                      rows.columns()
                          .get(KS1_ID)
                          .get(FOO_ID)
                          .iterator()
                          .next()
                          .getString("column_name"))
                  .isEqualTo("k");

              // Indexes
              assertThat(rows.indexes().keySet()).containsOnly(KS1_ID);
              assertThat(rows.indexes().get(KS1_ID).keySet()).containsOnly(FOO_ID);
              assertThat(
                      rows.indexes()
                          .get(KS1_ID)
                          .get(FOO_ID)
                          .iterator()
                          .next()
                          .getString("index_name"))
                  .isEqualTo("index");

              // Views
              assertThat(rows.views().keySet()).containsOnly(KS2_ID);
              assertThat(rows.views().get(KS2_ID)).hasSize(1);
              assertThat(rows.views().get(KS2_ID).iterator().next().getString("view_name"))
                  .isEqualTo("foo");

              // Functions
              assertThat(rows.functions().keySet()).containsOnly(KS2_ID);
              assertThat(rows.functions().get(KS2_ID)).hasSize(1);
              assertThat(rows.functions().get(KS2_ID).iterator().next().getString("function_name"))
                  .isEqualTo("add");

              // Aggregates
              assertThat(rows.aggregates().keySet()).containsOnly(KS2_ID);
              assertThat(rows.aggregates().get(KS2_ID)).hasSize(1);
              assertThat(
                      rows.aggregates().get(KS2_ID).iterator().next().getString("aggregate_name"))
                  .isEqualTo("add");
            });
  }

  @Test
  public void should_query_with_paging() {
    SchemaQueriesWithMockedChannel queries =
        new SchemaQueriesWithMockedChannel(driverChannel, node, config, "test");
    CompletionStage<SchemaRows> result = queries.execute();

    // Keyspace
    Call call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.keyspaces");
    call.result.complete(mockResult(mockRow("keyspace_name", "ks1")));

    // No types
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.types");
    call.result.complete(mockResult(/*empty*/ ));

    // Tables
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.tables");
    call.result.complete(mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo")));

    // Columns: paged
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.columns");

    AdminResult page2 =
        mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "v"));
    AdminResult page1 =
        mockResult(page2, mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "k"));
    call.result.complete(page1);

    // No indexes
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.indexes");
    call.result.complete(mockResult(/*empty*/ ));

    // No views
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.views");
    call.result.complete(mockResult(/*empty*/ ));

    // No functions
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.functions");
    call.result.complete(mockResult(/*empty*/ ));

    // No aggregates
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.aggregates");
    call.result.complete(mockResult(/*empty*/ ));

    channel.runPendingTasks();

    assertThatStage(result)
        .isSuccess(
            rows -> {
              assertThat(rows.columns().keySet()).containsOnly(KS1_ID);
              assertThat(rows.columns().get(KS1_ID).keySet()).containsOnly(FOO_ID);
              assertThat(rows.columns().get(KS1_ID).get(FOO_ID))
                  .extracting(r -> r.getString("column_name"))
                  .containsExactly("k", "v");
            });
  }

  @Test
  public void should_ignore_malformed_rows() {
    SchemaQueriesWithMockedChannel queries =
        new SchemaQueriesWithMockedChannel(driverChannel, node, config, "test");
    CompletionStage<SchemaRows> result = queries.execute();

    // Keyspace
    Call call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.keyspaces");
    call.result.complete(mockResult(mockRow("keyspace_name", "ks1")));

    // No types
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.types");
    call.result.complete(mockResult(/*empty*/ ));

    // Tables
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.tables");
    call.result.complete(
        mockResult(
            mockRow("keyspace_name", "ks", "table_name", "foo"),
            // Missing keyspace name:
            mockRow("table_name", "foo")));

    // Columns
    call = queries.calls.poll();
    call.result.complete(
        mockResult(
            mockRow("keyspace_name", "ks", "table_name", "foo", "column_name", "k"),
            // Missing keyspace name:
            mockRow("table_name", "foo", "column_name", "k"),
            // Missing table name:
            mockRow("keyspace_name", "ks", "column_name", "k")));

    AdminResult page2 =
        mockResult(mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "v"));
    AdminResult page1 =
        mockResult(page2, mockRow("keyspace_name", "ks1", "table_name", "foo", "column_name", "k"));
    call.result.complete(page1);

    // No indexes
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.indexes");
    call.result.complete(mockResult(/*empty*/ ));

    // No views
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.views");
    call.result.complete(mockResult(/*empty*/ ));

    // No functions
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.functions");
    call.result.complete(mockResult(/*empty*/ ));

    // No aggregates
    call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.aggregates");
    call.result.complete(mockResult(/*empty*/ ));

    channel.runPendingTasks();

    assertThatStage(result)
        .isSuccess(
            rows -> {
              assertThat(rows.tables().keySet()).containsOnly(KS_ID);
              assertThat(rows.tables().get(KS_ID)).hasSize(1);
              assertThat(rows.tables().get(KS_ID).iterator().next().getString("table_name"))
                  .isEqualTo("foo");

              assertThat(rows.columns().keySet()).containsOnly(KS_ID);
              assertThat(rows.columns().get(KS_ID).keySet()).containsOnly(FOO_ID);
              assertThat(
                      rows.columns()
                          .get(KS_ID)
                          .get(FOO_ID)
                          .iterator()
                          .next()
                          .getString("column_name"))
                  .isEqualTo("k");
            });
  }

  @Test
  public void should_abort_if_query_fails() {
    SchemaQueriesWithMockedChannel queries =
        new SchemaQueriesWithMockedChannel(driverChannel, node, config, "test");
    CompletionStage<SchemaRows> result = queries.execute();

    Exception mockQueryError = new Exception("mock query error");

    Call call = queries.calls.poll();
    assertThat(call.query).isEqualTo("SELECT * FROM system_schema.keyspaces");
    call.result.completeExceptionally(mockQueryError);

    channel.runPendingTasks();

    assertThatStage(result).isFailed(throwable -> assertThat(throwable).isEqualTo(mockQueryError));
  }

  /** Extends the class under test to mock the query execution logic. */
  static class SchemaQueriesWithMockedChannel extends Cassandra3SchemaQueries {

    final Queue<Call> calls = new LinkedBlockingDeque<>();

    SchemaQueriesWithMockedChannel(
        DriverChannel channel, Node node, DriverExecutionProfile config, String logPrefix) {
      super(channel, node, config, logPrefix);
    }

    @Override
    protected CompletionStage<AdminResult> query(String query) {
      Call call = new Call(query);
      calls.add(call);
      return call.result;
    }
  }
}
