tools/fqltool/src/org/apache/cassandra/fqltool/ResultStore.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.fqltool;


import java.io.File;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

import net.openhft.chronicle.bytes.BytesStore;
import net.openhft.chronicle.core.io.Closeable;
import net.openhft.chronicle.core.io.IORuntimeException;
import net.openhft.chronicle.queue.ChronicleQueue;
import net.openhft.chronicle.queue.impl.single.SingleChronicleQueueBuilder;
import net.openhft.chronicle.queue.ExcerptAppender;
import net.openhft.chronicle.wire.ReadMarshallable;
import net.openhft.chronicle.wire.ValueIn;
import net.openhft.chronicle.wire.ValueOut;
import net.openhft.chronicle.wire.WireIn;
import net.openhft.chronicle.wire.WireOut;
import net.openhft.chronicle.wire.WriteMarshallable;
import org.apache.cassandra.utils.Pair;
import org.apache.cassandra.utils.binlog.BinLog;

/**
 * note that we store each row as a separate chronicle document to be able to
 * avoid reading up the entire result set in memory when comparing
 *
 * document formats:
 * to mark the start of a new result set:
 * -------------------
 * version: int16
 * type: column_definitions
 * column_count: int32;
 * column_definition: text, text
 * column_definition: text, text
 * ....
 * --------------------
 *
 * to mark a failed query:
 * ---------------------
 * version: int16
 * type: query_failed
 * message: text
 * ---------------------
 *
 * row:
 * --------------------
 * version: int16
 * type: row
 * row_column_count: int32
 * column: bytes
 * ---------------------
 *
 * to mark the end of a result set:
 * -------------------
 * version: int16
 * type: end_resultset
 * -------------------
 *
 */
public class ResultStore
{
    private static final String VERSION = "version";
    private static final String TYPE = "type";
    // types:
    private static final String ROW = "row";
    private static final String END = "end_resultset";
    private static final String FAILURE = "query_failed";
    private static final String COLUMN_DEFINITIONS = "column_definitions";
    // fields:
    private static final String COLUMN_DEFINITION = "column_definition";
    private static final String COLUMN_COUNT = "column_count";
    private static final String MESSAGE = "message";
    private static final String ROW_COLUMN_COUNT = "row_column_count";
    private static final String COLUMN = "column";

    private static final int CURRENT_VERSION = 0;

    private final List<ChronicleQueue> queues;
    private final List<ExcerptAppender> appenders;
    private final ChronicleQueue queryStoreQueue;
    private final ExcerptAppender queryStoreAppender;
    private final Set<Integer> finishedHosts = new HashSet<>();

    public ResultStore(List<File> resultPaths, File queryFilePath)
    {
        queues = resultPaths.stream().map(path -> SingleChronicleQueueBuilder.single(path).build()).collect(Collectors.toList());
        appenders = queues.stream().map(ChronicleQueue::acquireAppender).collect(Collectors.toList());
        queryStoreQueue = queryFilePath != null ? SingleChronicleQueueBuilder.single(queryFilePath).build() : null;
        queryStoreAppender = queryStoreQueue != null ? queryStoreQueue.acquireAppender() : null;
    }

    /**
     * Store the column definitions in cds
     *
     * the ColumnDefinitions at position x will get stored by the appender at position x
     *
     * Calling this method indicates that we are starting a new result set from a query, it must be called before
     * calling storeRows.
     *
     */
    public void storeColumnDefinitions(FQLQuery query, List<ResultHandler.ComparableColumnDefinitions> cds)
    {
        finishedHosts.clear();
        if (queryStoreAppender != null)
        {
            BinLog.ReleaseableWriteMarshallable writeMarshallableQuery = query.toMarshallable();
            queryStoreAppender.writeDocument(writeMarshallableQuery);
            writeMarshallableQuery.release();
        }
        for (int i = 0; i < cds.size(); i++)
        {
            ResultHandler.ComparableColumnDefinitions cd = cds.get(i);
            appenders.get(i).writeDocument(new ColumnDefsWriter(cd));
        }
    }

