Skip to content

Commit

Permalink
Postgres seed (#21)
Browse files Browse the repository at this point in the history
* Postgres seed

* Update CHANGELOG.md

* rename

* PR fixes
  • Loading branch information
nirapx authored Jul 16, 2024
1 parent a05fb30 commit a70935a
Show file tree
Hide file tree
Showing 3 changed files with 175 additions and 0 deletions.
6 changes: 6 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
144 changes: 144 additions & 0 deletions seed/postgres/component.go
Original file line number Diff line number Diff line change
@@ -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
}
25 changes: 25 additions & 0 deletions seed/postgres/config.go
Original file line number Diff line number Diff line change
@@ -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:"setup,omitempty"`

// 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_name,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"`
}

0 comments on commit a70935a

Please sign in to comment.