<commit_before>package migrator

import "github.com/semrekkers/sugar"

const (
	Initialize = ""
)

type Model interface {
	TableName() string
	Migrate(*sugar.Querier, string)
}

type DBInfo interface {
	HasTable(sugar.Executor, string) (bool, error)
	TableColumns(sugar.Executor, string) ([]string, error)
}

type Migrator struct {
	ex      sugar.Executor
	bindVar sugar.Formatter
	mapper  sugar.TypeMapper
	dbInfo  DBInfo
}

type Result struct {
	Tables, Columns []string
}

func New(ex sugar.Executor, bindVar sugar.Formatter, mapper sugar.TypeMapper, dbInfo DBInfo) *Migrator {
	return &Migrator{ex, bindVar, mapper, dbInfo}
}

func NewFromDB(db *sugar.DB, dbInfo DBInfo) *Migrator {
	return &Migrator{db.DB, db.BindVar, db.Mapper, dbInfo}
}

func (m *Migrator) Migrate(models ...Model) (*Result, error) {
	var result Result

	for _, model := range models {
		tableName := model.TableName()

		tableExists, newColumns, err := m.migrateModel(tableName, model)
		if err != nil {
			return nil, err
		}
		if tableExists {
			result.Tables = append(result.Tables, tableName)
		} else {
			result.Columns = append(result.Columns, newColumns...)
		}
	}

	return &result, nil
}

func (m *Migrator) migrateModel(tableName string, model Model) (tableExists bool, newColumns []string, err error) {
	tableExists, err = m.dbInfo.HasTable(m.ex, tableName)
	if err != nil {
		return
	}

	q := sugar.NewQuerier(m.ex, sugar.DefaultBindVar)
	fieldSelector := sugar.Fields(model).SetTypeMapper(m.mapper)
	if !tableExists {
		q.Writef("CREATE TABLE %s (", tableName)
		q.FieldDefinitions("", fieldSelector.Select())
		model.Migrate(q, Initialize)
		q.WriteString(")")
	} else {
		var columns []string
		columns, err = m.dbInfo.TableColumns(m.ex, tableName)
		if err != nil {
			return
		}
		fields := fieldSelector.Except(columns...).Select()
		if len(fields) == 0 {
			return
		}

		q.Writef("ALTER TABLE %s (", tableName)
		q.FieldDefinitions("ADD ", fields)
		for _, field := range fields {
			model.Migrate(q, field.Name)
			newColumns = append(newColumns, tableName+"."+field.Name)
		}
		q.WriteString(")")
	}
	err = q.Exec()

	return
}
<commit_msg>Update migrator<commit_after>package migrator

import "github.com/semrekkers/sugar"

const (
	Initialize = ""
)

type Model interface {
	TableName() string
	Migrate(*sugar.Querier, string)
}

type DBInfo interface {
	HasTable(sugar.Executor, string) (bool, error)
	TableColumns(sugar.Executor, string) ([]string, error)
}

type Migrator struct {
	ex      sugar.Executor
	bindVar sugar.Formatter
	mapper  sugar.TypeMapper
	dbInfo  DBInfo
}

type Result struct {
	Tables, Columns []string
}

func New(ex sugar.Executor, bindVar sugar.Formatter, mapper sugar.TypeMapper, dbInfo DBInfo) *Migrator {
	return &Migrator{ex, bindVar, mapper, dbInfo}
}

func NewFromDB(db *sugar.DB, dbInfo DBInfo) *Migrator {
	return &Migrator{db.DB, db.BindVar, db.Mapper, dbInfo}
}

func (m *Migrator) Migrate(models ...Model) (*Result, error) {
	var result Result

	for _, model := range models {
		tableName := model.TableName()

		tableExists, newColumns, err := m.migrateModel(tableName, model)
		if err != nil {
			return nil, err
		}
		if !tableExists {
			result.Tables = append(result.Tables, tableName)
		} else {
			result.Columns = append(result.Columns, newColumns...)
		}
	}

	return &result, nil
}

func (m *Migrator) migrateModel(tableName string, model Model) (tableExists bool, newColumns []string, err error) {
	tableExists, err = m.dbInfo.HasTable(m.ex, tableName)
	if err != nil {
		return
	}

	q := sugar.NewQuerier(m.ex, sugar.DefaultBindVar)
	fieldSelector := sugar.Fields(model).SetTypeMapper(m.mapper)
	if !tableExists {
		q.Writef("CREATE TABLE %s (", tableName)
		q.FieldDefinitions("", fieldSelector.Select())
		model.Migrate(q, Initialize)
		q.WriteString(")")
	} else {
		var columns []string
		columns, err = m.dbInfo.TableColumns(m.ex, tableName)
		if err != nil {
			return
		}
		fields := fieldSelector.Except(columns...).Select()
		if len(fields) == 0 {
			return
		}

		q.Writef("ALTER TABLE %s ", tableName)
		q.FieldDefinitions("ADD ", fields)
		for _, field := range fields {
			model.Migrate(q, field.Name)
			newColumns = append(newColumns, tableName+"."+field.Name)
		}
	}
	err = q.Exec()

	return
}

{'fraction_non_alphanumeric': 0.13577636833450787, 'fraction_numerical': 0.0004698144233027954, 'mean_word_length': 4.105515587529976, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4079050', 'n_tokens_mistral': 1597, 'n_tokens_neox': 1505, 'n_words': 360}