Prevent cascade deletion of excluded tables #8 (#17)

* Prevent cascade deletion of excluded tables

* Divide fetchTablesSchemas into fetching and filtering

* Added unit test of filterTableSchemas

* Applied format

* Put index_schema.go into table_schema.go

* Directly pass a query for fetching table schema

* Fix TestFilterTableSchemas

* Move target logic from fetch to filter

* Revert unnecessary change.

* Improve filterTableSchemas switch statement

* Change map[string]struct{} into map[string]bool.

* Separated TestFilterTableSchemas into TestTargetFilterTableSchemas and TestExcludeFilterTableSchemas.

* Improve switch case statement

* Refactored the filter logic to use recursion.

* Applied format

* Fix the bug building parentRelation and childRelation.

* Optimized redundant for statement.

* Eliminated the subtest structure.

* Remove cascade deletability check from constructTableLineages

* Add test to ensure parents of tables with ON DELETE NO ACTION are not excluded

README.md

@@ -53,7 +53,7 @@ Application Options:
   -d, --database= (required) Cloud Spanner Database ID. [$SPANNER_DATABASE_ID]
   -q, --quiet     Disable all interactive prompts.
   -t, --tables=   Comma separated table names to be truncated. Default to truncate all tables if not specified.
-  -e, --exclude-tables Comma separated table names to be exempted from truncating. 'tables' and 'exclude-tables' cannot co-exist.
+  -e, --exclude-tables Comma separated table names to be exempted from truncating. 'tables' and 'exclude-tables' cannot co-exist. If interleaved tables are specified, the parent table is also excluded.
 Help Options:
   -h, --help      Show this help message
 ```

truncate/run.go

@@ -56,14 +56,21 @@ func Run(ctx context.Context, projectID, instanceID, databaseID string, quiet bo
 // This function uses an externally passed Cloud Spanner client.
 func RunWithClient(ctx context.Context, client *spanner.Client, quiet bool, out io.Writer, targetTables, excludeTables []string) error {
 	fmt.Fprintf(out, "Fetching table schema from %s\n", client.DatabaseName())
-	schemas, err := fetchTableSchemas(ctx, client, targetTables, excludeTables)
+	schemas, err := fetchTableSchemas(ctx, client)
 	if err != nil {
 		return fmt.Errorf("failed to fetch table schema: %v", err)
 	}
+
+	schemas, err = filterTableSchemas(schemas, targetTables, excludeTables)
+	if err != nil {
+		return fmt.Errorf("failed to filter table schema: %v", err)
+	}
+
 	for _, schema := range schemas {
 		fmt.Fprintf(out, "%s\n", schema.tableName)
 	}
 	fmt.Fprintf(out, "\n")
+
 	if !quiet {
 		if !confirm(out, "Rows in these tables will be deleted. Do you want to continue?") {
 			return nil

truncate/table_schema.go

@@ -18,6 +18,7 @@ package truncate
 
 import (
 	"context"
+	"errors"
 
 	"cloud.google.com/go/spanner"
 )
@@ -43,6 +44,14 @@ type tableSchema struct {
 	referencedBy []string
 }
 
+func (t *tableSchema) isCascadeDeletable() bool {
+	return t.parentOnDeleteAction == deleteActionCascadeDelete
+}
+
+func (t *tableSchema) isRoot() bool {
+	return t.parentTableName == ""
+}
+
 // indexSchema represents secondary index metadata.
 type indexSchema struct {
 	indexName string
@@ -54,7 +63,17 @@ type indexSchema struct {
 	parentTableName string
 }
 
-func fetchTableSchemas(ctx context.Context, client *spanner.Client, targetTables, excludeTables []string) ([]*tableSchema, error) {
+// tableLineage represents a table schema and its ancestors and descendants.
+// This is used to represent inter-table relationships.
+// The order of ancestors and descendants are not guaranteed.
+type tableLineage struct {
+	tableSchema *tableSchema
+	ancestors   []*tableSchema
+	descendants []*tableSchema
+}
+
+// fetchTableSchemas fetches schema information from spanner database.
+func fetchTableSchemas(ctx context.Context, client *spanner.Client) ([]*tableSchema, error) {
 	// This query fetches the table metadata and relationships.
 	iter := client.Single().Query(ctx, spanner.NewStatement(`
 		WITH FKReferences AS (
@@ -71,19 +90,6 @@ func fetchTableSchemas(ctx context.Context, client *spanner.Client, targetTables
 		ORDER BY T.TABLE_NAME ASC
 	`))
 
-	truncateAll := true
-	targets := make(map[string]bool, len(targetTables))
-	excludes := make(map[string]bool, len(excludeTables))
-	if len(targetTables) > 0 || len(excludeTables) > 0 {
-		truncateAll = false
-		for _, t := range targetTables {
-			targets[t] = true
-		}
-		for _, t := range excludeTables {
-			excludes[t] = true
-		}
-	}
-
 	var tables []*tableSchema
 	if err := iter.Do(func(r *spanner.Row) error {
 		var (
@@ -96,18 +102,6 @@ func fetchTableSchemas(ctx context.Context, client *spanner.Client, targetTables
 			return err
 		}
 
-		if !truncateAll {
-			if len(excludes) != 0 {
-				if _, ok := excludes[tableName]; ok {
-					return nil
-				}
-			} else {
-				if _, ok := targets[tableName]; !ok {
-					return nil
-				}
-			}
-		}
-
 		var parentTableName string
 		if parent.Valid {
 			parentTableName = parent.StringVal
@@ -137,6 +131,130 @@ func fetchTableSchemas(ctx context.Context, client *spanner.Client, targetTables
 	return tables, nil
 }
 
+// filterTableSchemas filters tables with given targetTables and excludeTables.
+// If targetTables is not empty, it fetches only the specified tables.
+// If excludeTables is not empty, it excludes the specified tables.
+// TargetTables and excludeTables cannot be specified at the same time.
+func filterTableSchemas(tables []*tableSchema, targetTables, excludeTables []string) ([]*tableSchema, error) {
+	isExclude := len(excludeTables) > 0
+	isTarget := len(targetTables) > 0
+
+	switch {
+	case isTarget && isExclude:
+		return nil, errors.New("both targetTables and excludeTables cannot be specified at the same time")
+	case isTarget:
+		return targetFilterTableSchemas(tables, targetTables), nil
+	case isExclude:
+		return excludeFilterTableSchemas(tables, excludeTables), nil
+	default: // No target and exclude tables are specified.
+		return tables, nil
+	}
+}
+
+// targetFilterTableSchemas filters tables with given targetTables.
+// If targetTables is empty, it returns all tables.
+func targetFilterTableSchemas(tables []*tableSchema, targetTables []string) []*tableSchema {
+	if len(targetTables) == 0 {
+		return tables
+	}
+
+	isTarget := make(map[string]bool, len(tables))
+	for _, t := range targetTables {
+		isTarget[t] = true
+	}
+
+	// TODO: Add child tables that may be deleted in cascade (#18)
+
+	filtered := make([]*tableSchema, 0, len(tables))
+	for _, t := range tables {
+		if isTarget[t.tableName] {
+			filtered = append(filtered, t)
+		}
+	}
+
+	return filtered
+}
+
+// excludeFilterTableSchemas filters tables with given excludeTables.
+// If excludeTables is empty, it returns all tables.
+// When an exclude table is cascade deletable, its parent table is also excluded.
+func excludeFilterTableSchemas(tables []*tableSchema, excludeTableSchemas []string) []*tableSchema {
+	if len(excludeTableSchemas) == 0 {
+		return tables
+	}
+
+	isExclude := make(map[string]bool, len(tables))
+	for _, t := range excludeTableSchemas {
+		isExclude[t] = true
+	}
+
+	// Additionally exclude parent tables that may delete the exclude tables in cascade
+	lineages := constructTableLineages(tables)
+	for _, l := range lineages {
+		if isExclude[l.tableSchema.tableName] && l.tableSchema.isCascadeDeletable() {
+			for _, a := range l.ancestors {
+				isExclude[a.tableName] = true
+			}
+		}
+	}
+
+	filtered := make([]*tableSchema, 0, len(tables))
+	for _, t := range tables {
+		if !isExclude[t.tableName] {
+			filtered = append(filtered, t)
+		}
+	}
+
+	return filtered
+}
+
+// constructTableLineages returns a list of interleave Lineages.
+// This function creates tableLineage for each of all given tableSchemas.
+func constructTableLineages(tables []*tableSchema) []*tableLineage {
+	tableMap := make(map[string]*tableSchema, len(tables))
+	for _, t := range tables {
+		tableMap[t.tableName] = t
+	}
+
+	parentRelation := make(map[string]*tableSchema, len(tables))
+	childRelation := make(map[string][]*tableSchema, len(tables))
+	for _, t := range tables {
+		if !t.isRoot() {
+			parentRelation[t.tableName] = tableMap[t.parentTableName]
+			childRelation[t.parentTableName] = append(childRelation[t.parentTableName], t)
+		}
+	}
+
+	lineages := make([]*tableLineage, 0, len(tables))
+	for _, t := range tables {
+		lineages = append(lineages, &tableLineage{
+			tableSchema: t,
+			ancestors:   findAncestors(t, parentRelation, nil),
+			descendants: findDescendants(t, childRelation, nil),
+		})
+	}
+
+	return lineages
+}
+
+// findAncestors recursively finds all ancestors of the given table .
+func findAncestors(table *tableSchema, parentRelation map[string]*tableSchema, ancestors []*tableSchema) []*tableSchema {
+	if parent, ok := parentRelation[table.tableName]; ok {
+		ancestors = append(ancestors, parent)
+		return findAncestors(parent, parentRelation, ancestors)
+	}
+	return ancestors
+}
+
+// findDescendants recursively finds all descendants of the given table.
+func findDescendants(table *tableSchema, childRelation map[string][]*tableSchema, descendants []*tableSchema) []*tableSchema {
+	for _, child := range childRelation[table.tableName] {
+		descendants = append(descendants, child)
+		descendants = findDescendants(child, childRelation, descendants)
+	}
+	return descendants
+}
+
 func fetchIndexSchemas(ctx context.Context, client *spanner.Client) ([]*indexSchema, error) {
 	// This query fetches defined indexes.
 	iter := client.Single().Query(ctx, spanner.NewStatement(`