Merge pull request #9430 from Mause/c-api-extensions

feat(c): add functions for determining statement/return types
duckdb · Nov 20, 2023 · ce5a2ec · ce5a2ec
2 parents 3eb29b7 + d454ee9
commit ce5a2ec
Show file tree

Hide file tree

Showing 8 changed files with 251 additions and 0 deletions.
diff --git a/src/include/duckdb.h b/src/include/duckdb.h
@@ -312,6 +312,45 @@ typedef enum {
 	DUCKDB_PENDING_NO_TASKS_AVAILABLE = 3
 } duckdb_pending_state;
 
+typedef enum {
+	DUCKDB_RESULT_TYPE_INVALID,
+	DUCKDB_RESULT_TYPE_CHANGED_ROWS,
+	DUCKDB_RESULT_TYPE_NOTHING,
+	DUCKDB_RESULT_TYPE_QUERY_RESULT,
+} duckdb_result_type;
+
+typedef enum {
+	DUCKDB_STATEMENT_TYPE_INVALID,
+	DUCKDB_STATEMENT_TYPE_SELECT,
+	DUCKDB_STATEMENT_TYPE_INSERT,
+	DUCKDB_STATEMENT_TYPE_UPDATE,
+	DUCKDB_STATEMENT_TYPE_EXPLAIN,
+	DUCKDB_STATEMENT_TYPE_DELETE,
+	DUCKDB_STATEMENT_TYPE_PREPARE,
+	DUCKDB_STATEMENT_TYPE_CREATE,
+	DUCKDB_STATEMENT_TYPE_EXECUTE,
+	DUCKDB_STATEMENT_TYPE_ALTER,
+	DUCKDB_STATEMENT_TYPE_TRANSACTION,
+	DUCKDB_STATEMENT_TYPE_COPY,
+	DUCKDB_STATEMENT_TYPE_ANALYZE,
+	DUCKDB_STATEMENT_TYPE_VARIABLE_SET,
+	DUCKDB_STATEMENT_TYPE_CREATE_FUNC,
+	DUCKDB_STATEMENT_TYPE_DROP,
+	DUCKDB_STATEMENT_TYPE_EXPORT,
+	DUCKDB_STATEMENT_TYPE_PRAGMA,
+	DUCKDB_STATEMENT_TYPE_SHOW,
+	DUCKDB_STATEMENT_TYPE_VACUUM,
+	DUCKDB_STATEMENT_TYPE_CALL,
+	DUCKDB_STATEMENT_TYPE_SET,
+	DUCKDB_STATEMENT_TYPE_LOAD,
+	DUCKDB_STATEMENT_TYPE_RELATION,
+	DUCKDB_STATEMENT_TYPE_EXTENSION,
+	DUCKDB_STATEMENT_TYPE_LOGICAL_PLAN,
+	DUCKDB_STATEMENT_TYPE_ATTACH,
+	DUCKDB_STATEMENT_TYPE_DETACH,
+	DUCKDB_STATEMENT_TYPE_MULTI,
+} duckdb_statement_type;
+
 //===--------------------------------------------------------------------===//
 // Open/Connect
 //===--------------------------------------------------------------------===//
@@ -496,6 +535,14 @@ Returns `DUCKDB_TYPE_INVALID` if the column is out of range.
 */
 DUCKDB_API duckdb_type duckdb_column_type(duckdb_result *result, idx_t col);
 
+/*!
+Returns the statement type of the statement that was executed
+
+* result: The result object to fetch the statement type from.
+ * returns: duckdb_statement_type value or DUCKDB_STATEMENT_TYPE_INVALID
+ */
+DUCKDB_API duckdb_statement_type duckdb_result_statement_type(duckdb_result result);
+
 /*!
 Returns the logical column type of the specified column.
 
@@ -627,6 +674,14 @@ Returns the number of data chunks present in the result.
 */
 DUCKDB_API idx_t duckdb_result_chunk_count(duckdb_result result);
 
+/*!
+Returns the return_type of the given result, or DUCKDB_RETURN_TYPE_INVALID on error
+
+* result: The result object
+* returns: The return_type
+ */
+DUCKDB_API duckdb_result_type duckdb_result_return_type(duckdb_result result);
+
 // Safe fetch functions
 // These functions will perform conversions if necessary.
 // On failure (e.g. if conversion cannot be performed or if the value is NULL) a default value is returned.
