query_result.h

Go to the documentation of this file.

#pragma once
 
#include <string>
 
#include "api.h"
#include "arrow.h"
#include "database_lifecycle_manager.h"
#include "types.h"
#include "query_summary.h"
 
namespace lbug {
namespace processor {
class FlatTuple;
}
namespace main {
 
enum class QueryResultType {
    FTABLE = 0,
    ARROW = 1,
};

class QueryResult {
public:
    LBUG_API QueryResult();
    explicit QueryResult(QueryResultType type);
    QueryResult(QueryResultType type, std::vector<std::string> columnNames,
        std::vector<common::LogicalType> columnTypes);

    LBUG_API virtual ~QueryResult() = 0;
    LBUG_API bool isSuccess() const;
    LBUG_API std::string getErrorMessage() const;
    LBUG_API size_t getNumColumns() const;
    LBUG_API std::vector<std::string> getColumnNames() const;
    LBUG_API std::vector<common::LogicalType> getColumnDataTypes() const;
    LBUG_API QuerySummary* getQuerySummary() const;
    QuerySummary* getQuerySummaryUnsafe();
    LBUG_API bool hasNextQueryResult() const;
    LBUG_API QueryResult* getNextQueryResult();
    LBUG_API virtual uint64_t getNumTuples() const = 0;
    LBUG_API virtual bool hasNext() const = 0;
    LBUG_API virtual std::shared_ptr<processor::FlatTuple> getNext() = 0;
    LBUG_API virtual void resetIterator() = 0;
    LBUG_API virtual std::string toString() const = 0;
    LBUG_API std::unique_ptr<ArrowSchema> getArrowSchema() const;
    LBUG_API virtual bool hasNextArrowChunk() = 0;
    LBUG_API virtual std::unique_ptr<ArrowArray> getNextArrowChunk(int64_t chunkSize) = 0;
 
    QueryResultType getType() const { return type; }
 
    void setColumnNames(std::vector<std::string> columnNames);
    void setColumnTypes(std::vector<common::LogicalType> columnTypes);
 
    void addNextResult(std::unique_ptr<QueryResult> next_);
    std::unique_ptr<QueryResult> moveNextResult();
 
    void setQuerySummary(std::unique_ptr<QuerySummary> summary);
 
    void setDBLifeCycleManager(
        std::shared_ptr<common::DatabaseLifeCycleManager> dbLifeCycleManager);
 
    static std::unique_ptr<QueryResult> getQueryResultWithError(const std::string& errorMessage);
 
    template<class TARGET>
    TARGET& cast() {
        return common::ku_dynamic_cast<TARGET&>(*this);
    }
    template<class TARGET>
    const TARGET& constCast() const {
        return common::ku_dynamic_cast<const TARGET&>(*this);
    }
 
protected:
    void validateQuerySucceed() const;
    void checkDatabaseClosedOrThrow() const;
 
protected:
    class QueryResultIterator {
    public:
        QueryResultIterator() = default;
 
        explicit QueryResultIterator(QueryResult* startResult) : current(startResult) {}
 
        void operator++() {
            if (current) {
                current = current->nextQueryResult.get();
            }
        }
 
        bool isEnd() const { return current == nullptr; }
 
        bool hasNextQueryResult() const { return current->nextQueryResult != nullptr; }
 
        QueryResult* getCurrentResult() const { return current; }
 
    private:
        QueryResult* current;
    };
 
    QueryResultType type;
 
    bool success = true;
 
    std::string errMsg;
 
    std::vector<std::string> columnNames;
 
    std::vector<common::LogicalType> columnTypes;
 
    std::shared_ptr<processor::FlatTuple> tuple;
 
    std::unique_ptr<QuerySummary> querySummary;
 
    std::unique_ptr<QueryResult> nextQueryResult;
 
    QueryResultIterator queryResultIterator;
 
    std::shared_ptr<common::DatabaseLifeCycleManager> dbLifeCycleManager;
};
 
} // namespace main
} // namespace lbug