Skip to content

Commit c07fd43

Browse files
feat: random data for tables missing in yaml (#46)
1 parent cbc9ab7 commit c07fd43

File tree

1 file changed

+34
-2
lines changed

1 file changed

+34
-2
lines changed

custom.go

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111

1212
// Skeleton the main type with captures all the yaml data
1313
type Skeleton struct {
14+
RandomDataForMissingTables bool `yaml: "RandomDataForMissingTables"`
1415
Custom []TableModel `yaml:"Custom"`
1516
}
1617

@@ -277,12 +278,39 @@ func (c *Skeleton) LoadDataByConfiguration() {
277278
db := ConnectDB()
278279
defer db.Close()
279280

281+
if c.RandomDataForMissingTables {
282+
var w []string
283+
for _, s := range c.Custom {
284+
w = append(w, fmt.Sprintf("'%s.%s'", s.Table, s.Schema))
285+
}
286+
whereClause := fmt.Sprintf("AND (n.nspname || '.' || c.relname) NOT IN (%s)", strings.Join(w, ","))
287+
allTables := dbExtractTables(whereClause)
288+
289+
var tablesToAutoMock []DBTables
290+
AllTablesLoop:
291+
for _, t := range allTables {
292+
for _, s := range c.Custom {
293+
if t.Table == s.Table && t.Schema == s.Schema {
294+
continue AllTablesLoop
295+
}
296+
}
297+
Debugf("Table to automatically mock: %s", t.Table)
298+
tablesToAutoMock = append(tablesToAutoMock, t)
299+
}
300+
tableMocker(tablesToAutoMock)
301+
} else {
302+
BackupDDL() // tableMocker does it automatically
303+
}
304+
305+
280306
for _, s := range c.Custom {
281307
// Initialize the mocking process
282308
tab := GenerateTableName(s.Table, s.Schema)
283309
msg := fmt.Sprintf("Mocking Table %s", tab)
284310
bar := StartProgressBar(msg, cmdOptions.Rows)
285-
311+
if c.RandomDataForMissingTables {
312+
RemoveConstraints(tab)
313+
}
286314
// Name the for loop to break when we encounter error
287315
DataTypePickerLoop:
288316
// Loop through the row count and start loading the data
@@ -314,10 +342,14 @@ func (c *Skeleton) LoadDataByConfiguration() {
314342
col = append(col, v.Name)
315343
data = append(data, fmt.Sprintf("%v", d))
316344
}
317-
318345
// Copy the data to the table
319346
CopyData(tab, col, data, db)
320347
_ = bar.Add(1)
321348
}
322349
}
350+
351+
if c.RandomDataForMissingTables && !cmdOptions.IgnoreConstraint {
352+
FixConstraints()
353+
}
354+
323355
}

0 commit comments

Comments
 (0)