getProfilingInformation api

duckdb · Mar 4, 2025 · a87eb0f · a87eb0f
1 parent eb3c22c
commit a87eb0f
Show file tree

Hide file tree

Showing 8 changed files with 117 additions and 9 deletions.
diff --git a/.github/workflows/Java.yml b/.github/workflows/Java.yml
@@ -63,7 +63,7 @@ jobs:
         run: |
           cp build/release/duckdb_jdbc.jar duckdb_jdbc-linux-amd64.jar
           ./scripts/upload-assets-to-staging.sh github_release duckdb_jdbc-linux-amd64.jar
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: java-linux-amd64
           path: |
@@ -108,7 +108,7 @@ jobs:
           cp build/release/duckdb_jdbc.jar duckdb_jdbc-linux-aarch64.jar
          # ./scripts/upload-assets-to-staging.sh github_release duckdb_jdbc-linux-aarch64.jar
 
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: java-linux-aarch64
           path: |
@@ -150,7 +150,7 @@ jobs:
         run: |
           cp build/release/duckdb_jdbc.jar duckdb_jdbc-windows-amd64.jar
           ./scripts/upload-assets-to-staging.sh github_release duckdb_jdbc-windows-amd64.jar
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: java-windows-amd64
           path: |
@@ -192,7 +192,7 @@ jobs:
         run: |
           cp build/release/duckdb_jdbc.jar duckdb_jdbc-osx-universal.jar
           ./scripts/upload-assets-to-staging.sh github_release duckdb_jdbc-osx-universal.jar
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: java-osx-universal
           path: |
@@ -218,22 +218,22 @@ jobs:
       - shell: bash
         run: mkdir jdbc-artifacts
 
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: java-linux-aarch64
           path: jdbc-artifacts/java-linux-aarch64
 
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: java-linux-amd64
           path: jdbc-artifacts/java-linux-amd64
 
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: java-windows-amd64
           path: jdbc-artifacts/java-windows-amd64
 
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: java-osx-universal
           path: jdbc-artifacts/java-osx-universal
@@ -267,7 +267,7 @@ jobs:
 
       - name: Upload artifacts
         if: always()
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: java-jars
           path: |

diff --git a/src/jni/duckdb_java.cpp b/src/jni/duckdb_java.cpp
@@ -105,6 +105,13 @@ static jmethodID J_Object_toString;
 
 static jclass J_DuckDBTime;
 
+static jobject J_ProfilerPrintFormat_QUERY_TREE;
+static jobject J_ProfilerPrintFormat_JSON;
+static jobject J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER;
+static jobject J_ProfilerPrintFormat_NO_OUTPUT;
+static jobject J_ProfilerPrintFormat_HTML;
+static jobject J_ProfilerPrintFormat_GRAPHVIZ;
+
 void ThrowJNI(JNIEnv *env, const char *message) {
 	D_ASSERT(J_SQLException);
 	env->ThrowNew(J_SQLException, message);
@@ -123,6 +130,12 @@ static jclass GetClassRef(JNIEnv *env, const string &name) {
 	return globalRef;
 }
 
+static jobject GetObjectStaticField(JNIEnv *env, jclass clazz, const char *name, const char *sig) {
+	auto field = env->GetStaticFieldID(clazz, name, sig);
+	auto obj = env->GetStaticObjectField(clazz, field);
+	return env->NewGlobalRef(obj);
+}
+
 JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *reserved) {
 	// Get JNIEnv from vm
 	JNIEnv *env;
@@ -291,6 +304,21 @@ JNIEXPORT jint JNICALL JNI_OnLoad(JavaVM *vm, void *reserved) {
 	J_Object_toString = env->GetMethodID(tmpLocalRef, "toString", "()Ljava/lang/String;");
 	env->DeleteLocalRef(tmpLocalRef);
 
+	tmpLocalRef = env->FindClass("org/duckdb/ProfilerPrintFormat");
+
+	J_ProfilerPrintFormat_QUERY_TREE =
+	    GetObjectStaticField(env, tmpLocalRef, "QUERY_TREE", "Lorg/duckdb/ProfilerPrintFormat;");
+	J_ProfilerPrintFormat_JSON = GetObjectStaticField(env, tmpLocalRef, "JSON", "Lorg/duckdb/ProfilerPrintFormat;");
+	J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER =
+	    GetObjectStaticField(env, tmpLocalRef, "QUERY_TREE_OPTIMIZER", "Lorg/duckdb/ProfilerPrintFormat;");
+	J_ProfilerPrintFormat_NO_OUTPUT =
+	    GetObjectStaticField(env, tmpLocalRef, "NO_OUTPUT", "Lorg/duckdb/ProfilerPrintFormat;");
+	J_ProfilerPrintFormat_HTML = GetObjectStaticField(env, tmpLocalRef, "HTML", "Lorg/duckdb/ProfilerPrintFormat;");
+	J_ProfilerPrintFormat_GRAPHVIZ =
+	    GetObjectStaticField(env, tmpLocalRef, "GRAPHVIZ", "Lorg/duckdb/ProfilerPrintFormat;");
+
+	env->DeleteLocalRef(tmpLocalRef);
+
 	return JNI_VERSION;
 }
 
@@ -317,6 +345,12 @@ JNIEXPORT void JNICALL JNI_OnUnload(JavaVM *vm, void *reserved) {
 	env->DeleteGlobalRef(J_DuckResultSetMeta);
 	env->DeleteGlobalRef(J_DuckVector);
 	env->DeleteGlobalRef(J_ByteBuffer);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_QUERY_TREE);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_JSON);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_NO_OUTPUT);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_HTML);
