github.com/weaviate/weaviate@v1.24.6/usecases/schema/migrate/fs/file_structure_migration.go

//                           _       _
// __      _____  __ ___   ___  __ _| |_ ___
// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
//
//  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
//
//  CONTACT: hello@weaviate.io
//

package fs

import (
	"fmt"
	"os"
	"path"
	"path/filepath"
	"strings"

	"github.com/weaviate/weaviate/adapters/repos/db/helpers"
	entschema "github.com/weaviate/weaviate/entities/schema"
	"github.com/weaviate/weaviate/usecases/sharding"
)

const vectorIndexCommitLog = `hnsw.commitlog.d`

func MigrateToHierarchicalFS(rootPath string, s schemaGetter) error {
	root, err := os.ReadDir(rootPath)
	if err != nil {
		return fmt.Errorf("read source path %q: %w", rootPath, err)
	}
	fm := newFileMatcher(s, rootPath)
	plan, err := assembleFSMigrationPlan(root, rootPath, fm)
	if err != nil {
		return err
	}

	for newRoot, parts := range plan.partsByShard {
		for _, part := range parts {
			newPath := path.Join(newRoot, part.newRelPath)
			absDir, _ := filepath.Split(newPath)
			if err := os.MkdirAll(absDir, os.ModePerm); err != nil {
				return fmt.Errorf("mkdir %q: %w", absDir, err)
			}
			if err = os.Rename(part.oldAbsPath, newPath); err != nil {
				return fmt.Errorf("mv %s %s: %w", part.oldAbsPath, newPath, err)
			}
		}
	}

	return nil
}

type migrationPart struct {
	oldAbsPath string
	newRelPath string
}

type shardRoot = string

type migrationPlan struct {
	rootPath     string
	partsByShard map[shardRoot][]migrationPart
}

func newMigrationPlan(rootPath string) *migrationPlan {
	return &migrationPlan{rootPath: rootPath, partsByShard: make(map[string][]migrationPart)}
}

func (p *migrationPlan) append(class, shard, oldRootRelPath, newShardRelPath string) {
	shardRoot := path.Join(p.rootPath, strings.ToLower(class), shard)
	p.partsByShard[shardRoot] = append(p.partsByShard[shardRoot], migrationPart{
		oldAbsPath: path.Join(p.rootPath, oldRootRelPath),
		newRelPath: newShardRelPath,
	})
}

func (p *migrationPlan) prepend(class, shard, oldRootRelPath, newShardRelPath string) {
	shardRoot := path.Join(p.rootPath, strings.ToLower(class), shard)
	p.partsByShard[shardRoot] = append([]migrationPart{{
		oldAbsPath: path.Join(p.rootPath, oldRootRelPath),
		newRelPath: newShardRelPath,
	}}, p.partsByShard[shardRoot]...)
}

func assembleFSMigrationPlan(entries []os.DirEntry, rootPath string, fm *fileMatcher) (*migrationPlan, error) {
	plan := newMigrationPlan(rootPath)

	for _, entry := range entries {
		if ok, cs := fm.isShardLsmDir(entry); ok {
			// make sure lsm dir is moved first, otherwise os.Rename may fail
			// if directory already exists (created by other files/dirs moved before)
			plan.prepend(cs.class, cs.shard,
				entry.Name(),
				"lsm")
		} else if ok, cs, suffix := fm.isShardFile(entry); ok {
			plan.append(cs.class, cs.shard,
				entry.Name(),
				suffix)
		} else if ok, cs := fm.isShardCommitLogDir(entry); ok {
			plan.append(cs.class, cs.shard,
				entry.Name(),
				fmt.Sprintf("main.%s", vectorIndexCommitLog))
		} else if ok, csp := fm.isShardGeoCommitLogDir(entry); ok {
			plan.append(csp.class, csp.shard,
				entry.Name(),
				fmt.Sprintf("geo.%s.%s", csp.geoProp, vectorIndexCommitLog))
		} else if ok, css := fm.isPqDir(entry); ok {
			for _, cs := range css {
				plan.append(cs.class, cs.shard,
					path.Join(strings.ToLower(entry.Name()), cs.shard, "compressed_objects"),
					path.Join("lsm", helpers.VectorsCompressedBucketLSM))
			}

			// explicitly rename Class directory starting with uppercase to lowercase
			// as MkdirAll will not create lowercased dir if uppercased one exists
			oldClassRoot := path.Join(rootPath, entry.Name())
			newClassRoot := path.Join(rootPath, strings.ToLower(entry.Name()))
			if err := os.Rename(oldClassRoot, newClassRoot); err != nil {
				return nil, fmt.Errorf(
					"rename pq index dir to avoid collision, old: %q, new: %q, err: %w",
					oldClassRoot, newClassRoot, err)
			}
		}
	}
	return plan, nil
}

type classShard struct {
	class string
	shard string
}

type classShardGeoProp struct {
	class   string
	shard   string
	geoProp string
}

