1#include "duckdb/execution/operator/helper/physical_pragma.hpp"
2
3#include "duckdb/main/client_context.hpp"
4#include "duckdb/main/database.hpp"
5#include "duckdb/storage/storage_manager.hpp"
6#include "duckdb/storage/buffer_manager.hpp"
7#include "duckdb/common/operator/cast_operators.hpp"
8#include "duckdb/planner/expression_binder.hpp"
9
10#include <cctype>
11
12using namespace duckdb;
13using namespace std;
14
15static idx_t ParseMemoryLimit(string arg);
16
17void PhysicalPragma::GetChunkInternal(ClientContext &context, DataChunk &chunk, PhysicalOperatorState *state) {
18 auto &pragma = *info;
19 auto &keyword = pragma.name;
20 if (keyword == "enable_profile" || keyword == "enable_profiling") {
21 // enable profiling
22 if (pragma.pragma_type == PragmaType::ASSIGNMENT) {
23 // enable_profiling with assignment
24 // this is either enable_profiling = json, or enable_profiling = query_tree
25 string assignment = pragma.parameters[0].ToString();
26 if (assignment == "json") {
27 context.profiler.automatic_print_format = ProfilerPrintFormat::JSON;
28 } else if (assignment == "query_tree") {
29 context.profiler.automatic_print_format = ProfilerPrintFormat::QUERY_TREE;
30 } else {
31 throw ParserException("Unrecognized print format %s, supported formats: [json, query_tree]",
32 assignment.c_str());
33 }
34 } else if (pragma.pragma_type == PragmaType::NOTHING) {
35 context.profiler.automatic_print_format = ProfilerPrintFormat::QUERY_TREE;
36 } else {
37 throw ParserException("Cannot call PRAGMA enable_profiling");
38 }
39 context.profiler.Enable();
40 } else if (keyword == "disable_profile" || keyword == "disable_profiling") {
41 if (pragma.pragma_type != PragmaType::NOTHING) {
42 throw ParserException("disable_profiling cannot take parameters!");
43 }
44 // enable profiling
45 context.profiler.Disable();
46 context.profiler.automatic_print_format = ProfilerPrintFormat::NONE;
47 } else if (keyword == "profiling_output" || keyword == "profile_output") {
48 // set file location of where to save profiling output
49 if (pragma.pragma_type != PragmaType::ASSIGNMENT || pragma.parameters[0].type != TypeId::VARCHAR) {
50 throw ParserException(
51 "Profiling output must be an assignment (e.g. PRAGMA profile_output='/tmp/test.json')");
52 }
53 context.profiler.save_location = pragma.parameters[0].str_value;
54 } else if (keyword == "memory_limit") {
55 if (pragma.pragma_type != PragmaType::ASSIGNMENT) {
56 throw ParserException("Memory limit must be an assignment (e.g. PRAGMA memory_limit='1GB')");
57 }
58 if (pragma.parameters[0].type == TypeId::VARCHAR) {
59 idx_t new_limit = ParseMemoryLimit(pragma.parameters[0].str_value);
60 // set the new limit in the buffer manager
61 context.db.storage->buffer_manager->SetLimit(new_limit);
62 } else {
63 int64_t value = pragma.parameters[0].GetValue<int64_t>();
64 if (value < 0) {
65 // limit < 0, set limit to infinite
66 context.db.storage->buffer_manager->SetLimit();
67 } else {
68 throw ParserException(
69 "Memory limit must be an assignment with a memory unit (e.g. PRAGMA memory_limit='1GB')");
70 }
71 }
72 } else if (keyword == "collation" || keyword == "default_collation") {
73 if (pragma.pragma_type != PragmaType::ASSIGNMENT) {
74 throw ParserException("Collation must be an assignment (e.g. PRAGMA default_collation=NOCASE)");
75 }
76 auto collation_param = StringUtil::Lower(pragma.parameters[0].CastAs(TypeId::VARCHAR).str_value);
77 // bind the collation to verify that it exists
78 ExpressionBinder::PushCollation(context, nullptr, collation_param);
79 context.db.collation = collation_param;
80 } else {
81 throw ParserException("Unrecognized PRAGMA keyword: %s", keyword.c_str());
82 }
83}
84
85idx_t ParseMemoryLimit(string arg) {
86 // split based on the number/non-number
87 idx_t idx = 0;
88 while (std::isspace(arg[idx])) {
89 idx++;
90 }
91 idx_t num_start = idx;
92 while ((arg[idx] >= '0' && arg[idx] <= '9') || arg[idx] == '.' || arg[idx] == 'e' || arg[idx] == 'E' ||
93 arg[idx] == '-') {
94 idx++;
95 }
96 if (idx == num_start) {
97 throw ParserException("Memory limit must have a number (e.g. PRAGMA memory_limit=1GB");
98 }
99 string number = arg.substr(num_start, idx - num_start);
100
101 // try to parse the number
102 double limit = Cast::Operation<string_t, double>(number.c_str());
103
104 // now parse the memory limit unit (e.g. bytes, gb, etc)
105 while (std::isspace(arg[idx])) {
106 idx++;
107 }
108 idx_t start = idx;
109 while (idx < arg.size() && !std::isspace(arg[idx])) {
110 idx++;
111 }
112 if (limit < 0) {
113 // limit < 0, set limit to infinite
114 return (idx_t)-1;
115 }
116 string unit = StringUtil::Lower(arg.substr(start, idx - start));
117 idx_t multiplier;
118 if (unit == "byte" || unit == "bytes" || unit == "b") {
119 multiplier = 1;
120 } else if (unit == "kilobyte" || unit == "kilobytes" || unit == "kb" || unit == "k") {
121 multiplier = 1000LL;
122 } else if (unit == "megabyte" || unit == "megabytes" || unit == "mb" || unit == "m") {
123 multiplier = 1000LL * 1000LL;
124 } else if (unit == "gigabyte" || unit == "gigabytes" || unit == "gb" || unit == "g") {
125 multiplier = 1000LL * 1000LL * 1000LL;
126 } else if (unit == "terabyte" || unit == "terabytes" || unit == "tb" || unit == "t") {
127 multiplier = 1000LL * 1000LL * 1000LL * 1000LL;
128 } else {
129 throw ParserException("Unknown unit for memory_limit: %s (expected: b, mb, gb or tb)", unit.c_str());
130 }
131 return (idx_t)multiplier * limit;
132}
133