+	env->DeleteGlobalRef(J_ProfilerPrintFormat_GRAPHVIZ);
 
 	for (auto &clazz : toFree) {
 		env->DeleteGlobalRef(clazz);
@@ -1269,3 +1303,34 @@ void _duckdb_jdbc_create_extension_type(JNIEnv *env, jclass, jobject conn_buf) {
 	byte_test_type_type.SetAlias("byte_test_type");
 	ExtensionUtil::RegisterType(db_instance, "byte_test_type", byte_test_type_type);
 }
+
+static ProfilerPrintFormat GetProfilerPrintFormat(JNIEnv *env, jobject format) {
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_QUERY_TREE)) {
+		return ProfilerPrintFormat::QUERY_TREE;
+	}
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_JSON)) {
+		return ProfilerPrintFormat::JSON;
+	}
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_QUERY_TREE_OPTIMIZER)) {
+		return ProfilerPrintFormat::QUERY_TREE_OPTIMIZER;
+	}
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_NO_OUTPUT)) {
+		return ProfilerPrintFormat::NO_OUTPUT;
+	}
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_HTML)) {
+		return ProfilerPrintFormat::HTML;
+	}
+	if (env->IsSameObject(format, J_ProfilerPrintFormat_GRAPHVIZ)) {
+		return ProfilerPrintFormat::GRAPHVIZ;
+	}
+}
+
+jstring _duckdb_jdbc_get_profiling_information(JNIEnv *env, jclass, jobject conn_ref_buf, jobject j_format) {
+	auto connection = get_connection(env, conn_ref_buf);
+	if (!connection) {
+		throw InvalidInputException("Invalid connection");
+	}
+	auto format = GetProfilerPrintFormat(env, j_format);
+	auto profiling_info = connection->GetProfilingInformation(format);
+	return env->NewStringUTF(profiling_info.c_str());
+}
diff --git a/src/jni/functions.cpp b/src/jni/functions.cpp
@@ -396,3 +396,13 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extens
 
 	}
 }
+
+JNIEXPORT jstring JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1get_1profiling_1information(JNIEnv * env, jclass param0, jobject param1, jobject param2) {
+	try {
+		return _duckdb_jdbc_get_profiling_information(env, param0, param1, param2);
+	} catch (const std::exception &e) {
+		duckdb::ErrorData error(e);
+		ThrowJNI(env, error.Message().c_str());
+
+	}
+}
diff --git a/src/jni/functions.hpp b/src/jni/functions.hpp
@@ -160,3 +160,7 @@ JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1appender_1appe
 void _duckdb_jdbc_create_extension_type(JNIEnv * env, jclass param0, jobject param1);
 
 JNIEXPORT void JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1create_1extension_1type(JNIEnv * env, jclass param0, jobject param1);
+
+jstring _duckdb_jdbc_get_profiling_information(JNIEnv * env, jclass param0, jobject param1, jobject param2);
+
+JNIEXPORT jstring JNICALL Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1get_1profiling_1information(JNIEnv * env, jclass param0, jobject param1, jobject param2);
diff --git a/src/main/java/org/duckdb/DuckDBConnection.java b/src/main/java/org/duckdb/DuckDBConnection.java
@@ -371,4 +371,8 @@ public void registerArrowStream(String name, Object arrow_array_stream) {
         long array_stream_address = getArrowStreamAddress(arrow_array_stream);
         DuckDBNative.duckdb_jdbc_arrow_register(conn_ref, array_stream_address, name.getBytes(StandardCharsets.UTF_8));
     }
+
+    public String getProfilingInformation(ProfilerPrintFormat format) throws SQLException {
+        return DuckDBNative.duckdb_jdbc_get_profiling_information(conn_ref, format);
+    }
 }
diff --git a/src/main/java/org/duckdb/DuckDBNative.java b/src/main/java/org/duckdb/DuckDBNative.java
@@ -167,6 +167,10 @@ protected static native void duckdb_jdbc_appender_append_decimal(ByteBuffer appe
 
     protected static native void duckdb_jdbc_create_extension_type(ByteBuffer conn_ref) throws SQLException;
 
+    protected static native String duckdb_jdbc_get_profiling_information(ByteBuffer conn_ref,
+                                                                         ProfilerPrintFormat format)
+        throws SQLException;
+
     public static void duckdb_jdbc_create_extension_type(DuckDBConnection conn) throws SQLException {
         duckdb_jdbc_create_extension_type(conn.conn_ref);
     }

diff --git a/src/main/java/org/duckdb/ProfilerPrintFormat.java b/src/main/java/org/duckdb/ProfilerPrintFormat.java
@@ -0,0 +1,11 @@
+package org.duckdb;
+
+public enum ProfilerPrintFormat {
+
+    QUERY_TREE,
+    JSON,
+    QUERY_TREE_OPTIMIZER,
+    NO_OUTPUT,
+    HTML,
+    GRAPHVIZ
+}
diff --git a/src/test/java/org/duckdb/TestDuckDBJDBC.java b/src/test/java/org/duckdb/TestDuckDBJDBC.java
@@ -4676,6 +4676,16 @@ public static void test_blob_after_rs_next() throws Exception {
         }
     }
 
+    public static void test_get_profiling_information() throws Exception {
+        try (Connection conn = DriverManager.getConnection(JDBC_URL); Statement stmt = conn.createStatement()) {
+            stmt.execute("SET enable_profiling = 'no_output';");
+            try (ResultSet rs = stmt.executeQuery("SELECT 1+1")) {
+                String profile = ((DuckDBConnection) conn).getProfilingInformation(ProfilerPrintFormat.JSON);
+                assertTrue(profile.contains("\"query_name\": \"SELECT 1+1\","));
+            }
+        }
+    }
+
     public static void main(String[] args) throws Exception {
         System.exit(runTests(args, TestDuckDBJDBC.class, TestExtensionTypes.class));
     }