Skip to content
Draft
Show file tree
Hide file tree
Changes from 1 commit
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
1 change: 1 addition & 0 deletions packages/api/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ require (
github.com/dprotaso/go-yit v0.0.0-20220510233725-9ba8df137936 // indirect
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/edsrzf/mmap-go v1.2.0 // indirect
github.com/exaring/otelpgx v0.9.3 // indirect
github.com/facette/natsort v0.0.0-20181210072756-2cd4dd1e2dcb // indirect
github.com/fatih/color v1.18.0 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
Expand Down
2 changes: 2 additions & 0 deletions packages/api/go.sum

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

4 changes: 3 additions & 1 deletion packages/api/internal/cfg/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@ type Config struct {
NomadAddress string `env:"NOMAD_ADDRESS" envDefault:"http://localhost:4646"`
NomadToken string `env:"NOMAD_TOKEN"`

PostgresConnectionString string `env:"POSTGRES_CONNECTION_STRING,required,notEmpty"`
PostgresConnectionString string `env:"POSTGRES_CONNECTION_STRING,required,notEmpty"`
PostgresMaxConnections int `env:"POSTGRES_MAX_CONNECTIONS" envDefault:"40"`
PostgresMinIdleConnections int `env:"POSTGRES_MIN_IDLE_CONNECTIONS" envDefault:"5"`

PosthogAPIKey string `env:"POSTHOG_API_KEY"`

Expand Down
25 changes: 6 additions & 19 deletions packages/api/internal/handlers/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package handlers

import (
"context"
"database/sql"
"errors"
"fmt"
"net/http"
Expand All @@ -11,6 +12,7 @@ import (
"github.com/golang-jwt/jwt/v5"
"github.com/google/uuid"
nomadapi "github.com/hashicorp/nomad/api"
"github.com/jackc/pgx/v5/pgxpool"
middleware "github.com/oapi-codegen/gin-middleware"
"github.com/redis/go-redis/v9"
"go.uber.org/zap"
Expand Down Expand Up @@ -58,22 +60,15 @@ type APIStore struct {
clustersPool *edge.Pool
}

func NewAPIStore(ctx context.Context, tel *telemetry.Client, config cfg.Config) *APIStore {
func NewAPIStore(ctx context.Context, tel *telemetry.Client, config cfg.Config, sqlConn *sql.DB, pool *pgxpool.Pool) *APIStore {
zap.L().Info("Initializing API store and services")

dbClient, err := db.NewClient(40, 20)
if err != nil {
zap.L().Fatal("Initializing Supabase client", zap.Error(err))
}

sqlcDB, err := sqlcdb.NewClient(ctx, sqlcdb.WithMaxConnections(40), sqlcdb.WithMinIdle(5))
if err != nil {
zap.L().Fatal("Initializing SQLC client", zap.Error(err))
}
dbClient := db.NewClient(sqlConn)

zap.L().Info("Created database client")
sqlcDB := sqlcdb.NewClient(pool)

var clickhouseStore clickhouse.Clickhouse
var err error

clickhouseConnectionString := config.ClickhouseConnectionString
if clickhouseConnectionString == "" {
Expand Down Expand Up @@ -216,14 +211,6 @@ func (a *APIStore) Close(ctx context.Context) error {
errs = append(errs, fmt.Errorf("closing Template manager client: %w", err))
}

if err := a.sqlcDB.Close(); err != nil {
errs = append(errs, fmt.Errorf("closing sqlc database client: %w", err))
}

if err := a.db.Close(); err != nil {
errs = append(errs, fmt.Errorf("closing database client: %w", err))
}

return errors.Join(errs...)
}

Expand Down
13 changes: 1 addition & 12 deletions packages/api/internal/utils/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,7 @@ import (

const trackingTable = "_migrations"

func CheckMigrationVersion(connectionString string, expectedMigration int64) error {
db, err := sql.Open("postgres", connectionString)
if err != nil {
return fmt.Errorf("failed to connect: %w", err)
}
defer func() {
dbErr := db.Close()
if dbErr != nil {
zap.L().Error("Failed to close database connection", zap.Error(dbErr))
}
}()

func CheckMigrationVersion(db *sql.DB, expectedMigration int64) error {
goose.SetTableName(trackingTable)

version, err := goose.GetDBVersion(db)
Expand Down
16 changes: 14 additions & 2 deletions packages/api/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ import (
metricsMiddleware "github.com/e2b-dev/infra/packages/api/internal/middleware/otel/metrics"
tracingMiddleware "github.com/e2b-dev/infra/packages/api/internal/middleware/otel/tracing"
"github.com/e2b-dev/infra/packages/api/internal/utils"
"github.com/e2b-dev/infra/packages/shared/pkg/db"
"github.com/e2b-dev/infra/packages/shared/pkg/env"
l "github.com/e2b-dev/infra/packages/shared/pkg/logger"
sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
Expand Down Expand Up @@ -278,7 +279,18 @@ func run() int {
zap.L().Fatal("Error parsing config", zap.Error(err))
}

err = utils.CheckMigrationVersion(config.PostgresConnectionString, expectedMigration)
dbPool, err := db.NewPool(ctx,
db.WithMaxConnections(config.PostgresMaxConnections),
db.WithMinIdle(config.PostgresMinIdleConnections),
)
if err != nil {
logger.Fatal("failed to create database pool", zap.Error(err))
}

dbConn := db.Open(dbPool)
defer dbConn.Close()

err = utils.CheckMigrationVersion(dbConn, expectedMigration)
if err != nil {
logger.Fatal("failed to check migration version", zap.Error(err))
}
Expand Down Expand Up @@ -343,7 +355,7 @@ func run() int {
// Create an instance of our handler which satisfies the generated interface
// (use the outer context rather than the signal handling
// context so it doesn't exit first.)
apiStore := handlers.NewAPIStore(ctx, tel, config)
apiStore := handlers.NewAPIStore(ctx, tel, config, dbConn, dbPool)
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bug: Database Pool Not Closed on Shutdown

The database connection pool (dbPool) is created but not reliably closed across multiple services. This leads to resource leaks, as dbPool.Close() is either missing from cleanup routines or placed where it won't execute during graceful application shutdown.

Additional Locations (2)

Fix in Cursor Fix in Web

cleanupFns = append(cleanupFns, apiStore.Close)

// pass the signal context so that handlers know when shutdown is happening.
Expand Down
62 changes: 2 additions & 60 deletions packages/db/client/client.go
Original file line number Diff line number Diff line change
@@ -1,76 +1,18 @@
package client

import (
"context"

"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
_ "github.com/lib/pq"
"go.uber.org/zap"

database "github.com/e2b-dev/infra/packages/db/queries"
"github.com/e2b-dev/infra/packages/shared/pkg/utils"
)

type Client struct {
*database.Queries

conn *pgxpool.Pool
}

type Option func(config *pgxpool.Config)

func WithMaxConnections(maxConns int32) Option {
return func(config *pgxpool.Config) {
config.MaxConns = maxConns
}
}

func WithMinIdle(minIdle int32) Option {
return func(config *pgxpool.Config) {
config.MinIdleConns = minIdle
}
}

func NewClient(ctx context.Context, options ...Option) (*Client, error) {
databaseURL := utils.RequiredEnv("POSTGRES_CONNECTION_STRING", "Postgres connection string")

// Parse the connection pool configuration
config, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
zap.L().Error("Unable to parse database URL", zap.Error(err))

return nil, err
}

// Set the default number of connections
for _, option := range options {
option(config)
}

// Create the connection pool
pool, err := pgxpool.NewWithConfig(ctx, config)
if err != nil {
zap.L().Error("Unable to create connection pool", zap.Error(err))
}

func NewClient(pool *pgxpool.Pool) *Client {
queries := database.New(pool)

return &Client{Queries: queries, conn: pool}, nil
}

func (db *Client) Close() error {
db.conn.Close()
return nil
}

// WithTx runs the given function in a transaction.
func (db *Client) WithTx(ctx context.Context) (*Client, pgx.Tx, error) {
tx, err := db.conn.BeginTx(ctx, pgx.TxOptions{})
if err != nil {
return nil, nil, err
}

client := &Client{Queries: db.Queries.WithTx(tx), conn: db.conn}
return client, tx, nil
return &Client{Queries: queries}
}
11 changes: 8 additions & 3 deletions packages/db/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,17 @@ tool (
)

require (
github.com/e2b-dev/infra/packages/shared v0.0.0-20250324174051-3fb806938dc1
github.com/google/uuid v1.6.0
github.com/jackc/pgx/v5 v5.7.4
github.com/lib/pq v1.10.9
github.com/pressly/goose/v3 v3.24.2
go.uber.org/zap v1.27.0
)

require (
cel.dev/expr v0.24.0 // indirect
filippo.io/edwards25519 v1.1.0 // indirect
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.14.0 // indirect
github.com/Azure/azure-sdk-for-go/sdk/azidentity v1.7.0 // indirect
github.com/ClickHouse/ch-go v0.66.1 // indirect
github.com/ClickHouse/clickhouse-go/v2 v2.37.2 // indirect
github.com/andybalholm/brotli v1.1.1 // indirect
Expand Down Expand Up @@ -48,6 +48,7 @@ require (
github.com/joho/godotenv v1.5.1 // indirect
github.com/jonboulle/clockwork v0.5.0 // indirect
github.com/klauspost/compress v1.18.0 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mfridman/interpolate v0.0.2 // indirect
github.com/mfridman/xflag v0.1.0 // indirect
Expand All @@ -60,9 +61,11 @@ require (
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 // indirect
github.com/pingcap/log v1.1.0 // indirect
github.com/pingcap/tidb/pkg/parser v0.0.0-20250324122243-d51e00e5bbf0 // indirect
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
github.com/prometheus/procfs v0.16.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/riza-io/grpc-go v0.2.0 // indirect
github.com/rogpeppe/go-internal v1.14.1 // indirect
github.com/segmentio/asm v1.2.0 // indirect
github.com/sethvargo/go-retry v0.3.0 // indirect
github.com/shopspring/decimal v1.4.0 // indirect
Expand All @@ -78,14 +81,16 @@ require (
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241112172322-ea1f63298f77 // indirect
github.com/ydb-platform/ydb-go-sdk/v3 v3.104.7 // indirect
github.com/ziutek/mymysql v1.5.4 // indirect
go.opentelemetry.io/auto/sdk v1.2.0 // indirect
go.opentelemetry.io/otel v1.38.0 // indirect
go.opentelemetry.io/otel/sdk/metric v1.38.0 // indirect
go.opentelemetry.io/otel/trace v1.38.0 // indirect
go.uber.org/atomic v1.11.0 // indirect
go.uber.org/mock v0.5.2 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.27.0 // indirect
golang.org/x/crypto v0.42.0 // indirect
golang.org/x/exp v0.0.0-20250531010427-b6e5de432a8b // indirect
golang.org/x/mod v0.27.0 // indirect
golang.org/x/net v0.44.0 // indirect
golang.org/x/sync v0.17.0 // indirect
golang.org/x/sys v0.36.0 // indirect
Expand Down
3 changes: 3 additions & 0 deletions packages/db/go.sum

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

14 changes: 14 additions & 0 deletions packages/docker-reverse-proxy/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,28 @@ require (
github.com/agext/levenshtein v1.2.3 // indirect
github.com/apparentlymart/go-textseg/v15 v15.0.0 // indirect
github.com/dchest/uniuri v1.2.0 // indirect
github.com/exaring/otelpgx v0.9.3 // indirect
github.com/go-logr/logr v1.4.3 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/go-openapi/inflect v0.21.0 // indirect
github.com/go-test/deep v1.0.8 // indirect
github.com/google/go-cmp v0.7.0 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/hashicorp/hcl/v2 v2.19.1 // indirect
github.com/jackc/pgpassfile v1.0.0 // indirect
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
github.com/jackc/pgx/v5 v5.7.4 // indirect
github.com/jackc/puddle/v2 v2.2.2 // indirect
github.com/lib/pq v1.10.9 // indirect
github.com/mitchellh/go-wordwrap v1.0.1 // indirect
github.com/zclconf/go-cty v1.14.1 // indirect
go.opentelemetry.io/auto/sdk v1.2.0 // indirect
go.opentelemetry.io/otel v1.38.0 // indirect
go.opentelemetry.io/otel/metric v1.38.0 // indirect
go.opentelemetry.io/otel/trace v1.38.0 // indirect
go.uber.org/multierr v1.11.0 // indirect
go.uber.org/zap v1.27.0 // indirect
golang.org/x/crypto v0.42.0 // indirect
golang.org/x/mod v0.27.0 // indirect
golang.org/x/sync v0.17.0 // indirect
golang.org/x/text v0.29.0 // indirect
Expand Down
Loading
Loading