Add db and migrations

Author: simondanielsson

.env.db.example

@@ -0,0 +1,3 @@
+POSTGRES_DB=recite
+POSTGRES_USER=recite
+POSTGRES_PASSWORD=password

.env.example

@@ -1 +1,9 @@
 OPENAI_API_KEY=<your_openai_api_key>
+
+POSTGRES_DBNAME=recite
+POSTGRES_USER=recite
+POSTGRES_PASSWORD=password
+POSTGRES_HOST=localhost
+POSTGRES_PORT=5427
+DRIVER_NAME=pgx
+DATABASE_URL=postgres://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${POSTGRES_HOST}:${POSTGRES_PORT}/${POSTGRES_DBNAME}?sslmode=disable

.gitignore

@@ -26,3 +26,6 @@ go.work.sum
 
 # Executables
 bin/
+data/output/
+.env
+.env.db

README.md

@@ -36,3 +36,16 @@ Always verify your code using
 make ci
 ```
 
+### Adding database migrations
+
+We use `dbmate` for migrations. After installing, create a new revision through
+
+```bash
+dbmate new <name_of_revision>
+```
+
+Apply the revision by running
+
+```bash
+dbmate up
+```

cmd/cli/main.go

@@ -14,6 +14,8 @@ import (
 	"github.com/simondanielsson/recite/pkg/prompts"
 )
 
+const maxContentLength int = 3500
+
 func main() {
 	// TODO: use a proper CLI library
 	if len(os.Args) < 2 {
@@ -32,7 +34,7 @@ func main() {
 	if err != nil {
 		log.Fatalf("Failed reading article: %v", err)
 	}
-	articleContent = articleContent[:3500]
+	articleContent = articleContent[:min(maxContentLength, len(articleContent))]
 
 	augmentPrompts, err := prompts.NewAugmentArticlePrompts(articleContent)
 	if err != nil {

cmd/internal/config/config.go

@@ -12,13 +12,22 @@ func Load(getenv func(string) string, logger *log.Logger) (Config, error) {
 	if err := env.Load(); err != nil {
 		logger.Fatal("failed loading .env")
 	}
+	// TODO: read these values from yaml config
 	cfg := Config{
 		Server: ServerConfig{
 			AppEnv:       "local",
 			Addr:         "8999",
 			ReadTimeout:  30 * time.Second,
 			WriteTimeout: 30 * time.Second,
 		},
+		DB: DBConfig{
+			Name:     getenv("POSTGRES_DBNAME"),
+			User:     getenv("POSTGRES_USER"),
+			Password: getenv("POSTGRES_PASSWORD"),
+			Host:     getenv("POSTGRES_HOST"),
+			Port:     getenv("POSTGRES_PORT"),
+			Driver:   getenv("DRIVER_NAME"),
+		},
 		JWT: JWTConfig{
 			Secret: "123",
 		},

cmd/internal/config/models.go

@@ -10,6 +10,15 @@ type ServerConfig struct {
 	WriteTimeout time.Duration `mapstructure:"write_timeout"`
 }
 
+type DBConfig struct {
+	Name     string `mapstructure:"name"`
+	User     string `mapstructure:"user"`
+	Password string `mapstructure:"password"`
+	Host     string `mapstructure:"host"`
+	Port     string `mapstructure:"port"`
+	Driver   string `mapstructure:"driver"`
+}
+
 // JWTConfig represents the configuration of the JWT secret
 type JWTConfig struct {
 	Secret string `mapstructure:"secret"`
@@ -18,5 +27,6 @@ type JWTConfig struct {
 // Config represents a yaml configuration file.
 type Config struct {
 	Server ServerConfig
+	DB     DBConfig
 	JWT    JWTConfig
 }

cmd/internal/db/db.go

@@ -0,0 +1,22 @@
+package db
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/simondanielsson/recite/cmd/internal/config"
+)
+
+func New(ctx context.Context, config config.Config) (*pgxpool.Pool, error) {
+	connString := fmt.Sprintf("user=%s dbname=%s password=%s port=%s host=%s sslmode=disable", config.DB.User, config.DB.Name, config.DB.Password, config.DB.Port, config.DB.Host)
+	pool, err := pgxpool.New(ctx, connString)
+	if err != nil {
+		return nil, fmt.Errorf("could not connect to db: %w", err)
+	}
+	if err := pool.Ping(ctx); err != nil {
+		return nil, fmt.Errorf("failed pinging database: %w", err)
+	}
+
+	return pool, nil
+}

cmd/internal/db/middleware.go

@@ -0,0 +1,48 @@
+package db
+
+import (
+	"context"
+	"log"
+	"net/http"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/simondanielsson/recite/cmd/internal/queries"
+)
+
+type RepositoryKeyType string
+
+const RepositoryKey RepositoryKeyType = "repo_key"
+
+type statusRecorder struct {
+	http.ResponseWriter
+	status int
+}
+
+func AddDatabaseMiddleware(handler http.Handler, pool *pgxpool.Pool, logger *log.Logger) http.Handler {
+	return http.HandlerFunc(
+		func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+			tx, err := pool.BeginTx(ctx, pgx.TxOptions{
+				IsoLevel:   pgx.Serializable,
+				AccessMode: pgx.ReadWrite,
+			})
+			if err != nil {
+				http.Error(w, "failed to start transaction", http.StatusInternalServerError)
+			}
+			repository := queries.New(pool).WithTx(tx)
+
+			ctx = context.WithValue(ctx, RepositoryKey, repository)
+
+			rec := statusRecorder{ResponseWriter: w, status: http.StatusOK}
+			handler.ServeHTTP(rec, r.WithContext(ctx))
+
+			if rec.status >= http.StatusBadRequest {
+				_ = tx.Rollback(ctx)
+			} else if err := tx.Commit(ctx); err != nil {
+				// Response has already been sent - just log
+				logger.Printf("commit tx failed: %v", err)
+			}
+		},
+	)
+}