Skip to content

Commit a20a53e

Browse files
akoshchiystaticlibs
akoshchiy
authored and
staticlibs
committed
getProfilingInformation api
1 parent 0ae3204 commit a20a53e

File tree

9 files changed

+102
-0
lines changed

9 files changed

+102
-0
lines changed

‎src/jni/duckdb_java.cpp

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -910,3 +910,34 @@ void _duckdb_jdbc_create_extension_type(JNIEnv *env, jclass, jobject conn_buf) {
910910
byte_test_type_type.SetAlias("byte_test_type");
911911
ExtensionUtil::RegisterType(db_instance, "byte_test_type", byte_test_type_type);
912912
}
913+
914+
static ProfilerPrintFormat GetProfilerPrintFormat(JNIEnv *env, jobject format) {
915+
if (env->IsSameObject(format, J_ProfilerPrintFormat_QUERY_TREE)) {
916+
return ProfilerPrintFormat::QUERY_TREE;
917+
}
918+
if (env->IsSameObject(format, J_ProfilerPrintFormat_JSON)) {
919+
return ProfilerPrintFormat::JSON;
920+
}
921+
if (env->IsSameObject(format, J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER)) {
922+
return ProfilerPrintFormat::QUERY_TREE_OPTIMIZER;
923+
}
924+
if (env->IsSameObject(format, J_ProfilerPrintFormat_NO_OUTPUT)) {
925+
return ProfilerPrintFormat::NO_OUTPUT;
926+
}
927+
if (env->IsSameObject(format, J_ProfilerPrintFormat_HTML)) {
928+
return ProfilerPrintFormat::HTML;
929+
}
930+
if (env->IsSameObject(format, J_ProfilerPrintFormat_GRAPHVIZ)) {
931+
return ProfilerPrintFormat::GRAPHVIZ;
932+
}
933+
}
934+
935+
jstring _duckdb_jdbc_get_profiling_information(JNIEnv *env, jclass, jobject conn_ref_buf, jobject j_format) {
936+
auto connection = get_connection(env, conn_ref_buf);
937+
if (!connection) {
938+
throw InvalidInputException("Invalid connection");
939+
}
940+
auto format = GetProfilerPrintFormat(env, j_format);
941+
auto profiling_info = connection->GetProfilingInformation(format);
942+
return env->NewStringUTF(profiling_info.c_str());
943+
}

‎src/jni/functions.cpp

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -396,3 +396,13 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extens
396396

397397
}
398398
}
399+
400+
JNIEXPORT jstring JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1get_1profiling_1information(JNIEnv * env, jclass param0, jobject param1, jobject param2) {
401+
try {
402+
return _duckdb_jdbc_get_profiling_information(env, param0, param1, param2);
403+
} catch (const std::exception &e) {
404+
duckdb::ErrorData error(e);
405+
ThrowJNI(env, error.Message().c_str());
406+
407+
}
408+
}

‎src/jni/functions.hpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,3 +160,7 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1appe
160160
void _duckdb_jdbc_create_extension_type(JNIEnv * env, jclass param0, jobject param1);
161161

162162
JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extension_1type(JNIEnv * env, jclass param0, jobject param1);
163+
164+
jstring _duckdb_jdbc_get_profiling_information(JNIEnv * env, jclass param0, jobject param1, jobject param2);
165+
166+
JNIEXPORT jstring JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1get_1profiling_1information(JNIEnv * env, jclass param0, jobject param1, jobject param2);

‎src/jni/refs.cpp

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,14 @@ jmethodID J_Object_toString;
100100

101101
jclass J_DuckDBTime;
102102

103+
jclass J_ProfilerPrintFormat;
104+
jobject J_ProfilerPrintFormat_QUERY_TREE;
105+
jobject J_ProfilerPrintFormat_JSON;
106+
jobject J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER;
107+
jobject J_ProfilerPrintFormat_NO_OUTPUT;
108+
jobject J_ProfilerPrintFormat_HTML;
109+
jobject J_ProfilerPrintFormat_GRAPHVIZ;
110+
103111
static std::vector<jobject> global_refs;
104112

