RDS provider uses writer for dump output (#42)

nickschuch · web-flow · commit 4d8e577b5dc1 · 2025-09-08T19:26:01.000+10:00
* RDS provider uses writer for dump output

* Linting
diff --git a/internal/mysql/provider/provider.go b/internal/mysql/provider/provider.go
@@ -1,10 +1,13 @@
 package provider
 
-import "context"
+import (
+	"context"
+	"io"
+)
 
 // Interface implements the required functionality for a Provider.
 type Interface interface {
-	GetSelectQueryForTable(ctx context.Context, table string, params DumpParams) (string, error)
+	WriteTableData(ctx context.Context, w io.Writer, table string, params DumpParams) error
 }
 
 // DumpParams is used to pass parameters to the Dump function.
diff --git a/internal/mysql/provider/rds/provider.go b/internal/mysql/provider/rds/provider.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"io"
 	"log"
 	"strings"
 
@@ -31,11 +32,30 @@ func NewClient(conn *sql.Conn, logger *log.Logger, region, uri string) *Client {
 	}
 }
 
-// GetSelectQueryForTable will return a complete SELECT query to export data from a table.
-func (d *Client) GetSelectQueryForTable(ctx context.Context, table string, params provider.DumpParams) (string, error) {
+// WriteTableData will export the data from a table to S3 and write the LOAD DATA query to the provided writer.
+func (d *Client) WriteTableData(ctx context.Context, w io.Writer, table string, params provider.DumpParams) error {
+	// Push table data to s3.
+	err := d.exportTableData(ctx, table, params)
+	if err != nil {
+		return err
+	}
+
+	// Write the import query to the writer.
+	err = d.writeLoadQueryForTable(w, table)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Export the data from a table to S3.
+func (d *Client) exportTableData(ctx context.Context, table string, params provider.DumpParams) error {
+	d.Logger.Printf("Exporting data to S3 for table: %s", table)
+
 	cols, err := providerutils.QueryColumnsForTable(ctx, d.Conn, table, params)
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	query := fmt.Sprintf("SELECT %s", strings.Join(cols, ", "))
@@ -50,26 +70,28 @@ func (d *Client) GetSelectQueryForTable(ctx context.Context, table string, param
 	query = fmt.Sprintf("%s MANIFEST ON", query)
 	query = fmt.Sprintf("%s OVERWRITE ON", query)
 
-	importQuery, err := d.GetLoadQueryForTable(table)
+	_, err = d.Conn.QueryContext(ctx, query)
 	if err != nil {
-		return "", err
+		return fmt.Errorf("error exporting data to S3 for table %s: %w", table, err)
 	}
 
-	fmt.Println(importQuery)
-	return query, nil
+	return nil
 }
 
-// GetLoadQueryForTable will return a complete SELECT query to fetch data from a table.
-func (d *Client) GetLoadQueryForTable(table string) (string, error) {
+// Write a LOAD DATA FROM S3 query to the provided writer.
+func (d *Client) writeLoadQueryForTable(w io.Writer, table string) error {
 	if table == "" {
-		return "", fmt.Errorf("error: no table specified")
+		return fmt.Errorf("error: no table specified")
 	}
+
 	if d.Region == "" || len(strings.Split(d.Region, "-")) != 3 {
-		return "", fmt.Errorf("error: region is not configured correctly")
+		return fmt.Errorf("error: region is not configured correctly")
 	}
+
 	path := strings.TrimPrefix(d.URI, "s3://")
 	query := fmt.Sprintf("LOAD DATA FROM S3 MANIFEST 'S3-%s://%s/%s.csv.manifest' INTO TABLE `%s` CHARACTER SET utf8mb4", d.Region, path, table, table)
 	query = fmt.Sprintf("%s FIELDS TERMINATED BY ',' ENCLOSED BY '\"' LINES TERMINATED BY '\\n'", query)
 
-	return query, nil
+	_, err := fmt.Fprintln(w, query)
+	return err
 }
diff --git a/internal/mysql/provider/rds/provider_test.go b/internal/mysql/provider/rds/provider_test.go
@@ -1,35 +1,69 @@
 package rds
 
 import (
+	"bytes"
 	"context"
+	"io"
 	"log"
-	"os"
+	"regexp"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
-	"github.com/stretchr/testify/assert"
-
-	"github.com/skpr/mtk/internal/mysql/mock"
 	"github.com/skpr/mtk/internal/mysql/provider"
 )
 
-func TestMySQLGetExportSelectQueryFor(t *testing.T) {
-	db, mock := mock.GetDB(t)
-	dumper := NewClient(db, log.New(os.Stdout, "", 0), "ap-southheast-2", "s3://path/to/bucket")
-	mock.ExpectQuery("SELECT \\* FROM `table` LIMIT 1").WillReturnRows(
-		sqlmock.NewRows([]string{"c1", "c2"}).AddRow("a", "b"))
-	query, err := dumper.GetSelectQueryForTable(context.TODO(), "table", provider.DumpParams{
-		SelectMap: map[string]map[string]string{"table": {"c2": "NOW()"}},
-		WhereMap:  map[string]string{"table": "c1 > 0"},
-	})
-	assert.Nil(t, err)
-	assert.Equal(t, "SELECT `c1`, NOW() AS `c2` FROM `table` WHERE c1 > 0 INTO OUTFILE S3 's3://path/to/bucket/table.csv' FIELDS TERMINATED BY ',' ENCLOSED BY '\"' LINES TERMINATED BY '\\n' MANIFEST ON OVERWRITE ON", query)
-}
+func TestWriteTableData(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New error: %v", err)
+	}
+	defer db.Close()
+
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		t.Fatalf("db.Conn error: %v", err)
+	}
+
+	// 1) QueryColumnsForTable runs: expect the probe query and return columns.
+	mock.
+		ExpectQuery(regexp.QuoteMeta("SELECT * FROM `users` LIMIT 1")).
+		WillReturnRows(sqlmock.NewRows([]string{"id", "name"})) // Columns only
+
+	// 2) exportTableData runs: expect the exact OUTFILE S3 query with backticked cols.
+	exportSQL := "SELECT `id`, `name` FROM `users` INTO OUTFILE S3 's3://my-bucket/prefix/users.csv' " +
+		"FIELDS TERMINATED BY ',' ENCLOSED BY '\"' LINES TERMINATED BY '\\n' MANIFEST ON OVERWRITE ON"
+	mock.
+		ExpectQuery(regexp.QuoteMeta(exportSQL)).
+		WillReturnRows(sqlmock.NewRows([]string{}))
+
+	// Client + buffer for the LOAD DATA query
+	buf := &bytes.Buffer{}
+	c := &Client{
+		Conn:   conn,
+		Logger: log.New(io.Discard, "", 0),
+		Region: "ap-southeast-2",
+		URI:    "s3://my-bucket/prefix",
+	}
+
+	// Run the full flow
+	if err := c.WriteTableData(context.Background(), buf, "users", provider.DumpParams{
+		WhereMap:  map[string]string{},
+		SelectMap: map[string]map[string]string{},
+	}); err != nil {
+		t.Fatalf("WriteTableData error: %v", err)
+	}
+
+	// Assert the LOAD DATA query is exactly written (including newline)
+	got := buf.String()
+	want := "LOAD DATA FROM S3 MANIFEST 'S3-ap-southeast-2://my-bucket/prefix/users.csv.manifest' " +
+		"INTO TABLE `users` CHARACTER SET utf8mb4 FIELDS TERMINATED BY ',' ENCLOSED BY '\"' " +
+		"LINES TERMINATED BY '\\n'\n"
+	if got != want {
+		t.Fatalf("unexpected load query:\n--- got ---\n%s\n--- want ---\n%s", got, want)
+	}
 
-func TestMySQLGetLoadQueryFor(t *testing.T) {
-	db, _ := mock.GetDB(t)
-	dumper := NewClient(db, log.New(os.Stdout, "", 0), "ap-southeast-4", "s3://path/to/bucket")
-	query, err := dumper.GetLoadQueryForTable("table_name")
-	assert.Nil(t, err)
-	assert.Equal(t, "LOAD DATA FROM S3 MANIFEST 'S3-ap-southeast-4://path/to/bucket/table_name.csv.manifest' INTO TABLE `table_name` CHARACTER SET utf8mb4 FIELDS TERMINATED BY ',' ENCLOSED BY '\"' LINES TERMINATED BY '\\n'", query)
+	// Ensure all expectations met
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("unmet sqlmock expectations: %v", err)
+	}
 }
diff --git a/internal/mysql/provider/stdout/provider.go b/internal/mysql/provider/stdout/provider.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"io"
 	"log"
 	"strings"
 
@@ -26,11 +27,11 @@ func NewClient(conn *sql.Conn, logger *log.Logger) *Client {
 	}
 }
 
-// GetSelectQueryForTable will return a complete SELECT query to fetch data from a table.
-func (d *Client) GetSelectQueryForTable(ctx context.Context, table string, params provider.DumpParams) (string, error) {
+// WriteTableData will write the data from a table to the provided writer.
+func (d *Client) WriteTableData(ctx context.Context, w io.Writer, table string, params provider.DumpParams) error {
 	cols, err := providerutils.QueryColumnsForTable(ctx, d.Conn, table, params)
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	query := fmt.Sprintf("SELECT %s FROM `%s`", strings.Join(cols, ", "), table)
@@ -39,5 +40,72 @@ func (d *Client) GetSelectQueryForTable(ctx context.Context, table string, param
 		query = fmt.Sprintf("%s WHERE %s", query, where)
 	}
 
-	return query, nil
+	rows, err := d.Conn.QueryContext(ctx, query)
+	if err != nil {
+		return err
+	}
+
+	columns, err := rows.Columns()
+	if err != nil {
+		return err
+	}
+
+	defer rows.Close()
+
+	values := make([]*sql.RawBytes, len(columns))
+	scanArgs := make([]interface{}, len(values))
+
+	for i := range values {
+		scanArgs[i] = &values[i]
+	}
+
+	var (
+		counter  = 0
+		firstRun = true
+	)
+
+	for rows.Next() {
+		// We have already done a loop and need to close the previous insert statement.
+		if counter >= params.ExtendedInsertRows {
+			fmt.Fprintln(w, ";")
+			counter = 0
+		} else {
+			if !firstRun {
+				fmt.Fprint(w, ",")
+			}
+		}
+
+		if counter == 0 {
+			fmt.Fprintf(w, "INSERT INTO `%s` VALUES ", table)
+		}
+
+		counter++
+
+		firstRun = false
+
+		if err = rows.Scan(scanArgs...); err != nil {
+			return err
+		}
+
+		var vals []string
+
+		for _, col := range values {
+			val := "NULL"
+
+			if col != nil {
+				val, err = getValue(string(*col))
+				if err != nil {
+					return err
+				}
+			}
+
+			vals = append(vals, val)
+		}
+
+		fmt.Fprintf(w, "(%s)", strings.Join(vals, ","))
+	}
+
+	fmt.Fprintln(w, ";")
+
+	return nil
 }
diff --git a/internal/mysql/provider/stdout/provider_test.go b/internal/mysql/provider/stdout/provider_test.go
@@ -1,41 +1,72 @@
 package stdout
 
 import (
+	"bytes"
 	"context"
-	"errors"
+	"io"
 	"log"
-	"os"
+	"regexp"
 	"testing"
 
 	"github.com/DATA-DOG/go-sqlmock"
-	"github.com/stretchr/testify/assert"
-
-	"github.com/skpr/mtk/internal/mysql/mock"
 	"github.com/skpr/mtk/internal/mysql/provider"
 )
 
-func TestMySQLGetSelectQueryFor(t *testing.T) {
-	db, mock := mock.GetDB(t)
-	dumper := NewClient(db, log.New(os.Stdout, "", 0))
-	mock.ExpectQuery("SELECT \\* FROM `table` LIMIT 1").WillReturnRows(
-		sqlmock.NewRows([]string{"c1", "c2"}).AddRow("a", "b"))
-	query, err := dumper.GetSelectQueryForTable(context.TODO(), "table", provider.DumpParams{
-		SelectMap: map[string]map[string]string{"table": {"c2": "NOW()"}},
-		WhereMap:  map[string]string{"table": "c1 > 0"},
-	})
-	assert.Nil(t, err)
-	assert.Equal(t, "SELECT `c1`, NOW() AS `c2` FROM `table` WHERE c1 > 0", query)
-}
+func TestWriteTableData(t *testing.T) {
+	db, mock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New error: %v", err)
+	}
+	defer db.Close()
+
+	conn, err := db.Conn(context.Background())
+	if err != nil {
+		t.Fatalf("db.Conn error: %v", err)
+	}
+
+	// Probe from utils.QueryColumnsForTable
+	mock.ExpectQuery(regexp.QuoteMeta("SELECT * FROM `users` LIMIT 1")).
+		WillReturnRows(sqlmock.NewRows([]string{"id", "name"}))
+
+	// Main SELECT that GetSelectQueryForTable issues
+	mainSelect := "SELECT `id`, `name` FROM `users`"
+	rows := sqlmock.NewRows([]string{"id", "name"}).
+		AddRow(1, "Alice").
+		AddRow(2, "Bob").
+		AddRow(3, "Charlie").
+		AddRow(4, "Dana").
+		AddRow(5, "Eve")
+	mock.ExpectQuery(regexp.QuoteMeta(mainSelect)).WillReturnRows(rows)
+
+	var out bytes.Buffer
+	c := &Client{
+		Conn:   conn,
+		Logger: log.New(io.Discard, "", 0),
+	}
+
+	params := provider.DumpParams{
+		ExtendedInsertRows: 2,
+		WhereMap:           map[string]string{},
+		SelectMap:          map[string]map[string]string{},
+	}
+
+	if err := c.WriteTableData(context.Background(), &out, "users", params); err != nil {
+		t.Fatalf("WriteTableData returned error: %v", err)
+	}
+
+	got := out.String()
+
+	// Expect two batches of 2 and one batch of 1 (with final semicolon inline).
+	want :=
+		"INSERT INTO `users` VALUES (1,'Alice'),(2,'Bob');\n" +
+			"INSERT INTO `users` VALUES (3,'Charlie'),(4,'Dana');\n" +
+			"INSERT INTO `users` VALUES (5,'Eve');\n"
+
+	if got != want {
+		t.Fatalf("unexpected writer output:\n--- got ---\n%s--- want ---\n%s", got, want)
+	}
 
-func TestMySQLGetSelectQueryForHandlingError(t *testing.T) {
-	db, mock := mock.GetDB(t)
-	dumper := NewClient(db, log.New(os.Stdout, "", 0))
-	e := errors.New("broken")
-	mock.ExpectQuery("SELECT \\* FROM `table` LIMIT 1").WillReturnError(e)
-	query, err := dumper.GetSelectQueryForTable(context.TODO(), "table", provider.DumpParams{
-		SelectMap: map[string]map[string]string{"table": {"c2": "NOW()"}},
-		WhereMap:  map[string]string{"table": "c1 > 0"},
-	})
-	assert.Equal(t, e, err)
-	assert.Equal(t, "", query)
+	if err := mock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("unmet sqlmock expectations: %v", err)
+	}
 }
diff --git a/internal/mysql/provider/stdout/utils.go b/internal/mysql/provider/stdout/utils.go
@@ -1,4 +1,4 @@
-package mysql
+package stdout
 
 import (
 	"bytes"
diff --git a/internal/mysql/provider/stdout/utils_test.go b/internal/mysql/provider/stdout/utils_test.go
@@ -1,4 +1,4 @@
-package mysql
+package stdout
 
 import (
 	"testing"
diff --git a/internal/mysql/tables.go b/internal/mysql/tables.go
diff --git a/internal/mysql/write.go b/internal/mysql/write.go

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-package mysql`
	`1`	`+package stdout`
`2`	`2`
`3`	`3`	`import (`
`4`	`4`	`"bytes"`