client_context.h

Go to the documentation of this file.

#pragma once
 
#include <atomic>
#include <memory>
#include <mutex>
 
#include "progress_bar.h"
#include "timer.h"
#include "value.h"
#include "scan_replacement.h"
#include "client_config.h"
#include "statement.h"
#include "prepared_statement.h"
#include "query_result.h"
#include "transaction_context.h"
 
namespace kuzu {
 
namespace binder {
class Binder;
class ExpressionBinder;
} // namespace binder
 
namespace common {
class RandomEngine;
class TaskScheduler;
} // namespace common
 
namespace extension {
struct ExtensionOptions;
}
 
namespace main {
struct DBConfig;
class Database;
class DatabaseManager;
class AttachedKuzuDatabase;
 
struct ActiveQuery {
 explicit ActiveQuery();
 std::atomic<bool> interrupted;
 common::Timer timer;
 
 void reset();
};
 
class KUZU_API ClientContext {
 friend class Connection;
 friend class binder::Binder;
 friend class binder::ExpressionBinder;
 
public:
 explicit ClientContext(Database* database);
 ~ClientContext();
 
 // Client config
 const ClientConfig* getClientConfig() const { return &clientConfig; }
 ClientConfig* getClientConfigUnsafe() { return &clientConfig; }
 const DBConfig* getDBConfig() const { return &dbConfig; }
 DBConfig* getDBConfigUnsafe() { return &dbConfig; }
 common::Value getCurrentSetting(const std::string& optionName);
 bool isOptionSet(const std::string& optionName) const;
 // Timer and timeout
 void interrupt() { activeQuery.interrupted = true; }
 bool interrupted() const { return activeQuery.interrupted; }
 bool hasTimeout() const { return clientConfig.timeoutInMS != 0; }
 void setQueryTimeOut(uint64_t timeoutInMS);
 uint64_t getQueryTimeOut() const;
 void startTimer();
 uint64_t getTimeoutRemainingInMS() const;
 void resetActiveQuery() { activeQuery.reset(); }
 
 // Parallelism
 void setMaxNumThreadForExec(uint64_t numThreads);
 uint64_t getMaxNumThreadForExec() const;
 
 // Transaction.
 transaction::Transaction* getTx() const;
 transaction::TransactionContext* getTransactionContext() const;
 
 // Progress bar
 common::ProgressBar* getProgressBar() const;
 
 // Replace function.
 void addScanReplace(function::ScanReplacement scanReplacement);
 std::unique_ptr<function::ScanReplacementData> tryReplace(const std::string& objectName) const;
 // Extension
 void setExtensionOption(std::string name, common::Value value);
 extension::ExtensionOptions* getExtensionOptions() const;
 std::string getExtensionDir() const;
 
 // Environment.
 std::string getEnvVariable(const std::string& name);
 
 // Database component getters.
 std::string getDatabasePath() const;
 Database* getDatabase() const { return localDatabase; }
 common::TaskScheduler* getTaskScheduler() const;
 DatabaseManager* getDatabaseManager() const;
 storage::StorageManager* getStorageManager() const;
 storage::MemoryManager* getMemoryManager();
 storage::WAL* getWAL() const;
 catalog::Catalog* getCatalog() const;
 transaction::TransactionManager* getTransactionManagerUnsafe() const;
 common::VirtualFileSystem* getVFSUnsafe() const;
 common::RandomEngine* getRandomEngine();
 
 // Query.
 std::unique_ptr<PreparedStatement> prepare(std::string_view query);
 std::unique_ptr<QueryResult> executeWithParams(PreparedStatement* preparedStatement,
 std::unordered_map<std::string, std::unique_ptr<common::Value>> inputParams,
 std::optional<uint64_t> queryID = std::nullopt);
 std::unique_ptr<QueryResult> query(std::string_view queryStatement,
 std::optional<uint64_t> queryID = std::nullopt);
 void runQuery(std::string query);
 
 // TODO(Jiamin): should remove after supporting ddl in manual tx
 std::unique_ptr<PreparedStatement> prepareTest(std::string_view query);
 // only use for test framework
 std::vector<std::shared_ptr<parser::Statement>> parseQuery(std::string_view query);
 
 void setDefaultDatabase(AttachedKuzuDatabase* defaultDatabase_);
 bool hasDefaultDatabase();
 
 void runFuncInTransaction(const std::function<void(void)>& fun);
 void addScalarFunction(std::string name, function::function_set definitions);
 void removeScalarFunction(std::string name);
 
 void cleanUP();
 
private:
 std::unique_ptr<QueryResult> query(std::string_view query, std::string_view encodedJoin,
 bool enumerateAllPlans = true, std::optional<uint64_t> queryID = std::nullopt);
 
 std::unique_ptr<QueryResult> queryResultWithError(std::string_view errMsg);
 
 std::unique_ptr<PreparedStatement> preparedStatementWithError(std::string_view errMsg);
 
 // when we do prepare, we will start a transaction for the query
 // when we execute after prepare in a same context, we set requireNewTx to false and will not
 // commit the transaction in prepare when we only prepare a query statement, we set requireNewTx
 // to true and will commit the transaction in prepare
 std::unique_ptr<PreparedStatement> prepareNoLock(
 std::shared_ptr<parser::Statement> parsedStatement, bool enumerateAllPlans = false,
 std::string_view joinOrder = std::string_view(), bool requireNewTx = true,
 std::optional<std::unordered_map<std::string, std::shared_ptr<common::Value>>> inputParams =
 std::nullopt);
 
 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) {
 auto name = arg.first;
 auto val = std::make_unique<common::Value>((T)arg.second);
 params.insert({name, std::move(val)});
 return executeWithParams(preparedStatement, std::move(params), args...);
 }
 
 void bindParametersNoLock(PreparedStatement* preparedStatement,
 const std::unordered_map<std::string, std::unique_ptr<common::Value>>& inputParams);
 
 std::unique_ptr<QueryResult> executeAndAutoCommitIfNecessaryNoLock(
 PreparedStatement* preparedStatement, uint32_t planIdx = 0u,
 std::optional<uint64_t> queryID = std::nullopt);
 
 bool canExecuteWriteQuery();
 
 // Client side configurable settings.
 ClientConfig clientConfig;
 // Database configurable settings.
 DBConfig& dbConfig;
 // Current query.
 ActiveQuery activeQuery;
 // Transaction context.
 std::unique_ptr<transaction::TransactionContext> transactionContext;
 // Replace external object as pointer Value;
 std::vector<function::ScanReplacement> scanReplacements;
 // Extension configurable settings.
 std::unordered_map<std::string, common::Value> extensionOptionValues;
 // Random generator for UUID.
 std::unique_ptr<common::RandomEngine> randomEngine;
 // Local database.
 Database* localDatabase;
 // Remote database.
 AttachedKuzuDatabase* remoteDatabase;
 // Progress bar.
 std::unique_ptr<common::ProgressBar> progressBar;
 std::mutex mtx;
};
 
} // namespace main
} // namespace kuzu