type fileMatcher struct {
	rootPath            string
	shardLsmDirs        map[string]*classShard
	shardFilePrefixes   map[string]*classShard
	shardGeoDirPrefixes map[string]*classShardGeoProp
	classes             map[string][]*classShard
}

type schemaGetter interface {
	CopyShardingState(class string) *sharding.State
	GetSchemaSkipAuth() entschema.Schema
}

func newFileMatcher(schemaGetter schemaGetter, rootPath string) *fileMatcher {
	shardLsmDirs := make(map[string]*classShard)
	shardFilePrefixes := make(map[string]*classShard)
	shardGeoDirPrefixes := make(map[string]*classShardGeoProp)
	classes := make(map[string][]*classShard)

	sch := schemaGetter.GetSchemaSkipAuth()
	for _, class := range sch.Objects.Classes {
		shards := schemaGetter.CopyShardingState(class.Class).AllLocalPhysicalShards()
		lowercasedClass := strings.ToLower(class.Class)

		var geoProps []string
		for _, prop := range class.Properties {
			if dt, ok := entschema.AsPrimitive(prop.DataType); ok && dt == entschema.DataTypeGeoCoordinates {
				geoProps = append(geoProps, prop.Name)
			}
		}

		classes[class.Class] = make([]*classShard, 0, len(shards))
		for _, shard := range shards {
			cs := &classShard{class: class.Class, shard: shard}
			shardLsmDirs[fmt.Sprintf("%s_%s_lsm", lowercasedClass, shard)] = cs
			shardFilePrefixes[fmt.Sprintf("%s_%s", lowercasedClass, shard)] = cs
			classes[class.Class] = append(classes[class.Class], cs)

			for _, geoProp := range geoProps {
				csp := &classShardGeoProp{class: class.Class, shard: shard, geoProp: geoProp}
				shardGeoDirPrefixes[fmt.Sprintf("%s_%s_%s", lowercasedClass, shard, geoProp)] = csp
			}
		}
	}

	return &fileMatcher{
		rootPath:            rootPath,
		shardLsmDirs:        shardLsmDirs,
		shardFilePrefixes:   shardFilePrefixes,
		shardGeoDirPrefixes: shardGeoDirPrefixes,
		classes:             classes,
	}
}

// Checks if entry is directory with name (class is lowercased):
// class_shard_lsm
func (fm *fileMatcher) isShardLsmDir(entry os.DirEntry) (bool, *classShard) {
	if !entry.IsDir() {
		return false, nil
	}
	if cs, ok := fm.shardLsmDirs[entry.Name()]; ok {
		return true, cs
	}
	return false, nil
}

// Checks if entry is file with name (class is lowercased):
// class_shard.*
// (e.g. class_shard.version, class_shard.indexcount)
func (fm *fileMatcher) isShardFile(entry os.DirEntry) (bool, *classShard, string) {
	if !entry.Type().IsRegular() {
		return false, nil, ""
	}
	parts := strings.SplitN(entry.Name(), ".", 2)
	if len(parts) != 2 {
		return false, nil, ""
	}
	if cs, ok := fm.shardFilePrefixes[parts[0]]; ok {
		return true, cs, parts[1]
	}
	return false, nil, ""
}

// Checks if entry is directory with name (class is lowercased):
// class_shard.hnsw.commitlog.d
func (fm *fileMatcher) isShardCommitLogDir(entry os.DirEntry) (bool, *classShard) {
	if !entry.IsDir() {
		return false, nil
	}
	parts := strings.SplitN(entry.Name(), ".", 2)
	if len(parts) != 2 {
		return false, nil
	}
	if parts[1] != vectorIndexCommitLog {
		return false, nil
	}
	if cs, ok := fm.shardFilePrefixes[parts[0]]; ok {
		return true, cs
	}
	return false, nil
}

// Checks if entry is directory with name (class is lowercased):
// class_shard_prop.hnsw.commitlog.d
func (fm *fileMatcher) isShardGeoCommitLogDir(entry os.DirEntry) (bool, *classShardGeoProp) {
	if !entry.IsDir() {
		return false, nil
	}
	parts := strings.SplitN(entry.Name(), ".", 2)
	if len(parts) != 2 {
		return false, nil
	}
	if parts[1] != vectorIndexCommitLog {
		return false, nil
	}
	if csp, ok := fm.shardGeoDirPrefixes[parts[0]]; ok {
		return true, csp
	}
	return false, nil
}

// Checks if entry is directory containing PQ index:
// Class/shard/compressed_object
func (fm *fileMatcher) isPqDir(entry os.DirEntry) (bool, []*classShard) {
	if !entry.IsDir() {
		return false, nil
	}

	resultcss := []*classShard{}
	if css, ok := fm.classes[entry.Name()]; ok {
		for _, cs := range css {
			pqDir := path.Join(fm.rootPath, cs.class, cs.shard, "compressed_objects")
			if info, err := os.Stat(pqDir); err == nil && info.IsDir() {
				resultcss = append(resultcss, cs)
			}
		}
		return true, resultcss
	}
	return false, nil
}