feat(rain): add drizzle-like read replica routing

[cungminh2710]

What this PR for?

Add WithReplicas and Primary() to route builder-based reads to replicas while keeping writes, raw SQL, and transactions on primary. Includes sqlite-backed routing tests and docs/examples for replica reads, forced primary reads, and custom selector behavior.

What did you do for validating the changes?

Anything else you want to add? (Optional)

[greptile-apps]

Greptile Summary

This PR adds drizzle-style read replica routing to rain ORM. WithReplicas wraps a primary *DB and one or more replica handles into a single routed handle: Select() queries are dispatched to replicas (via a pluggable ReplicaSelector), while Insert, Update, Delete, Exec, Query, QueryRow, Begin, and RunInTx always target the primary. Primary() returns a transient view that forces even Select to go to primary, which is useful for read-after-write consistency.

Key design decisions:

• A shared dbSharedState struct carries the QueryCache across all handles; mutating it on any handle (primary, replica, or routed) propagates to all others.
• replicaRoute.close() uses a sync.Once to deduplicate sql.DB.Close() calls and is shared across all views of the same routed handle.
• A custom ReplicaSelector that returns an unknown handle gracefully falls back to randomReplicaSelector rather than panicking.
• Relation loading (WithRelations) correctly threads the chosen replica runner through batch sub-queries, so related rows are fetched from the same replica as the root query.

The implementation is clean and the test suite is thorough, covering validation, default/custom selectors, primary views, relation loading, writes, transactions, raw SQL, cache sharing, close idempotency, and concurrent reads.

One P2 concern: WithReplicas silently overwrites each replica handle's shared pointer with the primary's dbSharedState. Any QueryCache previously configured on a replica via replica.WithQueryCache(cache) is discarded without error. The safe pattern — configure cache only after calling WithReplicas — should be documented in the WithReplicas doc comment.

Confidence Score: 4/5

Safe to merge after addressing or documenting the replica QueryCache overwrite behaviour.

All routing logic is correct and well-tested. The single remaining concern is a P2 API footgun: pre-configured QueryCache on replica handles is silently discarded by WithReplicas. This won't break users who follow the intended pattern (set cache post-construction on the routed handle), but could surprise users who configure replicas individually before composing them. A doc comment fix is sufficient.

pkg/rain/rain.go — specifically the WithReplicas function and its interaction with pre-configured replica shared state.

Important Files Changed

Filename	Overview
pkg/rain/rain.go	Core routing logic is sound: WithReplicas, Primary(), selectRunner(), and replicaRoute.close() are all correctly implemented. One P2 concern: WithReplicas mutates replica handles to overwrite any previously configured QueryCache without warning.
pkg/rain/read_replica_internal_test.go	Comprehensive test coverage across all routing scenarios (default selector, custom selector, primary view, fallback, relation loading, writes, transactions, raw SQL, cache sharing, close deduplication, and concurrent reads). Tests are well-structured and use isolated SQLite files per test.
README.md	Documentation correctly illustrates replica setup, Primary() usage for read-after-write, and the routing contract for writes/transactions. Notes about replica lag are helpful.
examples/basic/main.go	Example updated to demonstrate WithReplicas and Primary().Select() usage. Error for ToSQL is intentionally discarded with _ which is acceptable in example code.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    Client["Caller"]

    subgraph RoutedDB["Routed DB (WithReplicas result)"]
        direction TB
        Select["Select()"]
        Write["Insert / Update / Delete"]
        RawSQL["Exec / Query / QueryRow / Begin / RunInTx"]
        PrimaryView["Primary().Select()"]
    end

    subgraph selectRunner["selectRunner()"]
        Check{"forcePrimaryReads\nor no replicaRoute?"}
        PickReplica["replicaRoute.pickReplica()\n(ReplicaSelector → fallback rand)"]
    end

    Primary["Primary DB\n(sql.DB)"]
    Replica1["Replica 1\n(sql.DB)"]
    ReplicaN["Replica N\n(sql.DB)"]

    Client --> Select
    Client --> Write
    Client --> RawSQL
    Client --> PrimaryView

    Select --> Check
    PrimaryView --> Check
    Check -- "yes (forced / no route)" --> Primary
    Check -- "no" --> PickReplica
    PickReplica --> Replica1
    PickReplica --> ReplicaN

    Write --> Primary
    RawSQL --> Primary

Loading

Prompt To Fix All With AI

This is a comment left during a code review.
Path: pkg/rain/rain.go
Line: 117

Comment:
**Replica's pre-configured `QueryCache` is silently discarded**

`replica.shared = shared` overwrites each replica's existing `shared` pointer with the primary's `dbSharedState`. Any `QueryCache` a caller had previously set on the replica handle via `replica.WithQueryCache(cache)` is silently dropped and replaced by the primary's (initially nil) cache.

The intended idiom is to call `routed.WithQueryCache(cache)` *after* `WithReplicas`, which then propagates through the shared pointer to every handle. But if a caller follows the intuitive pattern of configuring individual handles before composing them, their cache configuration is silently lost with no error or warning.

At minimum, the `WithReplicas` doc comment should warn that any `QueryCache` previously set on the replica handles will be discarded in favour of the primary's shared state, and that cache must be configured on the returned routed handle (or the primary) instead.

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (2): Last reviewed commit: "." | Re-trigger Greptile}