Only check disks in the table's storage policy

cmd/partforge/main.go

@@ -996,7 +996,7 @@ func runUploadFreeze(ctx context.Context, args []string) error {
 	slog.Info("validated source schema and destination table", "stage", "validate_schemas", "destination_schema_table", chhttp.TableSQL(destinationTableRef.Database, destinationTableRef.Table))
 
 	slog.Info("discovering local ClickHouse disks", "stage", "discover_disks")
-	disks, err := freeze.LocalDisks(ctx, ch)
+	disks, err := freeze.LocalDisks(ctx, ch, *database, *table)
 	if err != nil {
 		return err
 	}

internal/freeze/scan.go

@@ -26,10 +26,11 @@ type Disk struct {
 	Type string
 }
 
-func LocalDisks(ctx context.Context, ch chhttp.Client) ([]Disk, error) {
-	out, err := ch.QueryString(ctx, "SELECT name, path, type FROM system.disks ORDER BY name FORMAT TSV")
+// LocalDisks returns the disks in the storage policy of database.table; other disks on the server are ignored.
+func LocalDisks(ctx context.Context, ch chhttp.Client, database, table string) ([]Disk, error) {
+	out, err := ch.QueryString(ctx, "SELECT name, path, type FROM system.disks WHERE name IN (SELECT arrayJoin(disks) FROM system.storage_policies WHERE policy_name IN (SELECT storage_policy FROM system.tables WHERE database = "+chhttp.StringLiteral(database)+" AND name = "+chhttp.StringLiteral(table)+")) ORDER BY name FORMAT TSV")
 	if err != nil {
-		return nil, fmt.Errorf("query ClickHouse disks: %w", err)
+		return nil, fmt.Errorf("query ClickHouse disks for %s: %w", chhttp.TableSQL(database, table), err)
 	}
 	rows, err := chhttp.FormatTSVStrings(out, 3)
 	if err != nil {
@@ -44,7 +45,7 @@ func LocalDisks(ctx context.Context, ch chhttp.Client) ([]Disk, error) {
 		disks = append(disks, disk)
 	}
 	if len(disks) == 0 {
-		return nil, fmt.Errorf("ClickHouse reported no local disks")
+		return nil, fmt.Errorf("ClickHouse reported no disks for the storage policy of %s", chhttp.TableSQL(database, table))
 	}
 	return disks, nil
 }

internal/rewrite/compactor.go

@@ -294,7 +294,7 @@ func (c Compactor) Compact(ctx context.Context, item CompactWorkItem) (CompactRe
 	if err := c.ClickHouse.Exec(ctx, "ALTER TABLE "+chhttp.TableSQL(item.DestinationDatabase, item.DestinationTable)+" FREEZE WITH NAME "+chhttp.StringLiteral(freezeName)); err != nil {
 		return CompactResult{}, fmt.Errorf("freeze compact destination table %s: %w", chhttp.TableSQL(item.DestinationDatabase, item.DestinationTable), err)
 	}
-	disks, err := freeze.LocalDisks(ctx, c.ClickHouse)
+	disks, err := freeze.LocalDisks(ctx, c.ClickHouse, item.DestinationDatabase, item.DestinationTable)
 	if err != nil {
 		return CompactResult{}, err
 	}

internal/rewrite/processor.go

@@ -691,7 +691,7 @@ func (p Processor) rewritePart(ctx context.Context, m manifest.Manifest, sourceP
 	if err := p.ClickHouse.Exec(ctx, "ALTER TABLE "+chhttp.TableSQL(m.Dest.Database, m.Dest.Table)+" FREEZE WITH NAME "+chhttp.StringLiteral(freezeName)); err != nil {
 		return rewriteResult{}, fmt.Errorf("freeze destination table %s: %w", chhttp.TableSQL(m.Dest.Database, m.Dest.Table), err)
 	}
-	disks, err := freeze.LocalDisks(ctx, p.ClickHouse)
+	disks, err := freeze.LocalDisks(ctx, p.ClickHouse, m.Dest.Database, m.Dest.Table)
 	if err != nil {
 		return rewriteResult{}, err
 	}