1#include "duckdb/function/table/system_functions.hpp"
2
3#include "duckdb/catalog/catalog.hpp"
4#include "duckdb/catalog/catalog_entry/duck_table_entry.hpp"
5#include "duckdb/catalog/catalog_entry/schema_catalog_entry.hpp"
6#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp"
7#include "duckdb/common/exception.hpp"
8#include "duckdb/main/client_context.hpp"
9#include "duckdb/main/client_data.hpp"
10#include "duckdb/parser/constraint.hpp"
11#include "duckdb/parser/constraints/check_constraint.hpp"
12#include "duckdb/parser/constraints/unique_constraint.hpp"
13#include "duckdb/planner/constraints/bound_unique_constraint.hpp"
14#include "duckdb/planner/constraints/bound_check_constraint.hpp"
15#include "duckdb/planner/constraints/bound_not_null_constraint.hpp"
16#include "duckdb/planner/constraints/bound_foreign_key_constraint.hpp"
17#include "duckdb/storage/data_table.hpp"
18
19namespace duckdb {
20
21struct UniqueKeyInfo {
22 string schema;
23 string table;
24 vector<LogicalIndex> columns;
25
26 bool operator==(const UniqueKeyInfo &other) const {
27 return (schema == other.schema) && (table == other.table) && (columns == other.columns);
28 }
29};
30
31} // namespace duckdb
32
33namespace std {
34
35template <>
36struct hash<duckdb::UniqueKeyInfo> {
37 template <class X>
38 static size_t ComputeHash(const X &x) {
39 return hash<X>()(x);
40 }
41
42 size_t operator()(const duckdb::UniqueKeyInfo &j) const {
43 D_ASSERT(j.columns.size() > 0);
44 return ComputeHash(x: j.schema) + ComputeHash(x: j.table) + ComputeHash(x: j.columns[0].index);
45 }
46};
47
48} // namespace std
49
50namespace duckdb {
51
52struct DuckDBConstraintsData : public GlobalTableFunctionState {
53 DuckDBConstraintsData() : offset(0), constraint_offset(0), unique_constraint_offset(0) {
54 }
55
56 vector<reference<CatalogEntry>> entries;
57 idx_t offset;
58 idx_t constraint_offset;
59 idx_t unique_constraint_offset;
60 unordered_map<UniqueKeyInfo, idx_t> known_fk_unique_constraint_offsets;
61};
62
63static unique_ptr<FunctionData> DuckDBConstraintsBind(ClientContext &context, TableFunctionBindInput &input,
64 vector<LogicalType> &return_types, vector<string> &names) {
65 names.emplace_back(args: "database_name");
66 return_types.emplace_back(args: LogicalType::VARCHAR);
67
68 names.emplace_back(args: "database_oid");
69 return_types.emplace_back(args: LogicalType::BIGINT);
70
71 names.emplace_back(args: "schema_name");
72 return_types.emplace_back(args: LogicalType::VARCHAR);
73
74 names.emplace_back(args: "schema_oid");
75 return_types.emplace_back(args: LogicalType::BIGINT);
76
77 names.emplace_back(args: "table_name");
78 return_types.emplace_back(args: LogicalType::VARCHAR);
79
80 names.emplace_back(args: "table_oid");
81 return_types.emplace_back(args: LogicalType::BIGINT);
82
83 names.emplace_back(args: "constraint_index");
84 return_types.emplace_back(args: LogicalType::BIGINT);
85
86 // CHECK, PRIMARY KEY or UNIQUE
87 names.emplace_back(args: "constraint_type");
88 return_types.emplace_back(args: LogicalType::VARCHAR);
89
90 names.emplace_back(args: "constraint_text");
91 return_types.emplace_back(args: LogicalType::VARCHAR);
92
93 names.emplace_back(args: "expression");
94 return_types.emplace_back(args: LogicalType::VARCHAR);
95
96 names.emplace_back(args: "constraint_column_indexes");
97 return_types.push_back(x: LogicalType::LIST(child: LogicalType::BIGINT));
98
99 names.emplace_back(args: "constraint_column_names");
100 return_types.push_back(x: LogicalType::LIST(child: LogicalType::VARCHAR));
101
102 return nullptr;
103}
104
105unique_ptr<GlobalTableFunctionState> DuckDBConstraintsInit(ClientContext &context, TableFunctionInitInput &input) {
106 auto result = make_uniq<DuckDBConstraintsData>();
107
108 // scan all the schemas for tables and collect them
109 auto schemas = Catalog::GetAllSchemas(context);
110
111 for (auto &schema : schemas) {
112 vector<reference<CatalogEntry>> entries;
113
114 schema.get().Scan(context, type: CatalogType::TABLE_ENTRY, callback: [&](CatalogEntry &entry) {
115 if (entry.type == CatalogType::TABLE_ENTRY) {
116 entries.push_back(x: entry);
117 }
118 });
119
120 sort(first: entries.begin(), last: entries.end(), comp: [&](CatalogEntry &x, CatalogEntry &y) { return (x.name < y.name); });
121
122 result->entries.insert(position: result->entries.end(), first: entries.begin(), last: entries.end());
123 };
124
125 return std::move(result);
126}
127
128void DuckDBConstraintsFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) {
129 auto &data = data_p.global_state->Cast<DuckDBConstraintsData>();
130 if (data.offset >= data.entries.size()) {
131 // finished returning values
132 return;
133 }
134 // start returning values
135 // either fill up the chunk or return all the remaining columns
136 idx_t count = 0;
137 while (data.offset < data.entries.size() && count < STANDARD_VECTOR_SIZE) {
138 auto &entry = data.entries[data.offset].get();
139 D_ASSERT(entry.type == CatalogType::TABLE_ENTRY);
140
141 auto &table = entry.Cast<TableCatalogEntry>();
142 auto &constraints = table.GetConstraints();
143 bool is_duck_table = table.IsDuckTable();
144 for (; data.constraint_offset < constraints.size() && count < STANDARD_VECTOR_SIZE; data.constraint_offset++) {
145 auto &constraint = constraints[data.constraint_offset];
146 // return values:
147 // constraint_type, VARCHAR
148 // Processing this first due to shortcut (early continue)
149 string constraint_type;
150 switch (constraint->type) {
151 case ConstraintType::CHECK:
152 constraint_type = "CHECK";
153 break;
154 case ConstraintType::UNIQUE: {
155 auto &unique = constraint->Cast<UniqueConstraint>();
156 constraint_type = unique.is_primary_key ? "PRIMARY KEY" : "UNIQUE";
157 break;
158 }
159 case ConstraintType::NOT_NULL:
160 constraint_type = "NOT NULL";
161 break;
162 case ConstraintType::FOREIGN_KEY: {
163 if (!is_duck_table) {
164 continue;
165 }
166 auto &bound_constraints = table.GetBoundConstraints();
167 auto &bound_foreign_key = bound_constraints[data.constraint_offset]->Cast<BoundForeignKeyConstraint>();
168 if (bound_foreign_key.info.type == ForeignKeyType::FK_TYPE_PRIMARY_KEY_TABLE) {
169 // Those are already covered by PRIMARY KEY and UNIQUE entries
170 continue;
171 }
172 constraint_type = "FOREIGN KEY";
173 break;
174 }
175 default:
176 throw NotImplementedException("Unimplemented constraint for duckdb_constraints");
177 }
178
179 idx_t col = 0;
180 // database_name, LogicalType::VARCHAR
181 output.SetValue(col_idx: col++, index: count, val: Value(table.schema.catalog.GetName()));
182 // database_oid, LogicalType::BIGINT
183 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: table.schema.catalog.GetOid()));
184 // schema_name, LogicalType::VARCHAR
185 output.SetValue(col_idx: col++, index: count, val: Value(table.schema.name));
186 // schema_oid, LogicalType::BIGINT
187 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: table.schema.oid));
188 // table_name, LogicalType::VARCHAR
189 output.SetValue(col_idx: col++, index: count, val: Value(table.name));
190 // table_oid, LogicalType::BIGINT
191 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: table.oid));
192
193 // constraint_index, BIGINT
194 UniqueKeyInfo uk_info;
195
196 if (is_duck_table) {
197 auto &bound_constraint = *table.GetBoundConstraints()[data.constraint_offset];
198 switch (bound_constraint.type) {
199 case ConstraintType::UNIQUE: {
200 auto &bound_unique = bound_constraint.Cast<BoundUniqueConstraint>();
201 uk_info = {.schema: table.schema.name, .table: table.name, .columns: bound_unique.keys};
202 break;
203 }
204 case ConstraintType::FOREIGN_KEY: {
205 const auto &bound_foreign_key = bound_constraint.Cast<BoundForeignKeyConstraint>();
206 const auto &info = bound_foreign_key.info;
207 // find the other table
208 auto table_entry = Catalog::GetEntry<TableCatalogEntry>(
209 context, catalog_name: table.catalog.GetName(), schema_name: info.schema, name: info.table, if_not_found: OnEntryNotFound::RETURN_NULL);
210 if (!table_entry) {
211 throw InternalException("dukdb_constraints: entry %s.%s referenced in foreign key not found",
212 info.schema, info.table);
213 }
214 vector<LogicalIndex> index;
215 for (auto &key : info.pk_keys) {
216 index.push_back(x: table_entry->GetColumns().PhysicalToLogical(index: key));
217 }
218 uk_info = {.schema: table_entry->schema.name, .table: table_entry->name, .columns: index};
219 break;
220 }
221 default:
222 break;
223 }
224 }
225
226 if (uk_info.columns.empty()) {
227 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: data.unique_constraint_offset++));
228 } else {
229 auto known_unique_constraint_offset = data.known_fk_unique_constraint_offsets.find(x: uk_info);
230 if (known_unique_constraint_offset == data.known_fk_unique_constraint_offsets.end()) {
231 data.known_fk_unique_constraint_offsets.insert(x: make_pair(x&: uk_info, y&: data.unique_constraint_offset));
232 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: data.unique_constraint_offset));
233 data.unique_constraint_offset++;
234 } else {
235 output.SetValue(col_idx: col++, index: count, val: Value::BIGINT(value: known_unique_constraint_offset->second));
236 }
237 }
238 output.SetValue(col_idx: col++, index: count, val: Value(constraint_type));
239
240 // constraint_text, VARCHAR
241 output.SetValue(col_idx: col++, index: count, val: Value(constraint->ToString()));
242
243 // expression, VARCHAR
244 Value expression_text;
245 if (constraint->type == ConstraintType::CHECK) {
246 auto &check = constraint->Cast<CheckConstraint>();
247 expression_text = Value(check.expression->ToString());
248 }
249 output.SetValue(col_idx: col++, index: count, val: expression_text);
250
251 vector<LogicalIndex> column_index_list;
252 if (is_duck_table) {
253 auto &bound_constraint = *table.GetBoundConstraints()[data.constraint_offset];
254 switch (bound_constraint.type) {
255 case ConstraintType::CHECK: {
256 auto &bound_check = bound_constraint.Cast<BoundCheckConstraint>();
257 for (auto &col_idx : bound_check.bound_columns) {
258 column_index_list.push_back(x: table.GetColumns().PhysicalToLogical(index: col_idx));
259 }
260 break;
261 }
262 case ConstraintType::UNIQUE: {
263 auto &bound_unique = bound_constraint.Cast<BoundUniqueConstraint>();
264 for (auto &col_idx : bound_unique.keys) {
265 column_index_list.push_back(x: col_idx);
266 }
267 break;
268 }
269 case ConstraintType::NOT_NULL: {
270 auto &bound_not_null = bound_constraint.Cast<BoundNotNullConstraint>();
271 column_index_list.push_back(x: table.GetColumns().PhysicalToLogical(index: bound_not_null.index));
272 break;
273 }
274 case ConstraintType::FOREIGN_KEY: {
275 auto &bound_foreign_key = bound_constraint.Cast<BoundForeignKeyConstraint>();
276 for (auto &col_idx : bound_foreign_key.info.fk_keys) {
277 column_index_list.push_back(x: table.GetColumns().PhysicalToLogical(index: col_idx));
278 }
279 break;
280 }
281 default:
282 throw NotImplementedException("Unimplemented constraint for duckdb_constraints");
283 }
284 }
285
286 vector<Value> index_list;
287 vector<Value> column_name_list;
288 for (auto column_index : column_index_list) {
289 index_list.push_back(x: Value::BIGINT(value: column_index.index));
290 column_name_list.emplace_back(args: table.GetColumn(idx: column_index).Name());
291 }
292
293 // constraint_column_indexes, LIST
294 output.SetValue(col_idx: col++, index: count, val: Value::LIST(child_type: LogicalType::BIGINT, values: std::move(index_list)));
295
296 // constraint_column_names, LIST
297 output.SetValue(col_idx: col++, index: count, val: Value::LIST(child_type: LogicalType::VARCHAR, values: std::move(column_name_list)));
298
299 count++;
300 }
301 if (data.constraint_offset >= constraints.size()) {
302 data.constraint_offset = 0;
303 data.offset++;
304 }
305 }
306 output.SetCardinality(count);
307}
308
309void DuckDBConstraintsFun::RegisterFunction(BuiltinFunctions &set) {
310 set.AddFunction(function: TableFunction("duckdb_constraints", {}, DuckDBConstraintsFunction, DuckDBConstraintsBind,
311 DuckDBConstraintsInit));
312}
313
314} // namespace duckdb
315