Optimizer extensions do not have access to query

[Miosp]

Hello,

after a solid couple of hours spent on debugging, I came to a conclusion that unfortunately behavior I encountered is a bug somewhere in DuckDB (my guess is duckdb-java, but it might as well be somewhere else).

I'm currently writing an extension, in which I wanted to capture user's queries. It worked just fine in DuckDB CLI, but when I loaded it into the JDBC version of DuckDB, it threw some nasty exceptions:

INTERNAL Error: Attempted to dereference unique_ptr that is NULL!

Stack Trace:

/home/miosp/.duckdb/extensions/v1.4.3/linux_amd64/my-extension.duckdb_extension(+0x782414) [0x7f004ccde414]
/home/miosp/.duckdb/extensions/v1.4.3/linux_amd64/my-extension.duckdb_extension(+0x7824bf) [0x7f004ccde4bf]
/home/miosp/.duckdb/extensions/v1.4.3/linux_amd64/my-extension.duckdb_extension(+0x784a71) [0x7f004cce0a71]
/home/miosp/.duckdb/extensions/v1.4.3/linux_amd64/my-extension.duckdb_extension(+0x20083a) [0x7f004c75c83a]
/home/miosp/.duckdb/extensions/v1.4.3/linux_amd64/my-extension.duckdb_extension(+0x47bc7f) [0x7f004c9d7c7f]
/tmp/libduckdb_java9599730312245461022.so(+0x128e924) [0x7f007a33b924]
/tmp/libduckdb_java9599730312245461022.so(+0x128ea5c) [0x7f007a33ba5c]
/tmp/libduckdb_java9599730312245461022.so(+0x1290674) [0x7f007a33d674]
/tmp/libduckdb_java9599730312245461022.so(+0x118488f) [0x7f007a23188f]
/tmp/libduckdb_java9599730312245461022.so(+0x1184ba1) [0x7f007a231ba1]
/tmp/libduckdb_java9599730312245461022.so(+0x1184fac) [0x7f007a231fac]
/tmp/libduckdb_java9599730312245461022.so(+0x1167ffc) [0x7f007a214ffc]
/tmp/libduckdb_java9599730312245461022.so(+0x11681ed) [0x7f007a2151ed]
/tmp/libduckdb_java9599730312245461022.so(+0x118e2c5) [0x7f007a23b2c5]
/tmp/libduckdb_java9599730312245461022.so(+0x118e36b) [0x7f007a23b36b]
/tmp/libduckdb_java9599730312245461022.so(+0x509315) [0x7f00795b6315]
/tmp/libduckdb_java9599730312245461022.so(Java_org_duckdb_DuckDBNative_duckdb_1jdbc_1prepare+0x11) [0x7f00795ba141]
[0x7f02e626eb6d]

I tried to dig into the issue and i developed a minimal PoC of the issue (strapped onto the extension template repository and switched both DuckDB and CI tools versions to 1.4.3 to be sure).

Then, I modified the quack_extension.cpp to:

#define DUCKDB_EXTENSION_MAIN

#include "quack_extension.hpp"
#include "duckdb.hpp"
#include "duckdb/common/exception.hpp"
#include "duckdb/function/scalar_function.hpp"
#include <duckdb/parser/parsed_data/create_scalar_function_info.hpp>

namespace duckdb {

void PreOptimize(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
	if (!plan) {
		return;
	}

	string query;
	try {
		query = input.context.GetCurrentQuery();
	} catch (...) {
		Printer::Print("Caught some Exception");
		return;
	}

	Printer::Print("QuackExtension PreOptimize called for query: " + query);

	return;
}

static void LoadInternal(ExtensionLoader &loader) {
	auto &db_instance = loader.GetDatabaseInstance();
	auto &config = db_instance.config;

	OptimizerExtension extension;
	extension.pre_optimize_function = &PreOptimize;

	config.optimizer_extensions.push_back(extension);
}

void QuackExtension::Load(ExtensionLoader &loader) {
	LoadInternal(loader);
}
std::string QuackExtension::Name() {
	return "quack";
}

std::string QuackExtension::Version() const {
#ifdef EXT_VERSION_QUACK
	return EXT_VERSION_QUACK;
#else
	return "";
#endif
}

} // namespace duckdb

extern "C" {

DUCKDB_CPP_EXTENSION_ENTRY(quack, loader) {
	duckdb::LoadInternal(loader);
}
}

When i launch the extension from the CLI i get:

duckdb -unsigned
DuckDB v1.4.3 (Andium) d1dc88f950
Enter ".help" for usage hints.
D LOAD '/home/miosp/proj/extension-template/build/release/repository/v1.4.3/linux_amd64/quack.duckdb_extension';
D pragma VERSION;
QuackExtension PreOptimize called for query: SELECT * FROM pragma_version();
┌─────────────────┬────────────┬──────────┐
│ library_version │ source_id  │ codename │
│     varchar     │  varchar   │ varchar  │
├─────────────────┼────────────┼──────────┤
│ v1.4.3          │ d1dc88f950 │ Andium   │
└─────────────────┴────────────┴──────────┘

But through the Java JDBC driver I get (after issuing the same commands programatically):

Caught some Exception

So the output is still correct, but the query was not attached to the context.

I also noticed that if I connect DuckLake to my JDBC DuckDB, the auxiliary queries are caught correctly (I believe this one is for SHOW ALL TABLES;):

QuackExtension PreOptimize called for query: BEGIN TRANSACTION
QuackExtension PreOptimize called for query: 
    SELECT * FROM postgres_query('__ducklake_metadata_miosp_ducklake',
        'SELECT snapshot_id, schema_version, next_catalog_id, next_file_id
         FROM "public".ducklake_snapshot WHERE snapshot_id = (
             SELECT MAX(snapshot_id) FROM "public".ducklake_snapshot
         );')
    
QuackExtension PreOptimize called for query: COMMIT
QuackExtension PreOptimize called for query: BEGIN TRANSACTION
QuackExtension PreOptimize called for query: 
    SELECT * FROM postgres_query('__ducklake_metadata_miosp_ducklake',
        'SELECT snapshot_id, schema_version, next_catalog_id, next_file_id
         FROM "public".ducklake_snapshot WHERE snapshot_id = (
             SELECT MAX(snapshot_id) FROM "public".ducklake_snapshot
         );')
    
Caught some Exception
QuackExtension PreOptimize called for query: COMMIT
Caught some Exception
Caught some Exception
QuackExtension PreOptimize called for query: BEGIN TRANSACTION
QuackExtension PreOptimize called for query: 
    SELECT * FROM postgres_query('__ducklake_metadata_miosp_ducklake',
        'SELECT snapshot_id, schema_version, next_catalog_id, next_file_id
         FROM "public".ducklake_snapshot WHERE snapshot_id = (
             SELECT MAX(snapshot_id) FROM "public".ducklake_snapshot
         );')
    
QuackExtension PreOptimize called for query: COMMIT
Caught some Exception

Am I doing something wrong, or this is a legitimate bug?

[staticlibs]

Hi, thanks for the report!

"Caught some Exception"

What is this exception that is being caught here? Is it the same every time?

[Miosp]

Take a look at the code, I just hardcoded-in the string. In general the active_query from the context seems to be nullptr, so the input.context.GetCurrentQuery(); fails. This always seems to be the cause of my original unique_ptr dereferencing exception.

[thijs-s]

we had similar problem, here is a potential fix duckdb/duckdb#20755