From db17409eb13667cd2e24a6e69951c8d22e068fd8 Mon Sep 17 00:00:00 2001
From: Janez Troha <dz0ny@ubuntu.si>
Date: Fri, 11 Aug 2023 21:31:10 +0200
Subject: [PATCH] feat: Refactor CLI code to handle force sync option

- Modified `String` method of `arrayForce` to join strings with a comma
- Added import statements for packages "os" and "github.com/rs/zerolog"
- Initialized logger with ConsoleWriter for standard error output
- Added command line flag for "force" option
- Logged message with tables to force sync if length of "forceSync" array is greater than 0
---
 README.md    | 17 +++++++++++++----
 cli/force.go | 27 +++++++++++++++++++++++++++
 cli/main.go  | 11 ++++++++++-
 flake.nix    |  1 -
 4 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 cli/force.go

diff --git a/README.md b/README.md
index 8751565..d67965e 100644
--- a/README.md
+++ b/README.md
@@ -23,17 +23,26 @@ Utilizing the native PostgreSQL COPY command, pg-subsetter performs data transfe
 ```
 Usage of subsetter:
   -dst string
-        Destination DSN
+        Destination database DSN
   -f float
         Fraction of rows to copy (default 0.05)
+  -force value
+        Query to copy required tables (users: id = 1)
   -src string
-        Source DSN
+        Source database DSN
 ```
 
 
-Example:
+### Example
 
-```pg-subsetter -src postgresql://:@/bigdb -dst postgresql://:@/littledb -f 0.05```
+```
+pg-subsetter 
+      -src postgresql://:@/bigdb 
+      -dst postgresql://:@/littledb 
+      -f 0.05 
+      -force "users: id = 1" 
+      -force "groups: id = 12"
+```
 
 # Installing
 
diff --git a/cli/force.go b/cli/force.go
new file mode 100644
index 0000000..d1ee7fc
--- /dev/null
+++ b/cli/force.go
@@ -0,0 +1,27 @@
+package main
+
+import (
+	"fmt"
+	"strings"
+)
+
+type Force struct {
+	Table string
+	Where string
+}
+
+type arrayForce []Force
+
+func (i *arrayForce) String() string {
+	return fmt.Sprintf("%v", *i)
+}
+
+func (i *arrayForce) Set(value string) error {
+	q := strings.SplitAfter(strings.TrimSpace(value), ":")
+
+	*i = append(*i, Force{
+		Table: q[0],
+		Where: q[1],
+	})
+	return nil
+}
diff --git a/cli/main.go b/cli/main.go
index dd07c59..08fdcc3 100644
--- a/cli/main.go
+++ b/cli/main.go
@@ -2,17 +2,22 @@ package main
 
 import (
 	"flag"
+	"os"
 
+	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 )
 
 var src = flag.String("src", "", "Source database DSN")
 var dst = flag.String("dst", "", "Destination database DSN")
 var fraction = flag.Float64("f", 0.05, "Fraction of rows to copy")
+var forceSync arrayForce
 
 func main() {
+	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
+
+	flag.Var(&forceSync, "force", "Query to copy required tables (users: id = 1)")
 	flag.Parse()
-	log.Info().Msg("Starting")
 
 	if *src == "" || *dst == "" {
 		log.Fatal().Msg("Source and destination DSNs are required")
@@ -22,4 +27,8 @@ func main() {
 		log.Fatal().Msg("Fraction must be between 0 and 1")
 	}
 
+	if len(forceSync) > 0 {
+		log.Info().Msg("Forcing sync for tables: " + forceSync.String())
+	}
+
 }
diff --git a/flake.nix b/flake.nix
index 29396f0..6367fbc 100644
--- a/flake.nix
+++ b/flake.nix
@@ -37,7 +37,6 @@
               golangci-lint
               postgresql
               process-compose
-              shellcheck
               nixpkgs-fmt
               pgweb
             ];