Add support for plugins

.github/workflows/lint.yml

@@ -18,6 +18,6 @@ jobs:
         with:
           go-version: stable
       - name: golangci-lint
-        uses: golangci/golangci-lint-action@v6
+        uses: golangci/golangci-lint-action@v7
         with:
-          version: v1.64.5
+          version: v2.11.3

README.md

@@ -125,6 +125,77 @@ The latter format is especially useful if you have generated columns on every ta
 
 In the future, additional flags may be added to allow you to include tables, add arbitrary `WHERE` conditions, modify the row id/key, export multiple files, or use existing templates.
 
+## Plugins
+
+If you would like to implement your own `valueFuncs`, you can do so by writing a ripoff plugin.
+
+Plugins are local unauthenticated TCP servers that consume and emit newline-separated JSON messages from ripoff.
+
+### Writing a plugin
+
+Plugins must listen to a local TCP port and provide a TCP stream (loop of receiving and sending messages) to clients.
+
+On startup, plugins must output the string `READY` in its first line of output to indicate to ripoff that it is ready to receive TCP messges.
+
+Each incoming message will be a single line of JSON in the following types:
+
+#### Return a value
+
+Your plugin must process an arbitrary `valueFunc` and return a string value. You can decide how to handle functions you do not expect/provide, by either returning an empty value or disconnecting the client.
+
+Message from ripoff:
+
+```json
+{"type": "valueFunc", "valueFunc": "someFuncName", "args": ["some", "argument", "list"]}
+```
+
+Response from your TCP server:
+
+```json
+{"value": "someString"}
+```
+
+#### Exit your process
+
+Ripoff will send a kill signal to your process, but if you'd like to clean up before that an exit message will be sent beforehand.
+
+Request message:
+
+```json
+{"type": "exit"}
+```
+
+#### Example
+
+An example plugin can be found at `cmd/helloplugin/helloplugin.go`. although TCP servers in other languages may be much easier to implement.
+
+### Using a plugin
+
+Plugins are defined in your ripoff files, which instruct ripoff to spawn a process to start your TCP server, then later connect to it with a single TCP stream. Here's an example from ripoff's tests:
+
+```yml
+# A list of plugins to register with ripoff.
+plugins:
+  # An arbitrary name for the plugin. Used only to handle merging ripoff files
+  # that may define duplicate plugins, which would otherwise conflict.
+  helloplugin:
+    # An arbitrary command to execute, which should start your TCP server.
+    # The command must output READY in its first line of stdout.
+    command: [go, run, cmd/helloplugin/helloplugin.go]
+    # A TCP address to connect to after your command is ready. Note that a
+    # single connection is used to avoid a handshake per valueFunc call.
+    address: localhost:6767
+    # The list of valueFuncs this plugin provides. If ripoff encounters these
+    # it will call out to your plugin. Note that these take precedence over
+    # built in valueFuncs, so you can override ripoff's defaults (like uuid()).
+    valueFuncs: [sayHello]
+rows:
+  users:uuid(fooBar):
+    # In your ripoff files, you can now call your plugin's registered
+    # valueFuncs the same as any other valueFunc.
+    name: sayHello(World)
+```
+
 ## Installation
 
 1. Run `go install github.com/mortenson/ripoff/cmd/ripoff-export@latest`

cmd/helloplugin/helloplugin.go

@@ -0,0 +1,90 @@
+package main
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"log"
+	"net"
+	"os"
+)
+
+func main() {
+	listener, err := net.Listen("tcp", ":6767")
+	if err != nil {
+		log.Fatal("Error listening:", err)
+	}
+	defer func() {
+		err := listener.Close()
+		if err != nil {
+			log.Println("Error closing listener", err)
+		}
+	}()
+	fmt.Println("READY")
+
+	for {
+		conn, err := listener.Accept()
+		if err != nil {
+			log.Println("Error accepting connection:", err)
+			continue
+		}
+		go handleConnection(conn)
+	}
+}
+
+type Request struct {
+	Type      string   `json:"type"`
+	ValueFunc string   `json:"valueFunc"`
+	Args      []string `json:"args"`
+}
+
+type Response struct {
+	Value string `json:"value"`
+}
+
+func handleConnection(conn net.Conn) {
+	defer func() {
+		err := conn.Close()
+		if err != nil {
+			log.Println("Error closing connection", err)
+		}
+	}()
+
+	scanner := bufio.NewScanner(conn)
+	for scanner.Scan() {
+		message := scanner.Bytes()
+		r := Request{}
+		err := json.Unmarshal(message, &r)
+		if err != nil {
+			log.Println("Error parsing body:", err)
+			return
+		}
+		if r.Type == "exit" {
+			os.Exit(0)
+			return
+		}
+		if r.ValueFunc != "sayHello" {
+			log.Println("Unknown value func:", r.ValueFunc)
+			return
+		}
+		if len(r.Args) == 0 {
+			log.Println("No args provided")
+			return
+		}
+		resp, err := json.Marshal(Response{
+			Value: fmt.Sprintf("Hello %s", r.Args[0]),
+		})
+		if err != nil {
+			log.Println("Could not marshal message:", r)
+			return
+		}
+		_, err = conn.Write(append(resp, '\n'))
+		if err != nil {
+			log.Println("Could not send message:", r.ValueFunc)
+			return
+		}
+	}
+	if err := scanner.Err(); err != nil {
+		log.Println("Scanner error:", err)
+	}
+}

