Refactor PostgreSQL connections #153
Conversation
|
""" WalkthroughThe changes introduce persistent PostgreSQL connection management to replace per-operation connection creation. A new connection utility class is added, and shared storage logic is refactored to reuse connections for read and write actions, with explicit lifecycle management. Conditional connection closure is implemented in the bot base class, and comprehensive tests are provided for the new connection logic and its integration. Changes
Sequence Diagram(s)sequenceDiagram
participant Bot as Bas::Bot::Base
participant Reader as SharedStorageReader
participant Writer as SharedStorageWriter
participant ConnUtil as Utils::Postgres::Connection
Bot->>Reader: read(...)
activate Reader
Reader->>ConnUtil: establish_connection(:read)
ConnUtil-->>Reader: returns connection
Reader->>ConnUtil: query(...)
ConnUtil-->>Reader: returns results
deactivate Reader
Bot->>Writer: write(...)
activate Writer
Writer->>ConnUtil: establish_connection(:write)
ConnUtil-->>Writer: returns connection
Writer->>ConnUtil: query(...)
ConnUtil-->>Writer: returns results
deactivate Writer
alt Option: close_connections_after_process == true
Bot->>Reader: close_connections
Reader->>ConnUtil: finish
Bot->>Writer: close_connections
Writer->>ConnUtil: finish
end
Estimated code review effort3 (~45 minutes) Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (5)
lib/bas/utils/postgres/connection.rb (1)
15-25: Consider adding connection health checks and error handlingThe query method should handle potential connection errors and provide meaningful error messages. Also, consider checking connection health before executing queries.
def query(query) + raise "Connection is closed" if @connection.nil? + results = if query.is_a? String @connection.exec(query) else sentence, params = query @connection.exec_params(sentence, params) end results.map { |result| result.transform_keys(&:to_sym) } +rescue PG::Error => e + raise "Query execution failed: #{e.message}" endspec/bas/bot/base_spec.rb (1)
77-92: Consider extracting common test setup to reduce duplicationBoth connection closing tests have similar setup code. Consider using a shared context or helper method to reduce duplication.
+shared_context "bot execution setup" do + before do + allow(@shared_storage_reader).to receive(:read).and_return(read_response) + allow(@shared_storage_writer).to receive(:write).and_return({}) + allow(@shared_storage_reader).to receive(:close_connections).and_return(true) + allow(@shared_storage_writer).to receive(:close_connections).and_return(true) + allow_any_instance_of(described_class).to receive(:process).and_return({ success: "ok" }) + end +end + it "closes the connections when close_connections_after_process is true" do - allow(@shared_storage_reader).to receive(:read).and_return(read_response) - allow(@shared_storage_writer).to receive(:write).and_return({}) - allow(@shared_storage_reader).to receive(:close_connections).and_return(true) - allow(@shared_storage_writer).to receive(:close_connections).and_return(true) - allow_any_instance_of(described_class).to receive(:process).and_return({ success: "ok" }) + include_context "bot execution setup" allow(@shared_storage_reader).to receive(:respond_to?).with(:close_connections).and_return(true) allow(@shared_storage_writer).to receive(:respond_to?).with(:close_connections).and_return(true)lib/bas/bot/base.rb (1)
35-35: Use truthy check instead of strict equalityThe strict equality check with
truemight miss truthy values. Consider using a more flexible check.-close_connections if @process_options[:close_connections_after_process].eql?(true) +close_connections if @process_options[:close_connections_after_process]lib/bas/shared_storage/postgres.rb (1)
20-26: Consider handling connection errors gracefullyWhile the connection is established lazily, query execution might fail if the connection is lost. Based on the codebase's error handling pattern (from learnings), errors should bubble up, but you might want to add connection health checks for long-running processes.
Add a comment to document the error handling behavior:
def read establish_connection(:read) + # Errors from query execution are intentionally not caught here + # to allow explicit error handling at higher levels first_result = @read_connection.query(read_query).first || {}spec/bas/utils/postgres/connection_spec.rb (1)
26-33: Consider making the mock data more flexibleThe mock setup hardcodes specific test data. Consider parameterizing the mock data or using a factory pattern to make tests more maintainable and flexible for different test scenarios.
- allow(pg_result).to receive(:map) do - # Simulate the actual behavior where each result hash gets its keys transformed to symbols - raw_data = [ - { "id" => "1", "name" => "John Doe", "email" => "john@example.com" }, - { "id" => "2", "name" => "Jane Smith", "email" => "jane@example.com" } - ] - raw_data.map { |result| result.transform_keys(&:to_sym) } - end + # Default mock data that can be overridden in specific tests + default_raw_data = [ + { "id" => "1", "name" => "John Doe", "email" => "john@example.com" }, + { "id" => "2", "name" => "Jane Smith", "email" => "jane@example.com" } + ] + allow(pg_result).to receive(:map) do + # Allow tests to override this behavior if needed + (defined?(@custom_raw_data) ? @custom_raw_data : default_raw_data).map { |result| result.transform_keys(&:to_sym) } + end
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
lib/bas/bot/base.rb(3 hunks)lib/bas/shared_storage/POSTGRES.md(1 hunks)lib/bas/shared_storage/base.rb(1 hunks)lib/bas/shared_storage/postgres.rb(4 hunks)lib/bas/utils/postgres/connection.rb(1 hunks)spec/bas/bot/base_spec.rb(2 hunks)spec/bas/shared_storage/postgres_spec.rb(3 hunks)spec/bas/utils/postgres/connection_spec.rb(1 hunks)
🧰 Additional context used
🧠 Learnings (1)
spec/bas/bot/base_spec.rb (1)
Learnt from: juanpabloxk
PR: kommitters/bas#138
File: lib/bas/shared_storage/elasticsearch.rb:25-35
Timestamp: 2025-07-02T19:18:20.943Z
Learning: In the Bas codebase, errors from shared storage operations (like Elasticsearch requests) should be thrown to the caller rather than being caught and handled within the storage methods themselves. This allows for explicit error handling at higher levels in the application.
🧬 Code Graph Analysis (4)
lib/bas/shared_storage/base.rb (2)
lib/bas/bot/base.rb (1)
close_connections(66-69)lib/bas/shared_storage/postgres.rb (1)
close_connections(34-39)
lib/bas/bot/base.rb (2)
lib/bas/shared_storage/base.rb (1)
close_connections(24-27)lib/bas/shared_storage/postgres.rb (1)
close_connections(34-39)
spec/bas/shared_storage/postgres_spec.rb (1)
lib/bas/shared_storage/postgres.rb (6)
read(16-109)read(19-26)write(28-32)close_connections(34-39)set_in_process(41-45)set_processed(47-51)
spec/bas/utils/postgres/connection_spec.rb (1)
lib/bas/utils/postgres/connection.rb (2)
query(15-25)finish(27-30)
🔇 Additional comments (3)
lib/bas/shared_storage/postgres.rb (1)
34-39: LGTM! Proper connection cleanup implementationThe
close_connectionsmethod correctly closes both connections and clears the instance variables to prevent reuse of closed connections.lib/bas/shared_storage/POSTGRES.md (1)
1-91: Documentation does not reflect the refactored connection management approachThis sequence diagram documents the old request-based architecture where each operation creates a new connection via
Utils::Postgres::Request. However, according to the PR objectives, the code has been refactored to use persistent connections with:
- Memoized
@read_connectionand@write_connectioninstance variables- Connection reuse across operations
- Explicit
close_connectionsmethodUtils::Postgres::Connectioninstead ofUtils::Postgres::RequestThe documentation should be updated to reflect the new connection management pattern.
Likely an incorrect or invalid review comment.
spec/bas/shared_storage/postgres_spec.rb (1)
1-227: Excellent test coverage for the refactored connection managementThe test suite comprehensively validates the new persistent connection behavior:
- Connection reuse for multiple operations
- Proper connection closing and re-establishment
- Separate connections for read/write operations
- Edge cases like empty results and nil IDs
The tests effectively verify that the refactoring achieves its goal of reducing connection overhead.
Pull Request Test Coverage Report for Build 16427129026Details
💛 - Coveralls |
There was a problem hiding this comment.
LGTM! 🚀
@juanpabloxk could you check the coveralls and make sure they cover 100%?
juanpabloxk
force-pushed
the
fix-postgres-pooling
branch
from
3523c4a to
a03d3db
Compare
@FelipeGuzmanSierra Sure, done ✅ |
Closes #147
Previous implementation
sequenceDiagram participant User as User/Code participant SS as SharedStorage::Postgres participant Req as Utils::Postgres::Request participant PG as Postgres Server User->>SS: new(read_options, write_options) activate SS Note over SS: Instance created with options User->>SS: read() activate SS SS->>Req: execute({connection: read_options[:connection], query: read_query}) activate Req Req->>PG: PG::Connection.new(connection params) activate PG Note over Req,PG: Connection established Req->>PG: exec_params(sentence, params) or exec(query) PG->>Req: Query results Req->>Req: map results to symbols deactivate PG Note over Req,PG: Connection not explicitly closed (auto-closed on GC) Req->>SS: Return mapped results deactivate Req SS->>SS: Build Read response SS->>User: Return Read object deactivate SS User->>SS: set_in_process() activate SS alt avoid_process == true or id nil SS->>User: Return nil else SS->>Req: execute({connection: read_options[:connection], query: update_query(id, "in process")}) activate Req Req->>PG: PG::Connection.new(connection params) activate PG Note over Req,PG: Connection established Req->>PG: exec_params(sentence, params) PG->>Req: Update result deactivate PG Note over Req,PG: Connection not explicitly closed Req->>SS: Return result deactivate Req SS->>User: Return result end deactivate SS User->>SS: write(data) activate SS SS->>Req: execute({connection: write_options[:connection], query: write_query(data)}) activate Req Req->>PG: PG::Connection.new(connection params) activate PG Note over Req,PG: Connection established Req->>PG: exec_params(sentence, params) PG->>Req: Insert result deactivate PG Note over Req,PG: Connection not explicitly closed Req->>SS: Return result deactivate Req SS->>User: Return result deactivate SS User->>SS: set_processed() activate SS alt avoid_process == true or id nil SS->>User: Return nil else SS->>Req: execute({connection: read_options[:connection], query: update_query(id, "processed")}) activate Req Req->>PG: PG::Connection.new(connection params) activate PG Note over Req,PG: Connection established Req->>PG: exec_params(sentence, params) PG->>Req: Update result deactivate PG Note over Req,PG: Connection not explicitly closed Req->>SS: Return result deactivate Req SS->>User: Return result end deactivate SS deactivate SSNew Implementation
sequenceDiagram participant User as User/Code participant SS as SharedStorage::Postgres participant RConn as Read Connection (Utils::Postgres::Connection) participant WConn as Write Connection (Utils::Postgres::Connection) participant PG as Postgres Server User->>SS: new(options) activate SS Note over SS: Instance created with @read_options and @write_options User->>SS: read() activate SS alt @read_connection nil SS->>RConn: new(read_options[:connection]) activate RConn RConn->>PG: PG::Connection.new(params) activate PG Note over RConn,PG: Read connection established end SS->>RConn: query(read_query) # [sentence, params] RConn->>PG: exec_params(sentence, params) PG->>RConn: Query results RConn->>RConn: map results to symbols RConn->>SS: Return mapped results SS->>SS: Build Read response SS->>User: Return Read object deactivate SS User->>SS: set_in_process() activate SS alt avoid_process == true or read_response.id nil SS->>User: Return nil else alt @read_connection nil SS->>RConn: new(read_options[:connection]) activate RConn RConn->>PG: PG::Connection.new(params) activate PG end SS->>RConn: query(update_query(id, "in process")) RConn->>PG: exec_params(sentence, params) PG->>RConn: Update result RConn->>SS: Return result SS->>User: Return result end deactivate SS User->>SS: write(data) activate SS alt @write_connection nil SS->>WConn: new(write_options[:connection]) activate WConn WConn->>PG: PG::Connection.new(params) activate PG Note over WConn,PG: Write connection established (may be same or different PG) end SS->>WConn: query(write_query(data)) WConn->>PG: exec_params(sentence, params) PG->>WConn: Insert result WConn->>SS: Return result SS->>User: Return result deactivate SS User->>SS: set_processed() activate SS alt avoid_process == true or read_response.id nil SS->>User: Return nil else SS->>RConn: query(update_query(id, "processed")) RConn->>PG: exec_params(sentence, params) PG->>RConn: Update result RConn->>SS: Return result SS->>User: Return result end deactivate SS User->>SS: close_connections() activate SS alt @read_connection exists SS->>RConn: finish() RConn->>PG: finish deactivate PG deactivate RConn end alt @write_connection exists SS->>WConn: finish() WConn->>PG: finish deactivate PG deactivate WConn end deactivate SS deactivate SSTesting
I tested this new implementation with two use cases from bas_use_cases
Adding some logging to know when connections are instantiated:
Only one connection is instantiated no matter how many items are processed:
Summary by CodeRabbit
New Features
Bug Fixes
Tests