Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 43 additions & 30 deletions go/api/gen/kagent/config/v1/config.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

23 changes: 19 additions & 4 deletions go/core/internal/database/client_postgres.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"time"

"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
Expand All @@ -12,21 +13,35 @@ import (
// Client persists control-plane state in PostgreSQL. Callers define the narrow
// interfaces they need; SQL rows and protobuf encoding stay inside the store.
type Client struct {
db *pgxpool.Pool
db retryingPool
vectorCosineOperator string
}

// ClientOption configures a Client built by NewClient.
type ClientOption func(*Client)

// WithConnectRetry makes each store call retry a failed PostgreSQL login for up
// to timeout. A call can take up to timeout plus one attempt. Rejected
// credentials are not retried. Zero makes one attempt.
func WithConnectRetry(timeout time.Duration) ClientOption {
return func(c *Client) { c.db.retryFor = timeout }
}

// NewClient wraps an existing PostgreSQL pool without connecting or migrating. The caller
// owns the pool and must close it. vectorSchema is the schema pgvector is installed in.
func NewClient(db *pgxpool.Pool, vectorSchema string) *Client {
return &Client{db: db, vectorCosineOperator: "OPERATOR(" + pgx.Identifier{vectorSchema}.Sanitize() + ".<=>)"}
func NewClient(db *pgxpool.Pool, vectorSchema string, opts ...ClientOption) *Client {
c := &Client{db: retryingPool{pool: db}, vectorCosineOperator: "OPERATOR(" + pgx.Identifier{vectorSchema}.Sanitize() + ".<=>)"}
for _, opt := range opts {
opt(c)
}
return c
}

// withTx commits all callback writes together on success and rolls them back on failure,
// returning callback or transaction errors. The callback must keep external network work
// outside the transaction.
func (c *Client) withTx(ctx context.Context, fn func(pgx.Tx) error) error {
tx, err := c.db.Begin(ctx)
tx, err := c.db.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
return fmt.Errorf("failed to begin transaction: %w", err)
}
Expand Down
Loading
Loading