cmd/ripoff-export/ripoff_export.go

@@ -49,7 +49,12 @@ func main() {
 		slog.Error("Could not connect to database", errAttr(err))
 		os.Exit(1)
 	}
-	defer conn.Close(ctx)
+	defer func() {
+		err := conn.Close(ctx)
+		if err != nil {
+			slog.Error("Could not close database connection", errAttr(err))
+		}
+	}()
 
 	exportDirectory := path.Clean(args[0])
 	dirInfo, err := os.Stat(exportDirectory)

cmd/ripoff/ripoff.go

@@ -44,7 +44,12 @@ func main() {
 		slog.Error("Could not connect to database", errAttr(err))
 		os.Exit(1)
 	}
-	defer conn.Close(ctx)
+	defer func() {
+		err := conn.Close(ctx)
+		if err != nil {
+			slog.Error("Could not close database connection", errAttr(err))
+		}
+	}()
 
 	tx, err := conn.Begin(ctx)
 	if err != nil {

db.go

@@ -21,12 +21,18 @@ import (
 
 // Runs ripoff from start to finish, without committing the transaction.
 func RunRipoff(ctx context.Context, tx pgx.Tx, totalRipoff RipoffFile) error {
+	manager, err := NewPluginManager(totalRipoff.Plugins)
+	if err != nil {
+		return err
+	}
+	defer manager.Close()
+
 	primaryKeys, err := getPrimaryKeys(ctx, tx)
 	if err != nil {
 		return err
 	}
 
-	queries, err := buildQueriesForRipoff(primaryKeys, totalRipoff)
+	queries, err := buildQueriesForRipoff(manager, primaryKeys, totalRipoff)
 	if err != nil {
 		return err
 	}
@@ -111,14 +117,18 @@ func GetEnumValues(ctx context.Context, tx pgx.Tx) (EnumValuesResult, error) {
 var valueFuncRegex = regexp.MustCompile(`([a-zA-Z]+)\((.*)\)$`)
 var referenceRegex = regexp.MustCompile(`^[a-zA-Z0-9_]+:[a-zA-Z]+\(`)
 
-func prepareValue(rawValue string) (string, error) {
+func prepareValue(manager *PluginManager, rawValue string) (string, error) {
 	valueFuncMatches := valueFuncRegex.FindStringSubmatch(rawValue)
 	if len(valueFuncMatches) != 3 {
 		return rawValue, nil
 	}
 	methodName := valueFuncMatches[1]
 	value := valueFuncMatches[2]
-	valueParts := strings.Split(strings.ReplaceAll(" ", "", valueFuncMatches[2]), ",")
+	valueParts := strings.Split(strings.ReplaceAll(valueFuncMatches[2], " ", ""), ",")
+
+	if manager.Supports(methodName) {
+		return manager.Call(methodName, valueParts...)
+	}
 
 	// Create a new random seed based on a sha256 hash of the value.
 	h := sha256.New()
@@ -153,7 +163,7 @@ func prepareValue(rawValue string) (string, error) {
 	return fakerResult, nil
 }
 
-func buildQueryForRow(primaryKeys PrimaryKeysResult, rowId string, row Row, dependencyGraph map[string][]string) (string, error) {
+func buildQueryForRow(manager *PluginManager, primaryKeys PrimaryKeysResult, rowId string, row Row, dependencyGraph map[string][]string) (string, error) {
 	parts := strings.Split(rowId, ":")
 	if len(parts) < 2 {
 		return "", fmt.Errorf("invalid id: %s", rowId)
@@ -225,7 +235,7 @@ func buildQueryForRow(primaryKeys PrimaryKeysResult, rowId string, row Row, depe
 			}
 
 			columns = append(columns, pq.QuoteIdentifier(column))
-			valuePrepared, err := prepareValue(value)
+			valuePrepared, err := prepareValue(manager, value)
 			if err != nil {
 				return "", err
 			}
@@ -257,7 +267,7 @@ func buildQueryForRow(primaryKeys PrimaryKeysResult, rowId string, row Row, depe
 }
 
 // Returns a sorted array of queries to run based on a given ripoff file.
-func buildQueriesForRipoff(primaryKeys PrimaryKeysResult, totalRipoff RipoffFile) ([]string, error) {
+func buildQueriesForRipoff(manager *PluginManager, primaryKeys PrimaryKeysResult, totalRipoff RipoffFile) ([]string, error) {
 	dependencyGraph := map[string][]string{}
 	queries := map[string]string{}
 
@@ -268,7 +278,7 @@ func buildQueriesForRipoff(primaryKeys PrimaryKeysResult, totalRipoff RipoffFile
 
 	// Build queries.
 	for rowId, row := range totalRipoff.Rows {
-		query, err := buildQueryForRow(primaryKeys, rowId, row, dependencyGraph)
+		query, err := buildQueryForRow(manager, primaryKeys, rowId, row, dependencyGraph)
 		if err != nil {
 			return []string{}, err
 		}
@@ -400,7 +410,7 @@ func topologicalSort(digraph map[string][]string) ([]string, error) {
 	visit = func(u string) {
 		if temporaryMark[u] {
 			acyclic = false
-		} else if !(temporaryMark[u] || permanentMark[u]) {
+		} else if !temporaryMark[u] && !permanentMark[u] {
 			temporaryMark[u] = true
 			for _, v := range digraph[u] {
 				visit(v)

db_test.go

@@ -47,11 +47,11 @@ func runTestData(t *testing.T, ctx context.Context, tx pgx.Tx, testDir string) {
 	validationFile, err := os.ReadFile(path.Join(testDir, "validate.sql"))
 	if err == nil {
 		row := tx.QueryRow(ctx, string(validationFile))
-		var success int
+		var success bool
 		var debug string
 		err := row.Scan(&success, &debug)
 		require.NoError(t, err)
-		if success != 1 {
+		if !success {
 			t.Fatalf("Validation failed with debug content: %s", debug)
 		}
 	}
@@ -64,10 +64,11 @@ func TestRipoff(t *testing.T) {
 	}
 	ctx := context.Background()
 	conn, err := pgx.Connect(ctx, envUrl)
-	if err != nil {
+	require.NoError(t, err)
+	defer func() {
+		err := conn.Close(ctx)
 		require.NoError(t, err)
-	}
-	defer conn.Close(ctx)
+	}()
 
 	_, filename, _, _ := runtime.Caller(0)
 	dir := path.Join(path.Dir(filename), "testdata", "import")

export_test.go

@@ -38,6 +38,8 @@ func runExportTestData(t *testing.T, ctx context.Context, tx pgx.Tx, testDir str
 	expectedRipoffFile := &RipoffFile{}
 	err = yaml.Unmarshal(expectedRipoffYaml, expectedRipoffFile)
 	require.NoError(t, err)
+	expectedRipoffFile.Plugins = nil
+	newRipoffFile.Plugins = nil
 	require.Equal(t, expectedRipoffFile, newRipoffFile)
 
 	// Wipe database.
@@ -72,10 +74,11 @@ func TestRipoffExport(t *testing.T) {
 	}
 	ctx := context.Background()
 	conn, err := pgx.Connect(ctx, envUrl)
-	if err != nil {
+	require.NoError(t, err)
+	defer func() {
+		err := conn.Close(ctx)
 		require.NoError(t, err)
-	}
-	defer conn.Close(ctx)
+	}()
 
 	_, filename, _, _ := runtime.Caller(0)
 	dir := path.Join(path.Dir(filename), "testdata", "export")
@@ -102,10 +105,11 @@ func TestExcludeFlag(t *testing.T) {
 	}
 	ctx := context.Background()
 	conn, err := pgx.Connect(ctx, envUrl)
-	if err != nil {
+	require.NoError(t, err)
+	defer func() {
+		err := conn.Close(ctx)
 		require.NoError(t, err)
-	}
-	defer conn.Close(ctx)
+	}()
 
 	// Start a transaction that we'll roll back at the end
 	tx, err := conn.Begin(ctx)
@@ -265,10 +269,11 @@ func TestExcludeColumnsFlag(t *testing.T) {
 	}
 	ctx := context.Background()
 	conn, err := pgx.Connect(ctx, envUrl)
-	if err != nil {
+	require.NoError(t, err)
+	defer func() {
+		err := conn.Close(ctx)
 		require.NoError(t, err)
-	}
-	defer conn.Close(ctx)
+	}()
 
 	// Start a transaction that we'll roll back at the end
 	tx, err := conn.Begin(ctx)

go.mod

@@ -1,8 +1,8 @@
 module github.com/mortenson/ripoff
 
-go 1.22
+go 1.26
 
-toolchain go1.22.4
+toolchain go1.26.0
 
 require (
 	github.com/brianvoe/gofakeit/v7 v7.0.4