1//===----------------------------------------------------------------------===//
2// DuckDB
3//
4// duckdb/main/connection.hpp
5//
6//
7//===----------------------------------------------------------------------===//
8
9#pragma once
10
11#include "duckdb/common/enums/profiler_format.hpp"
12#include "duckdb/common/serializer/buffered_file_writer.hpp"
13#include "duckdb/common/winapi.hpp"
14#include "duckdb/function/udf_function.hpp"
15#include "duckdb/main/materialized_query_result.hpp"
16#include "duckdb/main/pending_query_result.hpp"
17#include "duckdb/main/prepared_statement.hpp"
18#include "duckdb/main/query_result.hpp"
19#include "duckdb/main/relation.hpp"
20#include "duckdb/main/stream_query_result.hpp"
21#include "duckdb/main/table_description.hpp"
22#include "duckdb/parser/sql_statement.hpp"
23
24namespace duckdb {
25
26class ColumnDataCollection;
27class ClientContext;
28
29class DatabaseInstance;
30class DuckDB;
31class LogicalOperator;
32class SelectStatement;
33struct BufferedCSVReaderOptions;
34
35typedef void (*warning_callback)(std::string);
36
37//! A connection to a database. This represents a (client) connection that can
38//! be used to query the database.
39class Connection {
40public:
41 DUCKDB_API explicit Connection(DuckDB &database);
42 DUCKDB_API explicit Connection(DatabaseInstance &database);
43 DUCKDB_API ~Connection();
44
45 shared_ptr<ClientContext> context;
46 warning_callback warning_cb;
47
48public:
49 //! Returns query profiling information for the current query
50 DUCKDB_API string GetProfilingInformation(ProfilerPrintFormat format = ProfilerPrintFormat::QUERY_TREE);
51
52 //! Interrupt execution of the current query
53 DUCKDB_API void Interrupt();
54
55 //! Enable query profiling
56 DUCKDB_API void EnableProfiling();
57 //! Disable query profiling
58 DUCKDB_API void DisableProfiling();
59
60 DUCKDB_API void SetWarningCallback(warning_callback);
61
62 //! Enable aggressive verification/testing of queries, should only be used in testing
63 DUCKDB_API void EnableQueryVerification();
64 DUCKDB_API void DisableQueryVerification();
65 //! Force parallel execution, even for smaller tables. Should only be used in testing.
66 DUCKDB_API void ForceParallelism();
67
68 //! Issues a query to the database and returns a QueryResult. This result can be either a StreamQueryResult or a
69 //! MaterializedQueryResult. The result can be stepped through with calls to Fetch(). Note that there can only be
70 //! one active StreamQueryResult per Connection object. Calling SendQuery() will invalidate any previously existing
71 //! StreamQueryResult.
72 DUCKDB_API unique_ptr<QueryResult> SendQuery(const string &query);
73 //! Issues a query to the database and materializes the result (if necessary). Always returns a
74 //! MaterializedQueryResult.
75 DUCKDB_API unique_ptr<MaterializedQueryResult> Query(const string &query);
76 //! Issues a query to the database and materializes the result (if necessary). Always returns a
77 //! MaterializedQueryResult.
78 DUCKDB_API unique_ptr<MaterializedQueryResult> Query(unique_ptr<SQLStatement> statement);
79 // prepared statements
80 template <typename... Args>
81 unique_ptr<QueryResult> Query(const string &query, Args... args) {
82 vector<Value> values;
83 return QueryParamsRecursive(query, values, args...);
84 }
85
86 //! Issues a query to the database and returns a Pending Query Result. Note that "query" may only contain
87 //! a single statement.
88 DUCKDB_API unique_ptr<PendingQueryResult> PendingQuery(const string &query, bool allow_stream_result = false);
89 //! Issues a query to the database and returns a Pending Query Result
90 DUCKDB_API unique_ptr<PendingQueryResult> PendingQuery(unique_ptr<SQLStatement> statement,
91 bool allow_stream_result = false);
92
93 //! Prepare the specified query, returning a prepared statement object
94 DUCKDB_API unique_ptr<PreparedStatement> Prepare(const string &query);
95 //! Prepare the specified statement, returning a prepared statement object
96 DUCKDB_API unique_ptr<PreparedStatement> Prepare(unique_ptr<SQLStatement> statement);
97
98 //! Get the table info of a specific table (in the default schema), or nullptr if it cannot be found
99 DUCKDB_API unique_ptr<TableDescription> TableInfo(const string &table_name);
100 //! Get the table info of a specific table, or nullptr if it cannot be found
101 DUCKDB_API unique_ptr<TableDescription> TableInfo(const string &schema_name, const string &table_name);
102
103 //! Extract a set of SQL statements from a specific query
104 DUCKDB_API vector<unique_ptr<SQLStatement>> ExtractStatements(const string &query);
105 //! Extract the logical plan that corresponds to a query
106 DUCKDB_API unique_ptr<LogicalOperator> ExtractPlan(const string &query);
107
108 //! Appends a DataChunk to the specified table
109 DUCKDB_API void Append(TableDescription &description, DataChunk &chunk);
110 //! Appends a ColumnDataCollection to the specified table
111 DUCKDB_API void Append(TableDescription &description, ColumnDataCollection &collection);
112
113 //! Returns a relation that produces a table from this connection
114 DUCKDB_API shared_ptr<Relation> Table(const string &tname);
115 DUCKDB_API shared_ptr<Relation> Table(const string &schema_name, const string &table_name);
116 //! Returns a relation that produces a view from this connection
117 DUCKDB_API shared_ptr<Relation> View(const string &tname);
118 DUCKDB_API shared_ptr<Relation> View(const string &schema_name, const string &table_name);
119 //! Returns a relation that calls a specified table function
120 DUCKDB_API shared_ptr<Relation> TableFunction(const string &tname);
121 DUCKDB_API shared_ptr<Relation> TableFunction(const string &tname, const vector<Value> &values,
122 const named_parameter_map_t &named_parameters);
123 DUCKDB_API shared_ptr<Relation> TableFunction(const string &tname, const vector<Value> &values);
124 //! Returns a relation that produces values
125 DUCKDB_API shared_ptr<Relation> Values(const vector<vector<Value>> &values);
126 DUCKDB_API shared_ptr<Relation> Values(const vector<vector<Value>> &values, const vector<string> &column_names,
127 const string &alias = "values");
128 DUCKDB_API shared_ptr<Relation> Values(const string &values);
129 DUCKDB_API shared_ptr<Relation> Values(const string &values, const vector<string> &column_names,
130 const string &alias = "values");
131
132 //! Reads CSV file
133 DUCKDB_API shared_ptr<Relation> ReadCSV(const string &csv_file);
134 DUCKDB_API shared_ptr<Relation> ReadCSV(const string &csv_file, BufferedCSVReaderOptions &options);
135 DUCKDB_API shared_ptr<Relation> ReadCSV(const string &csv_file, const vector<string> &columns);
136
137 //! Reads Parquet file
138 DUCKDB_API shared_ptr<Relation> ReadParquet(const string &parquet_file, bool binary_as_string);
139 //! Returns a relation from a query
140 DUCKDB_API shared_ptr<Relation> RelationFromQuery(const string &query, const string &alias = "queryrelation",
141 const string &error = "Expected a single SELECT statement");
142 DUCKDB_API shared_ptr<Relation> RelationFromQuery(unique_ptr<SelectStatement> select_stmt,
143 const string &alias = "queryrelation");
144
145 //! Returns a substrait BLOB from a valid query
146 DUCKDB_API string GetSubstrait(const string &query);
147 //! Returns a Query Result from a substrait blob
148 DUCKDB_API unique_ptr<QueryResult> FromSubstrait(const string &proto);
149 //! Returns a substrait BLOB from a valid query
150 DUCKDB_API string GetSubstraitJSON(const string &query);
151 //! Returns a Query Result from a substrait JSON
152 DUCKDB_API unique_ptr<QueryResult> FromSubstraitJSON(const string &json);
153 DUCKDB_API void BeginTransaction();
154 DUCKDB_API void Commit();
155 DUCKDB_API void Rollback();
156 DUCKDB_API void SetAutoCommit(bool auto_commit);
157 DUCKDB_API bool IsAutoCommit();
158 DUCKDB_API bool HasActiveTransaction();
159
160 //! Fetch a list of table names that are required for a given query
161 DUCKDB_API unordered_set<string> GetTableNames(const string &query);
162
163 template <typename TR, typename... Args>
164 void CreateScalarFunction(const string &name, TR (*udf_func)(Args...)) {
165 scalar_function_t function = UDFWrapper::CreateScalarFunction<TR, Args...>(name, udf_func);
166 UDFWrapper::RegisterFunction<TR, Args...>(name, function, *context);
167 }
168
169 template <typename TR, typename... Args>
170 void CreateScalarFunction(const string &name, vector<LogicalType> args, LogicalType ret_type,
171 TR (*udf_func)(Args...)) {
172 scalar_function_t function =
173 UDFWrapper::CreateScalarFunction<TR, Args...>(name, args, std::move(ret_type), udf_func);
174 UDFWrapper::RegisterFunction(name, args, ret_type, udf_function: function, context&: *context);
175 }
176
177 template <typename TR, typename... Args>
178 void CreateVectorizedFunction(const string &name, scalar_function_t udf_func,
179 LogicalType varargs = LogicalType::INVALID) {
180 UDFWrapper::RegisterFunction<TR, Args...>(name, udf_func, *context, std::move(varargs));
181 }
182
183 void CreateVectorizedFunction(const string &name, vector<LogicalType> args, LogicalType ret_type,
184 scalar_function_t udf_func, LogicalType varargs = LogicalType::INVALID) {
185 UDFWrapper::RegisterFunction(name, args: std::move(args), ret_type: std::move(ret_type), udf_function: udf_func, context&: *context,
186 varargs: std::move(varargs));
187 }
188
189 //------------------------------------- Aggreate Functions ----------------------------------------//
190 template <typename UDF_OP, typename STATE, typename TR, typename TA>
191 void CreateAggregateFunction(const string &name) {
192 AggregateFunction function = UDFWrapper::CreateAggregateFunction<UDF_OP, STATE, TR, TA>(name);
193 UDFWrapper::RegisterAggrFunction(aggr_function: function, context&: *context);
194 }
195
196 template <typename UDF_OP, typename STATE, typename TR, typename TA, typename TB>
197 void CreateAggregateFunction(const string &name) {
198 AggregateFunction function = UDFWrapper::CreateAggregateFunction<UDF_OP, STATE, TR, TA, TB>(name);
199 UDFWrapper::RegisterAggrFunction(aggr_function: function, context&: *context);
200 }
201
202 template <typename UDF_OP, typename STATE, typename TR, typename TA>
203 void CreateAggregateFunction(const string &name, LogicalType ret_type, LogicalType input_typeA) {
204 AggregateFunction function =
205 UDFWrapper::CreateAggregateFunction<UDF_OP, STATE, TR, TA>(name, ret_type, input_typeA);
206 UDFWrapper::RegisterAggrFunction(aggr_function: function, context&: *context);
207 }
208
209 template <typename UDF_OP, typename STATE, typename TR, typename TA, typename TB>
210 void CreateAggregateFunction(const string &name, LogicalType ret_type, LogicalType input_typeA,
211 LogicalType input_typeB) {
212 AggregateFunction function =
213 UDFWrapper::CreateAggregateFunction<UDF_OP, STATE, TR, TA, TB>(name, ret_type, input_typeA, input_typeB);
214 UDFWrapper::RegisterAggrFunction(aggr_function: function, context&: *context);
215 }
216
217 void CreateAggregateFunction(const string &name, vector<LogicalType> arguments, LogicalType return_type,
218 aggregate_size_t state_size, aggregate_initialize_t initialize,
219 aggregate_update_t update, aggregate_combine_t combine, aggregate_finalize_t finalize,
220 aggregate_simple_update_t simple_update = nullptr,
221 bind_aggregate_function_t bind = nullptr,
222 aggregate_destructor_t destructor = nullptr) {
223 AggregateFunction function =
224 UDFWrapper::CreateAggregateFunction(name, arguments, return_type, state_size, initialize, update, combine,
225 finalize, simple_update, bind, destructor);
226 UDFWrapper::RegisterAggrFunction(aggr_function: function, context&: *context);
227 }
228
229private:
230 unique_ptr<QueryResult> QueryParamsRecursive(const string &query, vector<Value> &values);
231
232 template <typename T, typename... Args>
233 unique_ptr<QueryResult> QueryParamsRecursive(const string &query, vector<Value> &values, T value, Args... args) {
234 values.push_back(Value::CreateValue<T>(value));
235 return QueryParamsRecursive(query, values, args...);
236 }
237};
238
239} // namespace duckdb
240