diff --git a/README.md b/README.md index 2e537bc..2c38345 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,8 @@ # pg-subsetter +[![lint](https://github.com/teamniteo/pg-subsetter/actions/workflows/lint.yml/badge.svg)](https://github.com/teamniteo/pg-subsetter/actions/workflows/lint.yml)[![build](https://github.com/teamniteo/pg-subsetter/actions/workflows/go.yml/badge.svg)](https://github.com/teamniteo/pg-subsetter/actions/workflows/go.yml) + + `pg-subsetter` is a powerful and efficient tool designed to synchronize a fraction of a PostgreSQL database to another PostgreSQL database on the fly, it does not copy the SCHEMA, this means that your target database has to have schema populated in some other way. ### Database Fraction Synchronization diff --git a/subsetter/info_test.go b/subsetter/info_test.go index 58d45b1..f61746b 100644 --- a/subsetter/info_test.go +++ b/subsetter/info_test.go @@ -18,8 +18,8 @@ func TestGetTargetSet(t *testing.T) { tables []Table want []Table }{ - {"simple", 0.5, []Table{{"simple", 1000}}, []Table{{"simple", 31}}}, - {"simple", 0.5, []Table{{"simple", 10}}, []Table{{"simple", 3}}}, + {"simple", 0.5, []Table{{"simple", 1000, []string{}}}, []Table{{"simple", 31, []string{}}}}, + {"simple", 0.5, []Table{{"simple", 10, []string{}}}, []Table{{"simple", 3, []string{}}}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/subsetter/query.go b/subsetter/query.go index 8054b52..4bd5856 100644 --- a/subsetter/query.go +++ b/subsetter/query.go @@ -9,8 +9,9 @@ import ( ) type Table struct { - Name string - Rows int + Name string + Rows int + Relations []string } func GetTables(conn *pgx.Conn) (tables []string, err error) { @@ -57,7 +58,7 @@ func CopyTableToString(table string, limit int, conn *pgx.Conn) (result string, } func CopyStringToTable(table string, data string, conn *pgx.Conn) (err error) { - q := fmt.Sprintf(`copy %s from stdout`, table) + q := fmt.Sprintf(`copy %s from stdin`, table) var buff bytes.Buffer buff.WriteString(data) if _, err = conn.PgConn().CopyFrom(context.Background(), &buff, q); err != nil { diff --git a/subsetter/query_test.go b/subsetter/query_test.go index dfad7c6..83eb5c5 100644 --- a/subsetter/query_test.go +++ b/subsetter/query_test.go @@ -41,7 +41,7 @@ func TestGetTablesWithRows(t *testing.T) { wantTables []Table wantErr bool }{ - {"With tables", conn, []Table{{"simple", 0}, {"relation", 0}}, false}, + {"With tables", conn, []Table{{"simple", 0, []string{}}, {"relation", 0, []string{}}}, false}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -57,7 +57,7 @@ func TestGetTablesWithRows(t *testing.T) { } } -func TestCopyRowToString(t *testing.T) { +func TestCopyTableToString(t *testing.T) { conn := getTestConnection() populateTests(conn) defer conn.Close(context.Background()) @@ -77,11 +77,11 @@ func TestCopyRowToString(t *testing.T) { t.Run(tt.name, func(t *testing.T) { gotResult, err := CopyTableToString(tt.table, 10, tt.conn) if (err != nil) != tt.wantErr { - t.Errorf("CopyRowToString() error = %v, wantErr %v", err, tt.wantErr) + t.Errorf("CopyTableToString() error = %v, wantErr %v", err, tt.wantErr) return } - if strings.Contains(gotResult, "test") != tt.wantResult { - t.Errorf("CopyRowToString() = %v, want %v", gotResult, tt.wantResult) + if strings.Contains(gotResult, "test") == tt.wantResult { + t.Errorf("CopyTableToString() = %v, want %v", gotResult, tt.wantResult) } }) } @@ -92,7 +92,6 @@ func TestCopyStringToTable(t *testing.T) { populateTests(conn) defer conn.Close(context.Background()) defer clearPopulateTests(conn) - populateTestsWithData(conn, "simple", 10) tests := []struct { name string @@ -103,7 +102,7 @@ func TestCopyStringToTable(t *testing.T) { wantErr bool }{ {"With tables", "simple", "cccc5f58-44d3-4d7a-bf37-a97d4f081a63 test\n", conn, 1, false}, - {"With more tables", "simple", "edcd63fe-303e-4d51-83ea-3fd00740ba2c test4\na170b0f5-3aec-469c-9589-cf25888a72e2 test7", conn, 2, false}, + {"With more tables", "simple", "edcd63fe-303e-4d51-83ea-3fd00740ba2c test4\na170b0f5-3aec-469c-9589-cf25888a72e2 test7", conn, 3, false}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -112,7 +111,8 @@ func TestCopyStringToTable(t *testing.T) { t.Errorf("CopyStringToTable() error = %v, wantErr %v", err, tt.wantErr) return } - if tt.wantResult == insertedRows(tt.table, tt.conn) { + gotInserted := insertedRows(tt.table, tt.conn) + if tt.wantResult != gotInserted { t.Errorf("CopyStringToTable() = %v, want %v", tt.wantResult, tt.wantResult) } @@ -121,11 +121,11 @@ func TestCopyStringToTable(t *testing.T) { } func insertedRows(s string, conn *pgx.Conn) int { - tables, _ := GetTablesWithRows(conn) - for _, table := range tables { - if table.Name == s { - return table.Rows - } + q := "SELECT count(*) FROM " + s + var count int + err := conn.QueryRow(context.Background(), q).Scan(&count) + if err != nil { + panic(err) } - return 0 + return count } diff --git a/subsetter/relations.go b/subsetter/relations.go index 3cb84d4..ea96639 100644 --- a/subsetter/relations.go +++ b/subsetter/relations.go @@ -6,23 +6,30 @@ import ( "github.com/jackc/pgx/v5" ) +type Relation struct { + Table string + Column string +} + // GetRelations returns a list of tables that have a foreign key for particular table. -func GetRelations(table string, conn *pgx.Conn) (relations []string, err error) { +func GetRelations(table string, conn *pgx.Conn) (relations []Relation, err error) { - q := `SELECT tc.table_name AS foreign_table_name - FROM - information_schema.table_constraints AS tc - JOIN information_schema.key_column_usage AS kcu - ON tc.constraint_name = kcu.constraint_name - JOIN information_schema.constraint_column_usage AS ccu - ON ccu.constraint_name = tc.constraint_name - WHERE tc.constraint_type = 'FOREIGN KEY' AND ccu.table_name = $1;` + q := `SELECT + kcu.table_name, + kcu.column_name + FROM + information_schema.table_constraints AS tc + JOIN information_schema.key_column_usage AS kcu ON tc.constraint_name = kcu.constraint_name + JOIN information_schema.constraint_column_usage AS ccu ON ccu.constraint_name = tc.constraint_name + WHERE + tc.constraint_type = 'FOREIGN KEY' + AND ccu.table_name = $1;` rows, err := conn.Query(context.Background(), q, table) for rows.Next() { - var table string - if err := rows.Scan(&table); err == nil { - relations = append(relations, table) + var rel Relation + if err := rows.Scan(&rel.Table, &rel.Column); err == nil { + relations = append(relations, rel) } } rows.Close() diff --git a/subsetter/relations_test.go b/subsetter/relations_test.go index eac50c0..bd42eb0 100644 --- a/subsetter/relations_test.go +++ b/subsetter/relations_test.go @@ -17,9 +17,9 @@ func TestGetRelations(t *testing.T) { name string table string conn *pgx.Conn - wantRelations []string + wantRelations []Relation }{ - {"With relation", "simple", conn, []string{"relation"}}, + {"With relation", "simple", conn, []Relation{{"relation", "simple_id"}}}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/subsetter/sync.go b/subsetter/sync.go new file mode 100644 index 0000000..1ba5a03 --- /dev/null +++ b/subsetter/sync.go @@ -0,0 +1,56 @@ +package subsetter + +import ( + "context" + + "github.com/jackc/pgx/v5" +) + +type Sync struct { + source *pgx.Conn + destination *pgx.Conn + fraction float64 + verbose bool +} + +func NewSync(source string, target string, fraction float64, verbose bool) (*Sync, error) { + src, err := pgx.Connect(context.Background(), source) + if err != nil { + return nil, err + } + dst, err := pgx.Connect(context.Background(), source) + if err != nil { + return nil, err + } + + return &Sync{ + source: src, + destination: dst, + fraction: fraction, + verbose: verbose, + }, nil +} + +func (s *Sync) Sync() (err error) { + var tables []Table + if tables, err = GetTablesWithRows(s.source); err != nil { + return + } + + var subset []Table + if subset = GetTargetSet(s.fraction, tables); err != nil { + return + } + + for _, table := range subset { + var data string + if data, err = CopyTableToString(table.Name, table.Rows, s.source); err != nil { + return + } + if err = CopyStringToTable(table.Name, data, s.destination); err != nil { + return + } + } + + return +}