diff --git a/docs/postgres.md b/docs/postgres.md index 3374065..a7194cd 100644 --- a/docs/postgres.md +++ b/docs/postgres.md @@ -149,3 +149,77 @@ if (!result) { } ``` +## Notice Processor + +PostgreSQL functions can emit NOTICE messages using `RAISE NOTICE` in PL/pgSQL. By default, libpq prints these to stderr. sqlgen allows you to capture these messages by providing a custom notice handler in the connection credentials. + +### Setting Up a Notice Handler + +```cpp +#include +#include + +// Create credentials with a notice handler +const auto creds = sqlgen::postgres::Credentials{ + .user = "myuser", + .password = "mypassword", + .host = "localhost", + .dbname = "mydatabase", + .notice_handler = [](const char* msg) { + std::cout << "[PostgreSQL Notice] " << msg; + } +}; + +auto conn = sqlgen::postgres::connect(creds); +``` + +### Integration with Logging Frameworks + +The notice handler can forward messages to your preferred logging framework: + +```cpp +#include + +const auto creds = sqlgen::postgres::Credentials{ + .user = "myuser", + .password = "mypassword", + .host = "localhost", + .dbname = "mydatabase", + .notice_handler = [](const char* msg) { + // Remove trailing newline if present + std::string_view sv(msg); + if (!sv.empty() && sv.back() == '\n') { + sv.remove_suffix(1); + } + spdlog::info("PostgreSQL: {}", sv); + } +}; +``` + +### Using with Connection Pools + +The notice handler is set per-connection, so all connections in a pool will use the same handler: + +```cpp +const auto creds = sqlgen::postgres::Credentials{ + .user = "myuser", + .password = "mypassword", + .host = "localhost", + .dbname = "mydatabase", + .notice_handler = [](const char* msg) { + my_logger::log(msg); + } +}; + +auto pool = sqlgen::make_connection_pool( + sqlgen::ConnectionPoolConfig{.size = 4}, + creds +); +``` + +### Notes + +- If no `notice_handler` is provided, libpq's default behavior (printing to stderr) is used +- The handler receives the full message including any trailing newline +- The handler should be thread-safe when used with connection pools, as multiple connections may invoke it concurrently + diff --git a/include/sqlgen/postgres/Credentials.hpp b/include/sqlgen/postgres/Credentials.hpp index 318041c..cf17bc8 100644 --- a/include/sqlgen/postgres/Credentials.hpp +++ b/include/sqlgen/postgres/Credentials.hpp @@ -1,6 +1,7 @@ #ifndef SQLGEN_POSTGRES_CREDENTIALS_HPP_ #define SQLGEN_POSTGRES_CREDENTIALS_HPP_ +#include #include namespace sqlgen::postgres { @@ -11,6 +12,7 @@ struct Credentials { std::string host; std::string dbname; int port = 5432; + std::function notice_handler; std::string to_str() const { return "postgresql://" + user + ":" + password + "@" + host + ":" + diff --git a/include/sqlgen/postgres/PostgresV2Connection.hpp b/include/sqlgen/postgres/PostgresV2Connection.hpp index 016a1c9..5c37514 100644 --- a/include/sqlgen/postgres/PostgresV2Connection.hpp +++ b/include/sqlgen/postgres/PostgresV2Connection.hpp @@ -3,6 +3,7 @@ #include +#include #include #include #include @@ -15,15 +16,22 @@ namespace sqlgen::postgres { class SQLGEN_API PostgresV2Connection { public: + using NoticeHandler = std::function; + PostgresV2Connection(PGconn* _ptr) : ptr_(Ref::make(std::shared_ptr(_ptr, &PQfinish)) .value()) {} + PostgresV2Connection(PGconn* _ptr, NoticeHandler _notice_handler); + ~PostgresV2Connection() = default; static rfl::Result make( const std::string& _conn_str) noexcept; + static rfl::Result make( + const std::string& _conn_str, NoticeHandler _notice_handler) noexcept; + static rfl::Result make(PGconn* _ptr) noexcept { try { return PostgresV2Connection(_ptr); @@ -37,6 +45,7 @@ class SQLGEN_API PostgresV2Connection { private: Ref ptr_; + std::shared_ptr notice_handler_; }; } // namespace sqlgen::postgres diff --git a/src/sqlgen/postgres/Connection.cpp b/src/sqlgen/postgres/Connection.cpp index ee74c5b..1494536 100644 --- a/src/sqlgen/postgres/Connection.cpp +++ b/src/sqlgen/postgres/Connection.cpp @@ -16,7 +16,9 @@ namespace sqlgen::postgres { Connection::Connection(const Conn& _conn) : conn_(_conn) {} Connection::Connection(const Credentials& _credentials) - : conn_(PostgresV2Connection::make(_credentials.to_str()).value()) {} + : conn_(PostgresV2Connection::make(_credentials.to_str(), + _credentials.notice_handler) + .value()) {} Result Connection::begin_transaction() noexcept { return execute("BEGIN TRANSACTION;"); @@ -182,7 +184,8 @@ Result Connection::insert_impl( rfl::Result> Connection::make( const Credentials& _credentials) noexcept { - return PostgresV2Connection::make(_credentials.to_str()) + return PostgresV2Connection::make(_credentials.to_str(), + _credentials.notice_handler) .transform([](auto&& _conn) { return Ref::make(_conn); }); } diff --git a/src/sqlgen/postgres/PostgresV2Connection.cpp b/src/sqlgen/postgres/PostgresV2Connection.cpp index 038ed72..48f5bfe 100644 --- a/src/sqlgen/postgres/PostgresV2Connection.cpp +++ b/src/sqlgen/postgres/PostgresV2Connection.cpp @@ -2,6 +2,28 @@ namespace sqlgen::postgres { +namespace { +void notice_trampoline(void* arg, const char* message) { + auto* handler = + static_cast(arg); + if (handler && *handler) { + (*handler)(message); + } +} +} // namespace + +PostgresV2Connection::PostgresV2Connection(PGconn* _ptr, + NoticeHandler _notice_handler) + : ptr_(Ref::make(std::shared_ptr(_ptr, &PQfinish)).value()), + notice_handler_(_notice_handler + ? std::make_shared( + std::move(_notice_handler)) + : nullptr) { + if (notice_handler_) { + PQsetNoticeProcessor(ptr_.get(), notice_trampoline, notice_handler_.get()); + } +} + rfl::Result PostgresV2Connection::make( const std::string& _conn_str) noexcept { auto conn = PQconnectdb(_conn_str.c_str()); @@ -14,4 +36,16 @@ rfl::Result PostgresV2Connection::make( return PostgresV2Connection(conn); } +rfl::Result PostgresV2Connection::make( + const std::string& _conn_str, NoticeHandler _notice_handler) noexcept { + auto conn = PQconnectdb(_conn_str.c_str()); + if (PQstatus(conn) != CONNECTION_OK) { + const auto msg = + std::string("Connection to postgres failed: ") + PQerrorMessage(conn); + PQfinish(conn); + return error(msg); + } + return PostgresV2Connection(conn, std::move(_notice_handler)); +} + } // namespace sqlgen::postgres diff --git a/tests/postgres/test_notice_processor.cpp b/tests/postgres/test_notice_processor.cpp new file mode 100644 index 0000000..a37f018 --- /dev/null +++ b/tests/postgres/test_notice_processor.cpp @@ -0,0 +1,109 @@ +#ifndef SQLGEN_BUILD_DRY_TESTS_ONLY + +#include + +#include +#include +#include +#include + +namespace test_notice_processor { + +TEST(postgres, notice_processor_captures_raise_notice) { + std::vector captured_notices; + std::mutex mtx; + + const auto credentials = sqlgen::postgres::Credentials{ + .user = "postgres", + .password = "password", + .host = "localhost", + .dbname = "postgres", + .notice_handler = [&](const char* msg) { + std::lock_guard lock(mtx); + captured_notices.push_back(msg); + }}; + + auto conn_result = sqlgen::postgres::connect(credentials); + ASSERT_TRUE(conn_result); + auto conn = conn_result.value(); + + // Execute a DO block that raises a notice + auto result = conn->execute(R"( + DO $$ + BEGIN + RAISE NOTICE 'Hello from PL/pgSQL'; + END + $$; + )"); + ASSERT_TRUE(result) << result.error().what(); + + // Verify notice was captured + std::lock_guard lock(mtx); + ASSERT_EQ(captured_notices.size(), 1); + EXPECT_TRUE(captured_notices[0].find("Hello from PL/pgSQL") != + std::string::npos); +} + +TEST(postgres, no_notice_handler_works) { + // Default behavior - no handler, should not crash + const auto credentials = sqlgen::postgres::Credentials{ + .user = "postgres", + .password = "password", + .host = "localhost", + .dbname = "postgres"}; + + auto conn_result = sqlgen::postgres::connect(credentials); + ASSERT_TRUE(conn_result); + auto conn = conn_result.value(); + + auto result = conn->execute(R"( + DO $$ + BEGIN + RAISE NOTICE 'This goes to stderr by default'; + END + $$; + )"); + ASSERT_TRUE(result); +} + +TEST(postgres, notice_processor_multiple_notices) { + std::vector captured_notices; + std::mutex mtx; + + const auto credentials = sqlgen::postgres::Credentials{ + .user = "postgres", + .password = "password", + .host = "localhost", + .dbname = "postgres", + .notice_handler = [&](const char* msg) { + std::lock_guard lock(mtx); + captured_notices.push_back(msg); + }}; + + auto conn_result = sqlgen::postgres::connect(credentials); + ASSERT_TRUE(conn_result); + auto conn = conn_result.value(); + + // Execute a DO block that raises multiple notices + auto result = conn->execute(R"( + DO $$ + BEGIN + RAISE NOTICE 'First notice'; + RAISE NOTICE 'Second notice'; + RAISE NOTICE 'Third notice'; + END + $$; + )"); + ASSERT_TRUE(result) << result.error().what(); + + // Verify all notices were captured + std::lock_guard lock(mtx); + ASSERT_EQ(captured_notices.size(), 3); + EXPECT_TRUE(captured_notices[0].find("First notice") != std::string::npos); + EXPECT_TRUE(captured_notices[1].find("Second notice") != std::string::npos); + EXPECT_TRUE(captured_notices[2].find("Third notice") != std::string::npos); +} + +} // namespace test_notice_processor + +#endif