diff --git a/include/pgduckdb/pg/explain.hpp b/include/pgduckdb/pg/explain.hpp new file mode 100644 index 00000000..97f73267 --- /dev/null +++ b/include/pgduckdb/pg/explain.hpp @@ -0,0 +1,7 @@ +#pragma once + +#include "pgduckdb/pg/declarations.hpp" + +namespace pgduckdb { +const char *ExplainPGQuery(const char *query_str); +} // namespace pgduckdb diff --git a/include/pgduckdb/pgduckdb_guc.h b/include/pgduckdb/pgduckdb_guc.h index 629de5ac..e9fbc7cb 100644 --- a/include/pgduckdb/pgduckdb_guc.h +++ b/include/pgduckdb/pgduckdb_guc.h @@ -1,6 +1,7 @@ #pragma once extern bool duckdb_force_execution; +extern bool duckdb_log_pg_subplans; extern int duckdb_maximum_threads; extern char *duckdb_maximum_memory; extern char *duckdb_disabled_filesystems; diff --git a/include/pgduckdb/scan/postgres_scan.hpp b/include/pgduckdb/scan/postgres_scan.hpp index 6e36a130..21bb1547 100644 --- a/include/pgduckdb/scan/postgres_scan.hpp +++ b/include/pgduckdb/scan/postgres_scan.hpp @@ -11,6 +11,8 @@ namespace pgduckdb { +bool IsExplaining(); + // Global State struct PostgresScanGlobalState : public duckdb::GlobalTableFunctionState { diff --git a/src/pg/explain.cpp b/src/pg/explain.cpp new file mode 100644 index 00000000..0c6cf27e --- /dev/null +++ b/src/pg/explain.cpp @@ -0,0 +1,33 @@ + +extern "C" { +#include "postgres.h" +#include "lib/stringinfo.h" +#include "executor/spi.h" +} + +#include "pgduckdb/pg/explain.hpp" + +namespace pgduckdb { +const char * +ExplainPGQuery(const char *query_str) { + StringInfo explain_query_si = makeStringInfo(); + // Must be created before SPI_connect to be in the right context + StringInfo explain_result_si = makeStringInfo(); + appendStringInfo(explain_query_si, "EXPLAIN (%s)", query_str); + + SPI_connect(); + int ret = SPI_exec(explain_query_si->data, 0); + if (ret != SPI_OK_UTILITY) { + elog(ERROR, "SPI_exec failed: error code %s", SPI_result_code_string(ret)); + } + + for (uint64_t proc = 0; proc < SPI_processed; ++proc) { + HeapTuple tuple = SPI_tuptable->vals[proc]; + char *row = SPI_getvalue(tuple, SPI_tuptable->tupdesc, 1); + appendStringInfo(explain_result_si, "%s\n", row); + } + + SPI_finish(); + return explain_result_si->data; +} +} // namespace pgduckdb diff --git a/src/pgduckdb.cpp b/src/pgduckdb.cpp index 763c6969..0fb5efae 100644 --- a/src/pgduckdb.cpp +++ b/src/pgduckdb.cpp @@ -15,6 +15,7 @@ extern "C" { static void DuckdbInitGUC(void); bool duckdb_force_execution = false; +bool duckdb_log_pg_subplans = false; int duckdb_max_workers_per_postgres_scan = 2; int duckdb_motherduck_enabled = MotherDuckEnabled::MOTHERDUCK_AUTO; char *duckdb_motherduck_token = strdup(""); @@ -182,4 +183,7 @@ DuckdbInitGUC(void) { DefineCustomVariable("duckdb.motherduck_default_database", "Which database in MotherDuck to designate as default (in place of my_db)", &duckdb_motherduck_default_database, PGC_POSTMASTER, GUC_SUPERUSER_ONLY); + + DefineCustomVariable("duckdb.log_pg_subplans", "Log plan when scanning Postgres relations", + &duckdb_log_pg_subplans); } diff --git a/src/pgduckdb_hooks.cpp b/src/pgduckdb_hooks.cpp index bc3d921f..23c3ac0b 100644 --- a/src/pgduckdb_hooks.cpp +++ b/src/pgduckdb_hooks.cpp @@ -4,6 +4,7 @@ #include "pgduckdb/pg/transactions.hpp" #include "pgduckdb/pgduckdb_xact.hpp" #include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/scan/postgres_scan.hpp" extern "C" { #include "postgres.h" @@ -186,6 +187,12 @@ IsAllowedStatement(Query *query, bool throw_error = false) { static PlannedStmt * DuckdbPlannerHook_Cpp(Query *parse, const char *query_string, int cursor_options, ParamListInfo bound_params) { + if (pgduckdb::IsExplaining()) { + // Don't use DuckDB when explaining a PG query. + auto planner = prev_planner_hook ? prev_planner_hook : standard_planner; + return planner(parse, query_string, cursor_options, bound_params); + } + if (pgduckdb::IsExtensionRegistered()) { if (NeedsDuckdbExecution(parse)) { IsAllowedStatement(parse, true); @@ -214,11 +221,8 @@ DuckdbPlannerHook_Cpp(Query *parse, const char *query_string, int cursor_options pgduckdb::MarkStatementNotTopLevel(); - if (prev_planner_hook) { - return prev_planner_hook(parse, query_string, cursor_options, bound_params); - } else { - return standard_planner(parse, query_string, cursor_options, bound_params); - } + auto planner = prev_planner_hook ? prev_planner_hook : standard_planner; + return planner(parse, query_string, cursor_options, bound_params); } static PlannedStmt * diff --git a/src/scan/postgres_scan.cpp b/src/scan/postgres_scan.cpp index 0e469479..8c3d421b 100644 --- a/src/scan/postgres_scan.cpp +++ b/src/scan/postgres_scan.cpp @@ -1,7 +1,9 @@ #include "pgduckdb/scan/postgres_scan.hpp" #include "pgduckdb/scan/postgres_table_reader.hpp" #include "pgduckdb/pgduckdb_types.hpp" +#include "pgduckdb/pgduckdb_guc.h" #include "pgduckdb/pgduckdb_utils.hpp" +#include "pgduckdb/pg/explain.hpp" #include "pgduckdb/pg/relations.hpp" #include "pgduckdb/pgduckdb_process_lock.hpp" @@ -9,6 +11,24 @@ namespace pgduckdb { +bool explaining = false; + +bool +IsExplaining() { + return explaining; +} + +class ExplainingToggle { +public: + ExplainingToggle() { + explaining = true; + } + + ~ExplainingToggle() { + explaining = false; + } +}; + // // PostgresScanGlobalState // @@ -111,8 +131,19 @@ PostgresScanGlobalState::PostgresScanGlobalState(Snapshot _snapshot, Relation _r : snapshot(_snapshot), rel(_rel), table_tuple_desc(RelationGetDescr(rel)), count_tuples_only(false), total_row_count(0) { ConstructTableScanQuery(input); - table_reader_global_state = - duckdb::make_shared_ptr(scan_query.str().c_str(), count_tuples_only); + + auto scan_query_str = scan_query.str(); + auto scan_query_cstr = scan_query_str.c_str(); + if (duckdb_log_pg_subplans) { + D_ASSERT(!IsExplaining()); + ExplainingToggle toggle; + pd_log(INFO, "Query:\n%s", scan_query_cstr); + PostgresScopedStackReset scoped_stack_reset; + auto output = PostgresFunctionGuard(ExplainPGQuery, scan_query_cstr); + pd_log(INFO, "Plan:\n%s", output); + } + + table_reader_global_state = duckdb::make_shared_ptr(scan_query_cstr, count_tuples_only); pd_log(DEBUG2, "(DuckDB/PostgresSeqScanGlobalState) Running %" PRIu64 " threads -- ", (uint64_t)MaxThreads()); }