Postgres seed

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.0.8](https://github.com/PerimeterX/envite/compare/v0.0.7...v0.0.8)
+
+### Added
+
+- Postgres Seed functionality.
+
 ## [0.0.7](https://github.com/PerimeterX/envite/compare/v0.0.6...v0.0.7)
 
 ### Fixed

seed/postgres/component.go

@@ -0,0 +1,144 @@
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"reflect"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+
+	"github.com/perimeterx/envite"
+)
+
+// ComponentType represents the type of the Postgres seed component.
+const ComponentType = "postgres seed"
+
+// SeedComponent is a component for seeding Postgres with data.
+type SeedComponent struct {
+	lock   sync.Mutex
+	config SeedConfig
+	status atomic.Value
+	writer *envite.Writer
+}
+
+// NewSeedComponent creates a new SeedComponent instance.
+func NewSeedComponent(config SeedConfig) *SeedComponent {
+	m := &SeedComponent{config: config}
+	m.status.Store(envite.ComponentStatusStopped)
+	return m
+}
+
+func (m *SeedComponent) Type() string {
+	return ComponentType
+}
+
+func (m *SeedComponent) AttachEnvironment(_ context.Context, _ *envite.Environment, writer *envite.Writer) error {
+	m.writer = writer
+	return nil
+}
+
+func (m *SeedComponent) Prepare(context.Context) error {
+	return nil
+}
+
+func (m *SeedComponent) Start(ctx context.Context) error {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	m.status.Store(envite.ComponentStatusStarting)
+
+	err := m.Seed()
+	if err != nil {
+		m.status.Store(envite.ComponentStatusFailed)
+		return err
+	}
+
+	m.status.Store(envite.ComponentStatusFinished)
+
+	return nil
+}
+
+func (m *SeedComponent) Seed() error {
+	if m.writer != nil {
+		m.writer.WriteString("starting postgres seed")
+	}
+
+	client, err := m.clientProvider()
+	if err != nil {
+		return err
+	}
+
+	if _, err = client.Exec(m.config.Setup); err != nil {
+		return err
+	}
+
+	for _, table := range m.config.Data {
+
+		if _, err = client.Exec(fmt.Sprintf("DELETE FROM %s", table.TableName)); err != nil {
+			return err
+		}
+
+		for _, row := range table.Rows {
+			sql, values := generateInsertSQL(table.TableName, row)
+			_, err := client.Exec(sql, values...)
+			if err != nil {
+				return err
+			}
+		}
+
+		if m.writer != nil {
+			m.writer.WriteString(fmt.Sprintf(
+				"inserted %s rows to %s",
+				m.writer.Color.Green(strconv.Itoa(len(table.Rows))),
+				m.writer.Color.Cyan(table.TableName),
+			))
+		}
+	}
+
+	return nil
+}
+
+func (m *SeedComponent) clientProvider() (*sql.DB, error) {
+	return m.config.ClientProvider()
+}
+
+func (m *SeedComponent) Stop(context.Context) error {
+	m.status.Store(envite.ComponentStatusStopped)
+	return nil
+}
+
+func (m *SeedComponent) Cleanup(context.Context) error {
+	return nil
+}
+
+func (m *SeedComponent) Status(context.Context) (envite.ComponentStatus, error) {
+	return m.status.Load().(envite.ComponentStatus), nil
+}
+
+func (m *SeedComponent) Config() any {
+	return m.config
+}
+
+func generateInsertSQL(table string, data any) (string, []any) {
+	v := reflect.ValueOf(data)
+	t := reflect.TypeOf(data)
+	var columns []string
+	var placeholders []string
+	var values []any
+	for i := 0; i < t.NumField(); i++ {
+		field := t.Field(i)
+		column := field.Tag.Get("column")
+		if column != "" {
+			columns = append(columns, column)
+			placeholders = append(placeholders, fmt.Sprintf("$%d", i+1))
+			values = append(values, v.Field(i).Interface())
+		}
+	}
+	columnsPart := strings.Join(columns, ", ")
+	placeholdersPart := strings.Join(placeholders, ", ")
+	sql := fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s)", table, columnsPart, placeholdersPart)
+	return sql, values
+}

seed/postgres/config.go

@@ -0,0 +1,25 @@
+package postgres
+
+import "database/sql"
+
+// SeedConfig represents the configuration for the Postgres seed component.
+type SeedConfig struct {
+	// ClientProvider - Provides a postgres client to use.
+	// available only via code, not available in config files.
+	ClientProvider func() (*sql.DB, error) `json:"-"`
+
+	// Setup - a string that contains the SQL setup script to run before seeding the data.
+	Setup string `json:"-"`
+
+	// Data - a list of objects, each represents a single postgres table and its data
+	Data []*SeedTableData `json:"data,omitempty"`
+}
+
+// SeedTableData represents data for a Postgres table.
+type SeedTableData struct {
+	// TableName - the name of the target postgres table
+	TableName string `json:"table,omitempty"`
+
+	// Rows - a list of rows to insert using the postgres Exec function (a `column` tag is required for each field):
+	Rows []any `json:"rows,omitempty"`
+}