Fix tests

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

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) {

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 {

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
 }

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()

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) {

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
+}