generate-import-fixtures: add region and nation tables

Add the two remaining TPC-H dimension tables (region with 5 rows, nation
with 25 rows) so the tool can generate complete fixture sets.

Since dbgen produces these small tables as a single file (e.g. region.tbl)
rather than sharded files (e.g. region.tbl.1..8), extract the shard
logic into a shardFiles() function that returns one entry for single-file
tables and 8 entries for sharded tables.

Epic: CRDB-62435

Release note: None

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sravotto

pkg/cmd/generate-import-fixtures/main.go

@@ -88,30 +88,29 @@ func main() {
 		}
 		fmt.Printf("processing table: %s\n", tableName)
 
-		for shardIdx := 1; shardIdx <= 8; shardIdx++ {
-			fileName := fmt.Sprintf("%s.tbl.%d", tableName, shardIdx)
-			filePath := filepath.Join(*inputDir, fileName)
+		for _, s := range shardFiles(tableName) {
+			filePath := filepath.Join(*inputDir, s.fileName)
 
 			f, err := os.Open(filePath)
 			if err != nil {
 				log.Fatalf("opening %s: %v", filePath, err)
 			}
 
-			writer, err := format.NewWriter(tableDef, *outputDir, shardIdx)
+			writer, err := format.NewWriter(tableDef, *outputDir, s.idx)
 			if err != nil {
 				f.Close()
-				log.Fatalf("creating writer for %s shard %d: %v", tableName, shardIdx, err)
+				log.Fatalf("creating writer for %s shard %d: %v", tableName, s.idx, err)
 			}
 
 			if _, err := processShard(f, tableDef, writer); err != nil {
 				_ = writer.Close()
 				f.Close()
-				log.Fatalf("processing %s shard %d: %v", tableName, shardIdx, err)
+				log.Fatalf("processing %s shard %d: %v", tableName, s.idx, err)
 			}
 
 			if err := writer.Close(); err != nil {
 				f.Close()
-				log.Fatalf("closing writer for %s shard %d: %v", tableName, shardIdx, err)
+				log.Fatalf("closing writer for %s shard %d: %v", tableName, s.idx, err)
 			}
 			f.Close()
 		}
@@ -128,6 +127,36 @@ func availableFormats() string {
 	return strings.Join(names, ", ")
 }
 
+// shardFile identifies a single input file and its shard index.
+type shardFile struct {
+	fileName string
+	idx      int
+}
+
+// singleFileTables lists tables where dbgen produces a single unsuffixed file
+// instead of per-shard files (e.g. region.tbl instead of region.tbl.1..8).
+var singleFileTables = map[string]bool{
+	"region": true,
+	"nation": true,
+}
+
+// shardFiles returns the list of input files for the given table. Small
+// dimension tables (region, nation) have a single file; all others are split
+// into 8 shards by dbgen.
+func shardFiles(tableName string) []shardFile {
+	if singleFileTables[tableName] {
+		return []shardFile{{fileName: tableName + ".tbl", idx: 1}}
+	}
+	shards := make([]shardFile, 8)
+	for i := range 8 {
+		shards[i] = shardFile{
+			fileName: fmt.Sprintf("%s.tbl.%d", tableName, i+1),
+			idx:      i + 1,
+		}
+	}
+	return shards
+}
+
 // processShard reads a pipe-delimited TPC-H file and streams parsed rows to the
 // writer in fixed-size batches. It returns the total number of rows processed.
 func processShard(r io.Reader, table TableDef, writer FormatWriter) (int, error) {

pkg/cmd/generate-import-fixtures/tpch.go

@@ -156,6 +156,23 @@ var tpchTables = map[string]TableDef{
 			stringCol("s_comment"),
 		},
 	},
+	"region": {
+		Name: "region",
+		Columns: []ColumnDef{
+			longCol("r_regionkey"),
+			stringCol("r_name"),
+			stringCol("r_comment"),
+		},
+	},
+	"nation": {
+		Name: "nation",
+		Columns: []ColumnDef{
+			longCol("n_nationkey"),
+			stringCol("n_name"),
+			longCol("n_regionkey"),
+			stringCol("n_comment"),
+		},
+	},
 }
 
 // allTPCHTables returns the names of all defined TPC-H tables.