connection.h

Go to the documentation of this file.

#pragma once
 
#include "client_context.h"
#include "database.h"
#include "udf_function.h"
 
namespace lbug {
namespace main {

class Connection {
    friend class testing::BaseGraphTest;
    friend class testing::PrivateGraphTest;
    friend class testing::TestHelper;
    friend class benchmark::Benchmark;
    friend class ConnectionExecuteAsyncWorker;
    friend class ConnectionQueryAsyncWorker;
 
public:
    LBUG_API explicit Connection(Database* database);
    LBUG_API ~Connection();
    LBUG_API void setMaxNumThreadForExec(uint64_t numThreads);
    LBUG_API uint64_t getMaxNumThreadForExec();

    LBUG_API std::unique_ptr<QueryResult> query(std::string_view query);
 
    LBUG_API std::unique_ptr<QueryResult> queryAsArrow(std::string_view query, int64_t chunkSize);

    LBUG_API std::unique_ptr<PreparedStatement> prepare(std::string_view query);

    LBUG_API std::unique_ptr<PreparedStatement> prepareWithParams(std::string_view query,
        std::unordered_map<std::string, std::unique_ptr<common::Value>> inputParams);

    template<typename... Args>
    inline std::unique_ptr<QueryResult> execute(PreparedStatement* preparedStatement,
        std::pair<std::string, Args>... args) {
        std::unordered_map<std::string, std::unique_ptr<common::Value>> inputParameters;
        return executeWithParams(preparedStatement, std::move(inputParameters), args...);
    }

    LBUG_API std::unique_ptr<QueryResult> executeWithParams(PreparedStatement* preparedStatement,
        std::unordered_map<std::string, std::unique_ptr<common::Value>> inputParams);
    LBUG_API void interrupt();

    LBUG_API void setQueryTimeOut(uint64_t timeoutInMS);
 
    template<typename TR, typename... Args>
    void createScalarFunction(std::string name, TR (*udfFunc)(Args...)) {
        addScalarFunction(name, function::UDF::getFunction<TR, Args...>(name, udfFunc));
    }
 
    template<typename TR, typename... Args>
    void createScalarFunction(std::string name, std::vector<common::LogicalTypeID> parameterTypes,
        common::LogicalTypeID returnType, TR (*udfFunc)(Args...)) {
        addScalarFunction(name, function::UDF::getFunction<TR, Args...>(name, udfFunc,
                                    std::move(parameterTypes), returnType));
    }
 
    void addUDFFunctionSet(std::string name, function::function_set func) {
        addScalarFunction(name, std::move(func));
    }
 
    void removeUDFFunction(std::string name) { removeScalarFunction(name); }
 
    template<typename TR, typename... Args>
    void createVectorizedFunction(std::string name, function::scalar_func_exec_t scalarFunc) {
        addScalarFunction(name,
            function::UDF::getVectorizedFunction<TR, Args...>(name, std::move(scalarFunc)));
    }
 
    void createVectorizedFunction(std::string name,
        std::vector<common::LogicalTypeID> parameterTypes, common::LogicalTypeID returnType,
        function::scalar_func_exec_t scalarFunc) {
        addScalarFunction(name, function::UDF::getVectorizedFunction(name, std::move(scalarFunc),
                                    std::move(parameterTypes), returnType));
    }
 
    ClientContext* getClientContext() { return clientContext.get(); };
 
private:
    template<typename T, typename... Args>
    std::unique_ptr<QueryResult> executeWithParams(PreparedStatement* preparedStatement,
        std::unordered_map<std::string, std::unique_ptr<common::Value>> params,
        std::pair<std::string, T> arg, std::pair<std::string, Args>... args) {
        return clientContext->executeWithParams(preparedStatement, std::move(params), arg, args...);
    }
 
    LBUG_API void addScalarFunction(std::string name, function::function_set definitions);
    LBUG_API void removeScalarFunction(std::string name);
 
    std::unique_ptr<QueryResult> queryWithID(std::string_view query, uint64_t queryID);
 
    std::unique_ptr<QueryResult> executeWithParamsWithID(PreparedStatement* preparedStatement,
        std::unordered_map<std::string, std::unique_ptr<common::Value>> inputParams,
        uint64_t queryID);
 
private:
    Database* database;
    std::unique_ptr<ClientContext> clientContext;
    std::shared_ptr<common::DatabaseLifeCycleManager> dbLifeCycleManager;
};
 
} // namespace main
} // namespace lbug