Subclass BufferingLogStorage

Author: evertlammerts

src/duckdb_py/include/duckdb_python/python_log_storage.hpp

@@ -9,29 +9,60 @@
 #pragma once
 
 #include "duckdb/logging/log_storage.hpp"
-#include "duckdb/logging/logging.hpp"
+#include "duckdb/common/map.hpp"
+#include "duckdb/common/unique_ptr.hpp"
 
 namespace duckdb {
 
-class PythonLogStorage : public LogStorage {
+class ColumnDataCollection;
+class DatabaseInstance;
+
+//! Scan state backing PythonLogStorage's in-memory buffers (so `duckdb_logs` can read them).
+//! We define our own rather than reuse the engine's InMemoryLogStorageScanState to avoid
+//! depending on whether that type's symbols are exported across platforms.
+class PythonLogStorageScanState : public LogStorageScanState {
+public:
+	explicit PythonLogStorageScanState(LoggingTargetTable table) : LogStorageScanState(table) {
+	}
+	~PythonLogStorageScanState() override = default;
+
+	ColumnDataScanState scan_state;
+};
+
+//! A composite log storage that does two things for every engine log entry:
+//!   1. forwards it to Python's standard `logging` module (logging.getLogger("duckdb")), and
+//!   2. retains it in-memory so `SELECT * FROM duckdb_logs` keeps working.
+//!
+//! It subclasses BufferingLogStorage with a buffer size of 1 so each entry is flushed (and
+//! therefore forwarded to Python) immediately, rather than batched until a 2048-entry buffer
+//! fills — engine WARNINGs are sparse and must surface inline to be useful.
+class PythonLogStorage : public BufferingLogStorage {
 public:
-	PythonLogStorage() = default;
-	~PythonLogStorage() override = default;
+	explicit PythonLogStorage(DatabaseInstance &db);
+	~PythonLogStorage() override;
 
 	const string GetStorageName() override {
 		return "python_log_storage";
 	}
 
-	void WriteLogEntry(timestamp_t timestamp, LogLevel level, const string &log_type, const string &log_message,
-	                   const RegisteredLoggingContext &context) override;
-	void WriteLogEntries(DataChunk &chunk, const RegisteredLoggingContext &context) override;
-	void FlushAll() override {
-	}
-	void Flush(LoggingTargetTable table) override {
-	}
-	bool IsEnabled(LoggingTargetTable table) override {
-		return true;
-	}
+	//! Single-threaded scan interface — mirrors InMemoryLogStorage so duckdb_logs can read us.
+	bool CanScan(LoggingTargetTable table) override;
+	unique_ptr<LogStorageScanState> CreateScanState(LoggingTargetTable table) const override;
+	bool Scan(LogStorageScanState &state, DataChunk &result) const override;
+	void InitializeScan(LogStorageScanState &state) const override;
+
+protected:
+	//! Stores the chunk for duckdb_logs and (for LOG_ENTRIES) forwards it to Python.
+	void FlushChunk(LoggingTargetTable table, DataChunk &chunk) override;
+	//! Clears the in-memory buffers.
+	void ResetAllBuffers() override;
+
+private:
+	ColumnDataCollection &GetBuffer(LoggingTargetTable table) const;
+	//! Forwards each row of a LOG_ENTRIES chunk to logging.getLogger("duckdb"). Never throws.
+	void ForwardEntriesToPython(DataChunk &chunk);
+
+	map<LoggingTargetTable, unique_ptr<ColumnDataCollection>> log_storage_buffers;
 };
 
 } // namespace duckdb

src/duckdb_py/pyconnection.cpp

@@ -2288,7 +2288,7 @@ shared_ptr<DuckDBPyConnection> DuckDBPyConnection::Connect(const py::object &dat
 	{
 		auto &db_instance = *res->con.GetDatabase().instance;
 		auto &log_manager = db_instance.GetLogManager();
-		auto storage = make_shared_ptr<PythonLogStorage>();
+		auto storage = make_shared_ptr<PythonLogStorage>(db_instance);
 		shared_ptr<LogStorage> storage_base = storage;
 		// RegisterLogStorage returns false if the name is already registered on this
 		// DatabaseInstance. Instances are cached and shared across connections/cursors, so

src/duckdb_py/python_log_storage.cpp

@@ -1,92 +1,120 @@
 #include "duckdb_python/python_log_storage.hpp"
 #include "duckdb_python/pybind11/pybind_wrapper.hpp"
+
+#include "duckdb/common/allocator.hpp"
+#include "duckdb/common/exception.hpp"
+#include "duckdb/common/types/column/column_data_collection.hpp"
 #include "duckdb/common/types/data_chunk.hpp"
 #include "duckdb/common/types/vector.hpp"
-#include "duckdb/logging/logging.hpp"
 
 namespace duckdb {
 
-static int LogLevelToPython(LogLevel level) {
-	switch (level) {
-	case LogLevel::LOG_TRACE:
-	case LogLevel::LOG_DEBUG:
-		return 10; // logging.DEBUG
-	case LogLevel::LOG_INFO:
-		return 20; // logging.INFO
-	case LogLevel::LOG_WARNING:
-		return 30; // logging.WARNING
-	case LogLevel::LOG_ERROR:
-		return 40; // logging.ERROR
-	case LogLevel::LOG_FATAL:
-		return 50; // logging.CRITICAL
-	default:
-		return 30;
-	}
-}
-
+// Maps the engine's textual log level (stored as VARCHAR in the LOG_ENTRIES chunk) to the
+// integer levels of Python's logging module.
 static int LevelStringToPython(const string &level_str) {
 	if (level_str == "TRACE" || level_str == "DEBUG") {
-		return 10;
+		return 10; // logging.DEBUG
 	}
 	if (level_str == "INFO") {
-		return 20;
+		return 20; // logging.INFO
 	}
 	if (level_str == "WARNING") {
-		return 30;
+		return 30; // logging.WARNING
 	}
 	if (level_str == "ERROR") {
-		return 40;
+		return 40; // logging.ERROR
 	}
 	if (level_str == "FATAL") {
-		return 50;
+		return 50; // logging.CRITICAL
 	}
 	return 30;
 }
 
-// Both write methods run on engine worker threads and invoke arbitrary user Python (the
-// handlers installed on the "duckdb" logger). The engine calls these directly from query
-// binding/execution with NO surrounding try/catch (see LogManager::WriteLogEntry), so an
-// exception escaping here would fail the user's query. Logging is a side effect — it must
-// never do that. Hence every body swallows all exceptions.
-//
-// Note also that the engine holds LogManager::lock (a non-recursive mutex) across this call.
-// A handler that re-enters DuckDB on the same thread and emits another log entry would
-// self-deadlock on that lock — outside our control, but worth knowing.
+PythonLogStorage::PythonLogStorage(DatabaseInstance &db) : BufferingLogStorage(db, 1, true) {
+	log_storage_buffers[LoggingTargetTable::LOG_ENTRIES] =
+	    make_uniq<ColumnDataCollection>(Allocator::DefaultAllocator(), GetSchema(LoggingTargetTable::LOG_ENTRIES));
+	log_storage_buffers[LoggingTargetTable::LOG_CONTEXTS] =
+	    make_uniq<ColumnDataCollection>(Allocator::DefaultAllocator(), GetSchema(LoggingTargetTable::LOG_CONTEXTS));
+}
 
-void PythonLogStorage::WriteLogEntry(timestamp_t, LogLevel level, const string &, const string &log_message,
-                                     const RegisteredLoggingContext &) {
-	if (!Py_IsInitialized()) {
-		return; // interpreter is finalizing — acquiring the GIL would crash
-	}
-	try {
-		py::gil_scoped_acquire gil;
-		auto logging = py::module::import("logging");
-		auto logger = logging.attr("getLogger")("duckdb");
-		logger.attr("log")(LogLevelToPython(level), log_message);
-	} catch (...) {
-		// Logging must not disrupt query execution.
+PythonLogStorage::~PythonLogStorage() {
+}
+
+ColumnDataCollection &PythonLogStorage::GetBuffer(LoggingTargetTable table) const {
+	auto res = log_storage_buffers.find(table);
+	if (res == log_storage_buffers.end()) {
+		throw InternalException("PythonLogStorage: failed to find buffer for logging target table");
 	}
+	return *res->second;
 }
 
-void PythonLogStorage::WriteLogEntries(DataChunk &chunk, const RegisteredLoggingContext &) {
+void PythonLogStorage::ForwardEntriesToPython(DataChunk &chunk) {
+	// This fires from engine worker threads with the GIL released, and from within both the
+	// LogManager lock and this storage's lock. It runs arbitrary user Python (logging
+	// handlers) and MUST NOT let an exception escape: the engine calls the write path with no
+	// try/catch, directly from query binding/execution, so a raising handler would otherwise
+	// fail the user's query. Hence we swallow everything here.
+	//
+	// Caveat: because a lock is held across this call, a handler that re-enters DuckDB on the
+	// same thread and emits another log entry can self-deadlock on the non-recursive lock.
+	// That is outside our control (and matches the engine's own contract for log storages).
 	if (!Py_IsInitialized()) {
 		return; // interpreter is finalizing — acquiring the GIL would crash
 	}
 	try {
 		py::gil_scoped_acquire gil;
 		auto logging = py::module::import("logging");
 		auto logger = logging.attr("getLogger")("duckdb");
-		// DataChunk is in LOG_ENTRIES format: context_id, timestamp, type, log_level, message.
-		// log_level (idx 3) and message (idx 4) are both VARCHAR; the chunk is freshly
-		// allocated by the engine so the vectors are flat.
+		// LOG_ENTRIES schema: context_id, timestamp, type, log_level (idx 3), message (idx 4).
+		// log_level and message are both VARCHAR; the buffer chunk is flat.
 		auto level_data = FlatVector::GetData<string_t>(chunk.data[3]);
 		auto message_data = FlatVector::GetData<string_t>(chunk.data[4]);
 		for (idx_t i = 0; i < chunk.size(); i++) {
 			logger.attr("log")(LevelStringToPython(level_data[i].GetString()), message_data[i].GetString());
 		}
 	} catch (...) {
-		// Logging must not disrupt query execution.
+		// Logging must never disrupt query execution.
 	}
 }
 
+void PythonLogStorage::FlushChunk(LoggingTargetTable table, DataChunk &chunk) {
+	D_ASSERT(table == LoggingTargetTable::LOG_ENTRIES || table == LoggingTargetTable::LOG_CONTEXTS);
+	// Retain the entry for duckdb_logs FIRST, so a misbehaving Python handler can never cost
+	// us a stored row.
+	log_storage_buffers[table]->Append(chunk);
+	// Forward only real log entries (not context metadata) to Python's logging module.
+	if (table == LoggingTargetTable::LOG_ENTRIES) {
+		ForwardEntriesToPython(chunk);
+	}
+}
+
+void PythonLogStorage::ResetAllBuffers() {
+	BufferingLogStorage::ResetAllBuffers();
+	for (const auto &buffer : log_storage_buffers) {
+		buffer.second->Reset();
+	}
+}
+
+bool PythonLogStorage::CanScan(LoggingTargetTable table) {
+	unique_lock<mutex> lck(lock);
+	return IsEnabledInternal(table);
+}
+
+unique_ptr<LogStorageScanState> PythonLogStorage::CreateScanState(LoggingTargetTable table) const {
+	return make_uniq<PythonLogStorageScanState>(table);
+}
+
+bool PythonLogStorage::Scan(LogStorageScanState &state, DataChunk &result) const {
+	unique_lock<mutex> lck(lock);
+	auto &python_scan_state = state.Cast<PythonLogStorageScanState>();
+	return GetBuffer(python_scan_state.table).Scan(python_scan_state.scan_state, result);
+}
+
+void PythonLogStorage::InitializeScan(LogStorageScanState &state) const {
+	unique_lock<mutex> lck(lock);
+	auto &python_scan_state = state.Cast<PythonLogStorageScanState>();
+	GetBuffer(python_scan_state.table)
+	    .InitializeScan(python_scan_state.scan_state, ColumnDataScanProperties::DISALLOW_ZERO_COPY);
+}
+
 } // namespace duckdb