105113
template <typename T>
@@ -252,6 +260,18 @@ void create_refs(JNIEnv *env) {
252260
J_DuckVector_varlen = get_field_id(env, J_DuckVector, "varlen_data", "[Ljava/lang/Object;");
253261

254262
J_ByteBuffer = make_class_ref(env, "java/nio/ByteBuffer");
263+
264+
J_ProfilerPrintFormat = make_class_ref(env, "org/duckdb/ProfilerPrintFormat");
265+
J_ProfilerPrintFormat_QUERY_TREE =
266+
make_static_object_field_ref(env, J_ProfilerPrintFormat, "QUERY_TREE", "Lorg/duckdb/ProfilerPrintFormat;");
267+
J_ProfilerPrintFormat_JSON = make_static_object_field_ref(env, J_ProfilerPrintFormat, "JSON", "Lorg/duckdb/ProfilerPrintFormat;");
268+
J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER =
269+
make_static_object_field_ref(env, J_ProfilerPrintFormat, "QUERY_TREE_OPTIMIZER", "Lorg/duckdb/ProfilerPrintFormat;");
270+
J_ProfilerPrintFormat_NO_OUTPUT =
271+
make_static_object_field_ref(env, J_ProfilerPrintFormat, "NO_OUTPUT", "Lorg/duckdb/ProfilerPrintFormat;");
272+
J_ProfilerPrintFormat_HTML = make_static_object_field_ref(env, J_ProfilerPrintFormat, "HTML", "Lorg/duckdb/ProfilerPrintFormat;");
273+
J_ProfilerPrintFormat_GRAPHVIZ =
274+
make_static_object_field_ref(env, J_ProfilerPrintFormat, "GRAPHVIZ", "Lorg/duckdb/ProfilerPrintFormat;");
255275
}
256276

257277
void delete_global_refs(JNIEnv *env) noexcept {

‎src/jni/refs.hpp

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,14 @@ extern jmethodID J_Object_toString;
9797

9898
extern jclass J_DuckDBTime;
9999

100+
extern jclass J_ProfilerPrintFormat;
101+
extern jobject J_ProfilerPrintFormat_QUERY_TREE;
102+
extern jobject J_ProfilerPrintFormat_JSON;
103+
extern jobject J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER;
104+
extern jobject J_ProfilerPrintFormat_NO_OUTPUT;
105+
extern jobject J_ProfilerPrintFormat_HTML;
106+
extern jobject J_ProfilerPrintFormat_GRAPHVIZ;
107+
100108
void create_refs(JNIEnv *env);
101109

102110
void delete_global_refs(JNIEnv *env) noexcept;

‎src/main/java/org/duckdb/DuckDBConnection.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,4 +384,8 @@ public void registerArrowStream(String name, Object arrow_array_stream) {
384384
long array_stream_address = getArrowStreamAddress(arrow_array_stream);
385385
DuckDBNative.duckdb_jdbc_arrow_register(conn_ref, array_stream_address, name.getBytes(StandardCharsets.UTF_8));
386386
}
387+
388+
public String getProfilingInformation(ProfilerPrintFormat format) throws SQLException {
389+
return DuckDBNative.duckdb_jdbc_get_profiling_information(conn_ref, format);
390+
}
387391
}

‎src/main/java/org/duckdb/DuckDBNative.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,10 @@ protected static native void duckdb_jdbc_appender_append_decimal(ByteBuffer appe
167167

168168
protected static native void duckdb_jdbc_create_extension_type(ByteBuffer conn_ref) throws SQLException;
169169

170+
protected static native String duckdb_jdbc_get_profiling_information(ByteBuffer conn_ref,
171+
ProfilerPrintFormat format)
172+
throws SQLException;
173+
170174
public static void duckdb_jdbc_create_extension_type(DuckDBConnection conn) throws SQLException {
171175
duckdb_jdbc_create_extension_type(conn.conn_ref);
172176
}

‎src/main/java/org/duckdb/ProfilerPrintFormat.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
package org.duckdb;
2+
3+
public enum ProfilerPrintFormat {
4+
5+
QUERY_TREE,
6+
JSON,
7+
QUERY_TREE_OPTIMIZER,
8+
NO_OUTPUT,
9+
HTML,
10+
GRAPHVIZ
11+
}

‎src/test/java/org/duckdb/TestDuckDBJDBC.java

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4876,6 +4876,16 @@ public static void test_spark_path_option_ignored() throws Exception {
48764876
conn.close();
48774877
}
48784878

4879+
public static void test_get_profiling_information() throws Exception {
4880+
try (Connection conn = DriverManager.getConnection(JDBC_URL); Statement stmt = conn.createStatement()) {
4881+
stmt.execute("SET enable_profiling = 'no_output';");
4882+
try (ResultSet rs = stmt.executeQuery("SELECT 1+1")) {
4883+
String profile = ((DuckDBConnection) conn).getProfilingInformation(ProfilerPrintFormat.JSON);
4884+
assertTrue(profile.contains("\"query_name\": \"SELECT 1+1\","));
4885+
}
4886+
}
4887+
}
4888+
48794889
public static void main(String[] args) throws Exception {
48804890
String arg1 = args.length > 0 ? args[0] : "";
48814891
final int statusCode;

0 commit comments

Comments
 (0)