    /**
     * Store rows
     *
     * the row at position x will get stored by appender at position x
     *
     * Before calling this for a new result set, storeColumnDefinitions must be called.
     */
    public void storeRows(List<ResultHandler.ComparableRow> rows)
    {
        for (int i = 0; i < rows.size(); i++)
        {
            ResultHandler.ComparableRow row = rows.get(i);
            if (row == null && !finishedHosts.contains(i))
            {
                appenders.get(i).writeDocument(wire -> {
                    wire.write(VERSION).int16(CURRENT_VERSION);
                    wire.write(TYPE).text(END);
                });
                finishedHosts.add(i);
            }
            else if (row != null)
            {
                appenders.get(i).writeDocument(new RowWriter(row));
            }
        }
    }

    public void close()
    {
        queues.forEach(Closeable::close);
        if (queryStoreQueue != null)
            queryStoreQueue.close();
    }

    static class ColumnDefsWriter implements WriteMarshallable
    {
        private final ResultHandler.ComparableColumnDefinitions defs;

        ColumnDefsWriter(ResultHandler.ComparableColumnDefinitions defs)
        {
            this.defs = defs;
        }

        public void writeMarshallable(WireOut wire)
        {
            wire.write(VERSION).int16(CURRENT_VERSION);
            if (!defs.wasFailed())
            {
                wire.write(TYPE).text(COLUMN_DEFINITIONS);
                wire.write(COLUMN_COUNT).int32(defs.size());
                for (ResultHandler.ComparableDefinition d : defs.asList())
                {
                    ValueOut vo = wire.write(COLUMN_DEFINITION);
                    vo.text(d.getName());
                    vo.text(d.getType());
                }
            }
            else
            {
                wire.write(TYPE).text(FAILURE);
                wire.write(MESSAGE).text(defs.getFailureException().getMessage());
            }
        }
    }

    static class ColumnDefsReader implements ReadMarshallable
    {
        boolean wasFailed;
        String failureMessage;
        List<Pair<String, String>> columnDefinitions = new ArrayList<>();

        public void readMarshallable(WireIn wire) throws IORuntimeException
        {
            int version = wire.read(VERSION).int16();
            String type = wire.read(TYPE).text();
            if (type.equals(FAILURE))
            {
                wasFailed = true;
                failureMessage = wire.read(MESSAGE).text();
            }
            else if (type.equals(COLUMN_DEFINITION))
            {
                int columnCount = wire.read(COLUMN_COUNT).int32();
                for (int i = 0; i < columnCount; i++)
                {
                    ValueIn vi = wire.read(COLUMN_DEFINITION);
                    String name = vi.text();
                    String dataType = vi.text();
                    columnDefinitions.add(Pair.create(name, dataType));
                }
            }
        }
    }

    /**
     * read a single row from the wire, or, marks itself finished if we read "end_resultset"
     */
    static class RowReader implements ReadMarshallable
    {
        boolean isFinished;
        List<ByteBuffer> rows = new ArrayList<>();

        public void readMarshallable(WireIn wire) throws IORuntimeException
        {
            int version = wire.read(VERSION).int32();
            String type = wire.read(TYPE).text();
            if (!type.equals(END))
            {
                isFinished = false;
                int rowColumnCount = wire.read(ROW_COLUMN_COUNT).int32();

                for (int i = 0; i < rowColumnCount; i++)
                {
                    byte[] b = wire.read(COLUMN).bytes();
                    rows.add(ByteBuffer.wrap(b));
                }
            }
            else
            {
                isFinished = true;
            }
        }
    }

    /**
     * Writes a single row to the given wire
     */
    static class RowWriter implements WriteMarshallable
    {
        private final ResultHandler.ComparableRow row;

        RowWriter(ResultHandler.ComparableRow row)
        {
            this.row = row;
        }

        public void writeMarshallable(WireOut wire)
        {
            wire.write(VERSION).int16(CURRENT_VERSION);
            wire.write(TYPE).text(ROW);
            wire.write(ROW_COLUMN_COUNT).int32(row.getColumnDefinitions().size());
            for (int jj = 0; jj < row.getColumnDefinitions().size(); jj++)
            {
                ByteBuffer bb = row.getBytesUnsafe(jj);
                if (bb != null)
                    wire.write(COLUMN).bytes(BytesStore.wrap(bb));
                else
                    wire.write(COLUMN).bytes("NULL".getBytes());
            }
        }
    }

}