cassandra/test/unit/org/apache/cassandra/cql3/QueryEventsTest.java

/*
 * 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 org.apache.cassandra.cql3;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import javax.annotation.Nullable;

import com.datastax.driver.core.PreparedStatement;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.SimpleStatement;
import com.datastax.driver.core.Statement;

import org.junit.Test;

import org.apache.cassandra.cql3.statements.BatchStatement;
import org.apache.cassandra.cql3.statements.ModificationStatement;
import org.apache.cassandra.cql3.statements.SelectStatement;
import org.apache.cassandra.db.ColumnFamilyStore;
import org.apache.cassandra.service.QueryState;
import org.apache.cassandra.transport.Message;
import org.apache.cassandra.transport.messages.ResultMessage;
import org.apache.cassandra.utils.ByteArrayUtil;

import static com.google.common.collect.Lists.newArrayList;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

public class QueryEventsTest extends CQLTester
{
    @Test
    public void queryTest() throws Throwable
    {
        createTable("create table %s (id int primary key, v int)");
        MockListener listener = new MockListener(getCurrentColumnFamilyStore());
        QueryEvents.instance.registerListener(listener);
        String query = formatQuery("insert into %s (id, v) values (1, 1)");
        executeNet(query);
        listener.verify("querySuccess", 1);
        assertEquals(query, listener.query);
        assertTrue(listener.statement instanceof ModificationStatement);

        query = formatQuery("select * from %s where id=1");
        executeNet(query);
        listener.verify("querySuccess", 2);
        assertEquals(query, listener.query);
        assertTrue(listener.statement instanceof SelectStatement);

        query = formatQuery("select xyz from %s where id=1");
        Exception expectedException = null;
        try
        {
            executeNet(query);
            fail("Query should fail");
        }
        catch (Exception e)
        {
            expectedException = e;
        }
        listener.verify(newArrayList("querySuccess", "queryFailure"), newArrayList(2, 1));
        assertEquals(query, listener.query);
        assertNotNull(listener.e);
        assertNotNull(expectedException);
    }

    @Test
    public void prepareExecuteTest()
    {
        createTable("create table %s (id int primary key, v int)");
        MockListener listener = new MockListener(getCurrentColumnFamilyStore());
        QueryEvents.instance.registerListener(listener);
        Session session = sessionNet();
        String query = formatQuery("select * from %s where id = 1");
        PreparedStatement ps = session.prepare(query);
        listener.verify("prepareSuccess", 1);
        assertEquals(query, listener.query);
        assertTrue(listener.statement instanceof SelectStatement);
        Statement s = ps.bind();
        session.execute(s);
        listener.verify(newArrayList("prepareSuccess", "executeSuccess"), newArrayList(1, 1));

        QueryProcessor.clearPreparedStatements(false);
        s = ps.bind();
        session.execute(s); // this re-prepares the query!!
        listener.verify(newArrayList("prepareSuccess", "executeSuccess", "executeFailure"), newArrayList(2, 2, 1));

        query = formatQuery("select abcdef from %s where id = 1");
        Exception expectedException = null;
        try
        {
            session.prepare(query);
            fail("should fail");
        }
        catch (Exception e)
        {
            expectedException = e;
        }
        listener.verify(newArrayList("prepareSuccess", "prepareFailure", "executeSuccess", "executeFailure"), newArrayList(2, 1, 2, 1));
        assertNotNull(listener.e);
        assertNotNull(expectedException);
    }

    @Test
    public void batchTest()
    {
        createTable("create table %s (id int primary key, v int)");
        BatchMockListener listener = new BatchMockListener(getCurrentColumnFamilyStore());
        QueryEvents.instance.registerListener(listener);
        Session session = sessionNet();
        com.datastax.driver.core.BatchStatement batch = new com.datastax.driver.core.BatchStatement(com.datastax.driver.core.BatchStatement.Type.UNLOGGED);
        String q1 = formatQuery("insert into %s (id, v) values (?, ?)");
        PreparedStatement ps = session.prepare(q1);
        batch.add(ps.bind(1,1));
        batch.add(ps.bind(2,2));
        String q2 = formatQuery("insert into %s (id, v) values (1,1)");
        batch.add(new SimpleStatement(formatQuery("insert into %s (id, v) values (1,1)")));
        session.execute(batch);

        listener.verify(newArrayList("prepareSuccess", "batchSuccess"), newArrayList(1, 1));
        assertEquals(3, listener.queries.size());
        assertEquals(BatchStatement.Type.UNLOGGED, listener.batchType);
        assertEquals(newArrayList(q1, q1, q2), listener.queries);
        assertListOfArraysEquals(newArrayList(new byte[][] {ByteArrayUtil.bytes(1), ByteArrayUtil.bytes(1)},
                                              new byte[][] {ByteArrayUtil.bytes(2), ByteArrayUtil.bytes(2)},
                                              new byte[][] {}), listener.values);

        batch.add(new SimpleStatement("insert into abc.def (id, v) values (1,1)"));
        try
        {
            session.execute(batch);
            fail("Batch should fail");
        }
        catch (Exception e)
        {
            // ok
        }
        listener.verify(newArrayList("prepareSuccess", "batchSuccess", "batchFailure"), newArrayList(1, 1, 1));
        assertEquals(3, listener.queries.size());
        assertEquals(BatchStatement.Type.UNLOGGED, listener.batchType);
        assertEquals(newArrayList(q1, q1, q2), listener.queries);
        assertListOfArraysEquals(newArrayList(new byte[][] { ByteArrayUtil.bytes(1), ByteArrayUtil.bytes(1)},
                                              new byte[][] {ByteArrayUtil.bytes(2), ByteArrayUtil.bytes(2)},
                                              new byte[][] {},
                                              new byte[][] {}), listener.values);
    }

    private static void assertListOfArraysEquals(List<byte[][]> expected, List<byte[][]> actual)
    {
        assertEquals(expected.size(), actual.size());
        for (int i = 0; i < expected.size(); i++)
            assertArrayEquals(expected.get(i), actual.get(i));
    }

    @Test
    public void errorInListenerTest() throws Throwable
    {
        createTable("create table %s (id int primary key, v int)");
        QueryEvents.Listener listener = new QueryEvents.Listener()
        {
            @Override
            public void querySuccess(CQLStatement statement, String query, QueryOptions options, QueryState state, long queryTime, Message.Response response)
            {
                throw new AssertionError("whoo");
            }

            @Override
            public void queryFailure(@Nullable CQLStatement statement, String query, QueryOptions options, QueryState state, Exception cause)
            {
                throw new AssertionError("whee");
            }
        };

        QueryEvents.instance.registerListener(listener);

        executeNet("insert into %s (id, v) values (2,2)");
        ResultSet rs = executeNet("select * from %s");
        assertEquals(2, rs.one().getInt("id"));

        // record the exception without the throwing listener:
        QueryEvents.instance.unregisterListener(listener);
        Exception expected = null;
        try
        {
            executeNet("select blabla from %s");
            fail("Query should throw");
        }
        catch (Exception e)
        {
            expected = e;
        }

        QueryEvents.instance.registerListener(listener);
        // and with the listener:
        try
        {
            executeNet("select blabla from %s");
            fail("Query should throw");
        }
        catch (Exception e)
        {
            // make sure we throw the same exception even if the listener throws;
            assertSame(expected.getClass(), e.getClass());
            assertEquals(expected.getMessage(), e.getMessage());
        }


    }

    private static class MockListener implements QueryEvents.Listener
    {
        private final String tableName;
        private Map<String, Integer> callCounts = new HashMap<>();
        private String query;
        private CQLStatement statement;
        private long start = System.currentTimeMillis();
        long queryTime;
        private Exception e;


        MockListener(ColumnFamilyStore currentColumnFamilyStore)
        {
            tableName = currentColumnFamilyStore.getTableName();
        }

        public void querySuccess(CQLStatement statement, String query, QueryOptions options, QueryState state, long queryTime, Message.Response response)
        {
            if (query.contains(tableName))
            {
                inc("querySuccess");
                assertNotNull(query);
                this.query = query;
                assertNotNull(statement);
                this.statement = statement;
                this.queryTime = queryTime;
            }
        }

        public void queryFailure(@Nullable CQLStatement statement, String query, QueryOptions options, QueryState state, Exception cause)
        {

            if (query.contains(tableName))
            {
                inc("queryFailure");
                assertNotNull(query);
                this.query = query;
                this.statement = statement;
                e = cause;
            }
        }

        public void executeSuccess(CQLStatement statement, String query, QueryOptions options, QueryState state, long queryTime, Message.Response response)
        {
            if (query.contains(tableName))
                inc("executeSuccess");
        }
        public void executeFailure(@Nullable CQLStatement statement, @Nullable String query, QueryOptions options, QueryState state, Exception cause)
        {
            inc("executeFailure");
            e = cause;
        }


        public void prepareSuccess(CQLStatement statement, String query, QueryState state, long queryTime, ResultMessage.Prepared response)
        {
            if (query.contains(tableName))
            {
                inc("prepareSuccess");
                assertNotNull(query);
                this.query = query;
                this.statement = statement;
                this.queryTime = queryTime;
            }
        }
        public void prepareFailure(@Nullable CQLStatement statement, String query, QueryState state, Exception cause)
        {
            if (query.contains(tableName))
            {
                inc("prepareFailure");
                assertNotNull(query);
                assertNull(statement);
                this.query = query;
            }
        }

        void inc(String key)
        {
            callCounts.compute(key, (k, v) -> v == null ? 1 : v + 1);
        }

        private void verify(String key, int expected)
        {
            verify(Collections.singletonList(key), Collections.singletonList(expected));
        }

        void verify(List<String> keys, List<Integer> expected)
        {
            assertEquals("key count not equal: "+keys+" : "+callCounts, keys.size(), callCounts.size());
            assertTrue(queryTime >= start && queryTime <= System.currentTimeMillis());
            for (int i = 0; i < keys.size(); i++)
            {
                assertEquals("expected count mismatch for "+keys.get(i), (int)expected.get(i), (int) callCounts.get(keys.get(i)));
            }
        }
    }

    private static class BatchMockListener extends MockListener
    {
        private List<String> queries;
        private List<byte[][]> values;
        private BatchStatement.Type batchType;

        BatchMockListener(ColumnFamilyStore currentColumnFamilyStore)
        {
            super(currentColumnFamilyStore);
        }

        public void batchSuccess(BatchStatement.Type batchType, List<? extends CQLStatement> statements, List<String> queries, List<byte[][]> values, QueryOptions options, QueryState state, long queryTime, Message.Response response)
        {
            inc("batchSuccess");
            this.queries = queries;
            this.values = values;
            this.batchType = batchType;
            this.queryTime = queryTime;
        }

        public void batchFailure(BatchStatement.Type batchType, List<? extends CQLStatement> statements, List<String> queries, List<byte[][]> values, QueryOptions options, QueryState state, Exception cause)
        {
            inc("batchFailure");
            this.queries = queries;
            this.values = values;
            this.batchType = batchType;
        }

    }
}