Go로 PostgreSQL 프록시 만들기 (5) - 통합, E2E 테스트, 회고

// Before: NewServer에서 Pool, Router, Cache를 초기화하지 않음
func NewServer(cfg *config.Config) *Server {
    return &Server{
        listenAddr: cfg.Proxy.Listen,
        writerAddr: fmt.Sprintf("%s:%d", cfg.Writer.Host, cfg.Writer.Port),
    }
}

// Before: 쿼리를 무조건 writer 1대로 직접 릴레이
func (s *Server) relayQueries(ctx context.Context, clientConn, backendConn net.Conn) {
    // ... 파싱이나 분류 없이 단순 전달
}

Client connects
    ↓
Read startup (SSL negotiate)
    ↓
Connect to WRITER (direct, per-client)
    ↓ Auth relay (bidirectional)
Session created (Router)
    ↓
Query Loop:
  classify → route → cache check → forward → cache store/invalidate

type DialFunc func() (net.Conn, error)

type Config struct {
    DialFunc          DialFunc // nil이면 기존 TCP dial 사용
    Addr              string
    MinConnections    int
    MaxConnections    int
    // ...
}

for _, addr := range readerAddrs {
    addr := addr
    p, _ := pool.New(pool.Config{
        DialFunc: func() (net.Conn, error) {
            return pgConnect(addr, cfg.Backend.User, cfg.Backend.Password, cfg.Backend.Database)
        },
        MinConnections: 0, // lazy creation
        MaxConnections: cfg.Pool.MaxConnections,
    })
    s.readerPools[addr] = p
}

func handleSCRAM(conn net.Conn, mechanisms []byte, user, password string) error {
    // 1. SASLInitialResponse: client-first-message
    clientNonce := generateNonce()
    clientFirstBare := fmt.Sprintf("n=%s,r=%s", user, clientNonce)

    // 2. Read SASLContinue: server-first-message (nonce, salt, iterations)
    // 3. Compute PBKDF2 → HMAC-SHA-256 → client proof
    saltedPassword := pbkdf2.Key([]byte(password), salt, iterations, 32, sha256.New)
    clientKey := hmacSHA256(saltedPassword, []byte("Client Key"))
    // ... XOR로 proof 생성

    // 4. Read SASLFinal: server signature 검증
}

Backend → AuthenticationMD5Password → Client
Client  → PasswordMessage          → Backend   ← 이 방향이 빠져있었다!
Backend → AuthenticationOk          → Client

func (s *Server) relayAuth(clientConn, backendConn net.Conn) error {
    for {
        msg, _ := protocol.ReadMessage(backendConn)
        protocol.WriteMessage(clientConn, msg.Type, msg.Payload)

        // 인증 요청이면 클라이언트 응답을 읽어서 백엔드로 전달
        if msg.Type == MsgAuthentication {
            authType := binary.BigEndian.Uint32(msg.Payload[0:4])
            if authNeedsResponse(authType) { // MD5(5), SCRAM(10,11)
                clientMsg, _ := protocol.ReadMessage(clientConn)
                protocol.WriteMessage(backendConn, clientMsg.Type, clientMsg.Payload)
            }
        }
    }
}

Parse(P) → Bind(B) → Execute(E) → Sync(S)

if msg.Type != protocol.MsgQuery {
    protocol.WriteMessage(writerConn, msg.Type, msg.Payload)
    // Sync일 때만 응답을 대기
    if msg.Type == protocol.MsgSync {
        s.relayUntilReady(clientConn, writerConn)
    }
    continue
}

query := protocol.ExtractQueryText(msg.Payload)
route := session.Route(query)  // Session이 트랜잭션/R-A-W 상태 추적

if route == RouteWriter {
    forwardAndRelay(clientConn, writerConn, msg)
    // 쓰기 쿼리면 캐시 무효화
    if Classify(query) == QueryWrite {
        for _, table := range ExtractTables(query) {
            queryCache.InvalidateTable(table)
        }
    }
} else { // RouteReader
    // 캐시 확인
    if cached := queryCache.Get(CacheKey(query)); cached != nil {
        clientConn.Write(cached) // 캐시 히트: PG 응답 바이트 그대로 전송
        continue
    }
    // Reader 풀에서 연결 획득 (RoundRobin)
    addr := balancer.Next()
    rConn, _ := readerPools[addr].Acquire(ctx)
    // 쿼리 전달 + 응답 수집 + 캐시 저장
    collected, _ := relayAndCollect(clientConn, rConn)
    queryCache.Set(CacheKey(query), collected, nil)
    readerPools[addr].Release(rConn)
}

# docker-compose.yml: Primary + Replica x2
services:
  primary:   # port 15432, wal_level=replica
  replica1:  # port 15433, pg_basebackup from primary
  replica2:  # port 15434, pg_basebackup from primary

func TestE2E_ProxyIntegration(t *testing.T) {
    db, _ := sql.Open("postgres", "postgres://postgres:postgres@127.0.0.1:15440/testdb?sslmode=disable")

    t.Run("ReadQuery", func(t *testing.T) {
        // SELECT → reader로 라우팅 → RoundRobin 분배
    })
    t.Run("WriteQuery", func(t *testing.T) {
        // INSERT → writer → Extended Query Protocol 정상 처리
    })
    t.Run("Transaction", func(t *testing.T) {
        // BEGIN~ROLLBACK → 전체 writer 라우팅
    })
    t.Run("CacheHit", func(t *testing.T) {
        // 동일 쿼리 두 번 → 두 번째는 캐시 히트
    })
}

query routed sql="SELECT name FROM users ORDER BY id LIMIT 3" route=reader
cache set   sql="SELECT name FROM users ORDER BY id LIMIT 3" size=98
query routed sql="DELETE FROM users WHERE name = 'dave'" route=writer
cache invalidated table=users
query routed sql="BEGIN READ WRITE" route=writer
query routed sql="INSERT INTO users ..." route=writer
query routed sql="SELECT name FROM users WHERE name = 'tx_user'" route=writer  ← 트랜잭션 내
query routed sql=ROLLBACK route=writer