From 4fbdd487a029d70736cf9e3857738e170ba9458b Mon Sep 17 00:00:00 2001
From: Janez T <hey@dz0ny.dev>
Date: Mon, 11 Nov 2024 10:56:37 +0100
Subject: [PATCH] Fix crash while parsing args

---
 cli/extra.go       | 18 ++++++++++++------
 cli/extra_test.go  |  2 ++
 subsetter/graph.go | 22 ++++++++++++++++++++++
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/cli/extra.go b/cli/extra.go
index e217000..27761ec 100644
--- a/cli/extra.go
+++ b/cli/extra.go
@@ -9,16 +9,22 @@ import (
 
 type arrayExtra []subsetter.Rule
 
-func (i *arrayExtra) String() string {
-	return fmt.Sprintf("%v", *i)
+func (ae *arrayExtra) String() string {
+	return fmt.Sprintf("%v", *ae)
 }
 
-func (i *arrayExtra) Set(value string) error {
+func (ae *arrayExtra) Set(value string) error {
 	q := strings.SplitN(strings.TrimSpace(value), ":", 2)
 
-	*i = append(*i, subsetter.Rule{
-		Table: strings.TrimSpace(q[0]),
-		Where: maybeAll(strings.TrimSpace(q[1])),
+	table := strings.TrimSpace(q[0])
+	where := ""
+	if len(q) > 1 {
+		where = strings.TrimSpace(q[1])
+	}
+
+	*ae = append(*ae, subsetter.Rule{
+		Table: table,
+		Where: maybeAll(where),
 	})
 	return nil
 }
diff --git a/cli/extra_test.go b/cli/extra_test.go
index 96812ca..0f6d5c5 100644
--- a/cli/extra_test.go
+++ b/cli/extra_test.go
@@ -14,6 +14,8 @@ func Test_arrayExtra_Set(t *testing.T) {
 		wantErr bool
 	}{
 		{"With tables", "simple: id < 10", arrayExtra{{Table: "simple", Where: "id < 10"}}, false},
+		{"With tables and all", "simple: all", arrayExtra{{Table: "simple", Where: "1=1"}}, false},
+		{"With tables only", "simple", arrayExtra{{Table: "simple"}}, false},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
diff --git a/subsetter/graph.go b/subsetter/graph.go
index a04157a..233e529 100644
--- a/subsetter/graph.go
+++ b/subsetter/graph.go
@@ -6,6 +6,17 @@ import (
 	"github.com/stevenle/topsort"
 )
 
+// TableGraph generates a topologically sorted list of table names based on their relations.
+// It takes a primary table name and a slice of Relation objects as input.
+// The function returns a slice of strings representing the sorted table names and an error if any.
+//
+// Parameters:
+//   - primary: The name of the primary table to start the topological sort from.
+//   - relations: A slice of Relation objects representing the relationships between tables.
+//
+// Returns:
+//   - l: A slice of strings representing the topologically sorted table names.
+//   - err: An error if the topological sort fails or if there is an issue adding edges to the graph.
 func TableGraph(primary string, relations []Relation) (l []string, err error) {
 	graph := topsort.NewGraph() // Create a new graph
 
@@ -25,6 +36,17 @@ func TableGraph(primary string, relations []Relation) (l []string, err error) {
 	return
 }
 
+// RequiredTableGraph generates a list of required tables in topological order
+// starting from the primary table. It uses the provided relations to build a
+// directed graph and performs a topological sort.
+//
+// Parameters:
+// - primary: The name of the primary table to start the topological sort from.
+// - relations: A slice of Relation structs that define the relationships between tables.
+//
+// Returns:
+// - l: A slice of strings representing the tables in topological order.
+// - err: An error if the graph construction or topological sort fails.
 func RequiredTableGraph(primary string, relations []Relation) (l []string, err error) {
 	graph := topsort.NewGraph() // Create a new graph