expose profiling in the pyconnection

Author: guillesd

.gitignore

@@ -14,6 +14,8 @@
 .sw?
 #OS X specific files.
 .DS_store
+#VSCode specifics
+.vscode/
 
 #==============================================================================#
 # Build artifacts

_duckdb-stubs/__init__.pyi

@@ -86,9 +86,11 @@ __all__: list[str] = [
     "default_connection",
     "description",
     "df",
+    "disable_profiling",
     "distinct",
     "dtype",
     "duplicate",
+    "enable_profiling",
     "enum_type",
     "execute",
     "executemany",
@@ -109,6 +111,7 @@ __all__: list[str] = [
     "from_df",
     "from_parquet",
     "from_query",
+    "get_profiling_information",
     "get_table_names",
     "install_extension",
     "interrupt",
@@ -313,6 +316,9 @@ class DuckDBPyConnection:
         repository_url: str | None = None,
         version: str | None = None,
     ) -> None: ...
+    def get_profiling_information(self, format: str = "json") -> str: ...
+    def enable_profiling(self) -> None: ...
+    def disable_profiling(self) -> None: ...
     def interrupt(self) -> None: ...
     def list_filesystems(self) -> list[str]: ...
     def list_type(self, type: sqltypes.DuckDBPyType) -> sqltypes.DuckDBPyType: ...
@@ -1227,6 +1233,9 @@ def limit(
     *,
     connection: DuckDBPyConnection | None = None,
 ) -> DuckDBPyRelation: ...
+def get_profiling_information(*, connection: DuckDBPyConnection | None = None, format: str = "json") -> str: ...
+def enable_profiling(*, connection: DuckDBPyConnection | None = None) -> None: ...
+def disable_profiling(*, connection: DuckDBPyConnection | None = None) -> None: ...
 def list_filesystems(*, connection: DuckDBPyConnection | None = None) -> list[str]: ...
 def list_type(
     type: sqltypes.DuckDBPyType, *, connection: DuckDBPyConnection | None = None

duckdb/__init__.py

@@ -84,9 +84,11 @@
     default_connection,
     description,
     df,
+    disable_profiling,
     distinct,
     dtype,
     duplicate,
+    enable_profiling,
     enum_type,
     execute,
     executemany,
@@ -107,6 +109,7 @@
     from_df,
     from_parquet,
     from_query,
+    get_profiling_information,
     get_table_names,
     install_extension,
     interrupt,
@@ -310,9 +313,11 @@
     "default_connection",
     "description",
     "df",
+    "disable_profiling",
     "distinct",
     "dtype",
     "duplicate",
+    "enable_profiling",
     "enum_type",
     "execute",
     "executemany",
@@ -333,6 +338,7 @@
     "from_df",
     "from_parquet",
     "from_query",
+    "get_profiling_information",
     "get_table_names",
     "install_extension",
     "interrupt",

scripts/connection_methods.json

@@ -1093,5 +1093,30 @@
 			}
 		],
 		"return": "None"
+	},
+	{
+		"name": "get_profiling_information",
+		"function": "GetProfilingInformation",
+		"docs": "Get profiling information for a query",
+		"args": [
+			{
+				"name": "format",
+				"default": "JSON",
+				"type": "Optional[str]"
+			}
+		],
+		"return": "str"
+	},
+	{
+		"name": "enable_profiling",
+		"function": "EnableProfiling",
+		"docs": "Enable profiling for a connection",
+		"return": "None"
+	},
+	{
+		"name": "disable_profiling",
+		"function": "DisableProfiling",
+		"docs": "Disable profiling for a connection",
+		"return": "None"
 	}
 ]

scripts/generate_connection_stubs.py

@@ -5,7 +5,7 @@
 os.chdir(Path(__file__).parent)
 
 JSON_PATH = "connection_methods.json"
-DUCKDB_STUBS_FILE = Path("..") / "duckdb" / "__init__.pyi"
+DUCKDB_STUBS_FILE = Path("..") / "_duckdb-stubs" / "__init__.pyi"
 
 START_MARKER = "    # START OF CONNECTION METHODS"
 END_MARKER = "    # END OF CONNECTION METHODS"

src/duckdb_py/duckdb_python.cpp

@@ -124,6 +124,34 @@ static void InitializeConnectionMethods(py::module_ &m) {
 	    },
 	    "Check if a filesystem with the provided name is currently registered", py::arg("name"), py::kw_only(),
 	    py::arg("connection") = py::none());