@@ -971,6 +1026,14 @@ Clear the params bind to the prepared statement.
 */
 DUCKDB_API duckdb_state duckdb_clear_bindings(duckdb_prepared_statement prepared_statement);
 
+/*!
+Returns the statement type of the statement to be executed
+
+ * statement: The prepared statement.
+ * returns: duckdb_statement_type value or DUCKDB_STATEMENT_TYPE_INVALID
+ */
+DUCKDB_API duckdb_statement_type duckdb_prepared_statement_type(duckdb_prepared_statement statement);
+
 /*!
 Binds a value to the prepared statement at the specified index.
 */
@@ -1343,6 +1406,15 @@ This should not be used with `DUCKDB_TYPE_DECIMAL`.
 */
 DUCKDB_API duckdb_logical_type duckdb_create_logical_type(duckdb_type type);
 
+/*!
+Returns the alias of a duckdb_logical_type, if one is set, else `NULL`
+You must free the result.
+
+* type: The logical type to return the alias of
+* returns: The alias or `NULL`
+ */
+DUCKDB_API char *duckdb_logical_type_get_alias(duckdb_logical_type type);
+
 /*!
 Creates a list type from its child type.
 The resulting type should be destroyed with `duckdb_destroy_logical_type`.

diff --git a/src/include/duckdb/main/capi/capi_internal.hpp b/src/include/duckdb/main/capi/capi_internal.hpp
@@ -76,5 +76,6 @@ LogicalTypeId ConvertCTypeToCPP(duckdb_type c_type);
 idx_t GetCTypeSize(duckdb_type type);
 duckdb_state duckdb_translate_result(unique_ptr<QueryResult> result, duckdb_result *out);
 bool deprecated_materialize_result(duckdb_result *result);
+duckdb_statement_type StatementTypeToC(duckdb::StatementType statement_type);
 
 } // namespace duckdb
diff --git a/src/main/capi/helper-c.cpp b/src/main/capi/helper-c.cpp
@@ -173,6 +173,71 @@ idx_t GetCTypeSize(duckdb_type type) {
 	} // LCOV_EXCL_STOP
 }
 
+duckdb_statement_type StatementTypeToC(duckdb::StatementType statement_type) {
+	switch (statement_type) {
+	case duckdb::StatementType::SELECT_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_SELECT;
+	case duckdb::StatementType::INVALID_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_INVALID;
+	case duckdb::StatementType::INSERT_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_INSERT;
+	case duckdb::StatementType::UPDATE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_UPDATE;
+	case duckdb::StatementType::EXPLAIN_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_EXPLAIN;
+	case duckdb::StatementType::DELETE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_DELETE;
+	case duckdb::StatementType::PREPARE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_PREPARE;
+	case duckdb::StatementType::CREATE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_CREATE;
+	case duckdb::StatementType::EXECUTE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_EXECUTE;
+	case duckdb::StatementType::ALTER_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_ALTER;
+	case duckdb::StatementType::TRANSACTION_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_TRANSACTION;
+	case duckdb::StatementType::COPY_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_COPY;
+	case duckdb::StatementType::ANALYZE_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_ANALYZE;
+	case duckdb::StatementType::VARIABLE_SET_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_VARIABLE_SET;
+	case duckdb::StatementType::CREATE_FUNC_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_CREATE_FUNC;
+	case duckdb::StatementType::DROP_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_DROP;
+	case duckdb::StatementType::EXPORT_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_EXPORT;
+	case duckdb::StatementType::PRAGMA_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_PRAGMA;
+	case duckdb::StatementType::SHOW_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_SHOW;
+	case duckdb::StatementType::VACUUM_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_VACUUM;
+	case duckdb::StatementType::CALL_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_CALL;
+	case duckdb::StatementType::SET_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_SET;
+	case duckdb::StatementType::LOAD_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_LOAD;
+	case duckdb::StatementType::RELATION_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_RELATION;
+	case duckdb::StatementType::EXTENSION_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_EXTENSION;
+	case duckdb::StatementType::LOGICAL_PLAN_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_LOGICAL_PLAN;
+	case duckdb::StatementType::ATTACH_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_ATTACH;
+	case duckdb::StatementType::DETACH_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_DETACH;
+	case duckdb::StatementType::MULTI_STATEMENT:
+		return DUCKDB_STATEMENT_TYPE_MULTI;
+	default:
+		return DUCKDB_STATEMENT_TYPE_INVALID;
+	}
+}
+
 } // namespace duckdb
 
 void *duckdb_malloc(size_t size) {

diff --git a/src/main/capi/logical_types-c.cpp b/src/main/capi/logical_types-c.cpp
@@ -276,6 +276,11 @@ char *duckdb_struct_type_child_name(duckdb_logical_type type, idx_t index) {
 	return strdup(duckdb::StructType::GetChildName(ltype, index).c_str());
 }
 
+char *duckdb_logical_type_get_alias(duckdb_logical_type type) {
+	auto &ltype = *(reinterpret_cast<duckdb::LogicalType *>(type));
+	return ltype.HasAlias() ? strdup(ltype.GetAlias().c_str()) : nullptr;
+}
+
 duckdb_logical_type duckdb_struct_type_child_type(duckdb_logical_type type, idx_t index) {
 	if (!AssertInternalType(type, duckdb::PhysicalType::STRUCT)) {
 		return nullptr;

diff --git a/src/main/capi/prepared-c.cpp b/src/main/capi/prepared-c.cpp
@@ -320,6 +320,15 @@ duckdb_state duckdb_execute_prepared(duckdb_prepared_statement prepared_statemen
 	return duckdb_translate_result(std::move(result), out_result);
 }
 
+duckdb_statement_type duckdb_prepared_statement_type(duckdb_prepared_statement statement) {
+	if (!statement) {
+		return DUCKDB_STATEMENT_TYPE_INVALID;
+	}
+	auto stmt = reinterpret_cast<PreparedStatementWrapper *>(statement);
+
+	return StatementTypeToC(stmt->statement->GetStatementType());
+}
+
 template <class T>
 void duckdb_destroy(void **wrapper) {
 	if (!wrapper) {

diff --git a/src/main/capi/result-c.cpp b/src/main/capi/result-c.cpp
@@ -515,3 +515,29 @@ bool duckdb_result_is_streaming(duckdb_result result) {
 	auto &result_data = *(reinterpret_cast<duckdb::DuckDBResultData *>(result.internal_data));
 	return result_data.result->type == duckdb::QueryResultType::STREAM_RESULT;
 }
+
+duckdb_result_type duckdb_result_return_type(duckdb_result result) {
+	if (!result.internal_data || duckdb_result_error(&result) != nullptr) {
+		return DUCKDB_RESULT_TYPE_INVALID;
+	}
+	auto &result_data = *(reinterpret_cast<duckdb::DuckDBResultData *>(result.internal_data));
+	switch (result_data.result->properties.return_type) {
+	case duckdb::StatementReturnType::CHANGED_ROWS:
+		return DUCKDB_RESULT_TYPE_CHANGED_ROWS;
+	case duckdb::StatementReturnType::NOTHING:
+		return DUCKDB_RESULT_TYPE_NOTHING;
+	case duckdb::StatementReturnType::QUERY_RESULT:
+		return DUCKDB_RESULT_TYPE_QUERY_RESULT;
+	default:
+		return DUCKDB_RESULT_TYPE_INVALID;
+	}
+}
+
+duckdb_statement_type duckdb_result_statement_type(duckdb_result result) {
+	if (!result.internal_data || duckdb_result_error(&result) != nullptr) {
+		return DUCKDB_STATEMENT_TYPE_INVALID;
+	}
+	auto &pres = *(reinterpret_cast<duckdb::DuckDBResultData *>(result.internal_data));
+
+	return StatementTypeToC(pres.result->statement_type);
+}
diff --git a/test/api/capi/test_capi_complex_types.cpp b/test/api/capi/test_capi_complex_types.cpp
@@ -1,4 +1,6 @@
 #include "capi_tester.hpp"
+#include "duckdb/main/database_manager.hpp"
+#include "duckdb/parser/parsed_data/create_type_info.hpp"
 
 using namespace duckdb;
 using namespace std;
@@ -218,3 +220,58 @@ TEST_CASE("Test enum types creation C API", "[capi]") {
 
 	REQUIRE(duckdb_create_enum_type(nullptr, 0) == nullptr);
 }
+
+TEST_CASE("Logical types with aliases", "[capi]") {
+	CAPITester tester;
+
+	REQUIRE(tester.OpenDatabase(nullptr));
+
+	Connection *connection = reinterpret_cast<Connection *>(tester.connection);
+
+	connection->BeginTransaction();
+
+	child_list_t<LogicalType> children = {{"hello", LogicalType::VARCHAR}};
+	auto id = LogicalType::STRUCT(children);
+	auto type_name = "test_type";
+	id.SetAlias(type_name);
+	CreateTypeInfo info(type_name, id);
+
+	auto &catalog_name = DatabaseManager::GetDefaultDatabase(*connection->context);
+	auto &catalog = Catalog::GetCatalog(*connection->context, catalog_name);
+	catalog.CreateType(*connection->context, info);
+
+	connection->Commit();
+
+	auto result = tester.Query("SELECT {hello: 'world'}::test_type");
+
+	REQUIRE(NO_FAIL(*result));
+
+	auto chunk = result->FetchChunk(0);
+	REQUIRE(chunk);
+
+	for (idx_t i = 0; i < result->ColumnCount(); i++) {
+		auto logical_type = duckdb_vector_get_column_type(duckdb_data_chunk_get_vector(chunk->GetChunk(), i));
+		REQUIRE(logical_type);
+
+		auto alias = duckdb_logical_type_get_alias(logical_type);
+		REQUIRE(string(alias) == "test_type");
+		duckdb_free(alias);
+
+		duckdb_destroy_logical_type(&logical_type);
+	}
+}
+
+TEST_CASE("Statement types", "[capi]") {
+	CAPITester tester;
+	REQUIRE(tester.OpenDatabase(nullptr));
+
+	duckdb_prepared_statement prepared;
+	REQUIRE(duckdb_prepare(tester.connection, "select ?", &prepared) == DuckDBSuccess);
+
+	REQUIRE(duckdb_prepared_statement_type(prepared) == DUCKDB_STATEMENT_TYPE_SELECT);
+	duckdb_destroy_prepare(&prepared);
+
+	auto result = tester.Query("CREATE TABLE t1 (id int)");
+
+	REQUIRE(duckdb_result_statement_type(result->InternalResult()) == DUCKDB_STATEMENT_TYPE_CREATE);
+}
diff --git a/test/api/capi/test_capi_data_chunk.cpp b/test/api/capi/test_capi_data_chunk.cpp
@@ -313,6 +313,22 @@ TEST_CASE("Test DataChunk result fetch in C API", "[capi]") {
 	REQUIRE(!chunk);
 }
 
+TEST_CASE("Test duckdb_result_return_type", "[capi]") {
+	CAPITester tester;
+	duckdb::unique_ptr<CAPIResult> result;
+
+	REQUIRE(tester.OpenDatabase(nullptr));
+
+	result = tester.Query("CREATE TABLE t (id INT)");
+	REQUIRE(duckdb_result_return_type(result->InternalResult()) == DUCKDB_RESULT_TYPE_NOTHING);
+
+	result = tester.Query("INSERT INTO t VALUES (42)");
+	REQUIRE(duckdb_result_return_type(result->InternalResult()) == DUCKDB_RESULT_TYPE_CHANGED_ROWS);
+
+	result = tester.Query("FROM t");
+	REQUIRE(duckdb_result_return_type(result->InternalResult()) == DUCKDB_RESULT_TYPE_QUERY_RESULT);
+}
+
 TEST_CASE("Test DataChunk populate ListVector in C API", "[capi]") {
 	REQUIRE(duckdb_list_vector_reserve(nullptr, 100) == duckdb_state::DuckDBError);
 	REQUIRE(duckdb_list_vector_set_size(nullptr, 200) == duckdb_state::DuckDBError);