+	m.def(
+	    "get_profiling_information",
+	    [](const py::str &format, shared_ptr<DuckDBPyConnection> conn = nullptr) {
+		    if (!conn) {
+			    conn = DuckDBPyConnection::DefaultConnection();
+		    }
+		    return conn->GetProfilingInformation(format);
+	    },
+	    "Get profiling information from a query", py::kw_only(), py::arg("format") = "json", 
+		py::arg("connection") = py::none());
+	m.def(
+	    "enable_profiling",
+	    [](shared_ptr<DuckDBPyConnection> conn = nullptr) {
+		    if (!conn) {
+			    conn = DuckDBPyConnection::DefaultConnection();
+		    }
+		    return conn->EnableProfiling();
+	    },
+	    "Enable profiling for the current connection", py::kw_only(), py::arg("connection") = py::none());
+	m.def(
+	    "disable_profiling",
+	    [](shared_ptr<DuckDBPyConnection> conn = nullptr) {
+		    if (!conn) {
+			    conn = DuckDBPyConnection::DefaultConnection();
+		    }
+		    return conn->DisableProfiling();
+	    },
+	    "Disable profiling for the current connection", py::kw_only(), py::arg("connection") = py::none());
 	m.def(
 	    "create_function",
 	    [](const string &name, const py::function &udf, const py::object &arguments = py::none(),

src/duckdb_py/include/duckdb_python/pyconnection/pyconnection.hpp

@@ -337,6 +337,11 @@ struct DuckDBPyConnection : public enable_shared_from_this<DuckDBPyConnection> {
 	py::list ListFilesystems();
 	bool FileSystemIsRegistered(const string &name);
 
+	// Profiling info
+	py::str GetProfilingInformation(const py::str &format = "json");
+	void EnableProfiling();
+	void DisableProfiling();
+
 	//! Default connection to an in-memory database
 	static DefaultConnectionHolder default_connection;
 	//! Caches and provides an interface to get frequently used modules+subtypes

src/duckdb_py/pyconnection.cpp

@@ -3,6 +3,7 @@
 #include "duckdb/catalog/default/default_types.hpp"
 #include "duckdb/common/arrow/arrow.hpp"
 #include "duckdb/common/enums/file_compression_type.hpp"
+#include "duckdb/common/enums/profiler_format.hpp"
 #include "duckdb/common/printer.hpp"
 #include "duckdb/common/types.hpp"
 #include "duckdb/common/types/vector.hpp"
@@ -285,6 +286,9 @@ static void InitializeConnectionMethods(py::class_<DuckDBPyConnection, shared_pt
 	      py::arg("extension"), py::kw_only(), py::arg("force_install") = false, py::arg("repository") = py::none(),
 	      py::arg("repository_url") = py::none(), py::arg("version") = py::none());
 	m.def("load_extension", &DuckDBPyConnection::LoadExtension, "Load an installed extension", py::arg("extension"));
+	m.def("get_profiling_information", &DuckDBPyConnection::GetProfilingInformation, "Get profiling information for a query", py::arg("format") = "json");
+	m.def("enable_profiling", &DuckDBPyConnection::EnableProfiling, "Enable profiling for subsequent queries");
+	m.def("disable_profiling", &DuckDBPyConnection::DisableProfiling, "Disable profiling for subsequent queries");
 } // END_OF_CONNECTION_METHODS
 
 void DuckDBPyConnection::UnregisterFilesystem(const py::str &name) {
@@ -331,6 +335,39 @@ py::list DuckDBPyConnection::ListFilesystems() {
 	return names;
 }
 
+py::str DuckDBPyConnection::GetProfilingInformation(const py::str &format) {
+	// We want to expose ProfilerPrintFormat as a string to Python users
+	ProfilerPrintFormat format_enum;
+	if (format == "query_tree") {
+		format_enum = ProfilerPrintFormat::QUERY_TREE;
+	} else if (format == "json") {
+		format_enum = ProfilerPrintFormat::JSON;
+	} else if (format == "query_tree_optimizer") {
+		format_enum = ProfilerPrintFormat::QUERY_TREE_OPTIMIZER;
+	} else if (format == "no_output") {
+		format_enum = ProfilerPrintFormat::NO_OUTPUT;
+	} else if (format == "html") {
+		format_enum = ProfilerPrintFormat::HTML;
+	} else if (format == "graphviz") {
+		format_enum = ProfilerPrintFormat::GRAPHVIZ;
+	} else {
+		throw InvalidInputException("Invalid ProfilerPrintFormat string: " + std::string(format) + ". Valid options are: query_tree, json, query_tree_optimizer, no_output, html, graphviz.");
+	}
+	auto &connection = con.GetConnection();
+	py::str profiling_info = connection.GetProfilingInformation(format_enum);
+	return profiling_info;
+}
+
+void DuckDBPyConnection::EnableProfiling() {
+	auto &connection = con.GetConnection();
+	connection.EnableProfiling();
+}
+
+void DuckDBPyConnection::DisableProfiling() {
+	auto &connection = con.GetConnection();
+	connection.DisableProfiling();
+}
+
 py::list DuckDBPyConnection::ExtractStatements(const string &query) {
 	py::list result;
 	auto &connection = con.GetConnection();