knowledge-governance/cmd/release-worker/runtime.go

package main

import (
	"context"
	"errors"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/jackc/pgx/v5/pgxpool"

	"knowledge-governance/internal/platform/weknora"
	v071 "knowledge-governance/internal/releaseadapter/v071"
	"knowledge-governance/internal/releasesnapshot"
	releaseworkflow "knowledge-governance/internal/workflows/release"
)

type releaseRuntime struct {
	publisher *releaseworkflow.RoutingPublisher
	pools     []*pgxpool.Pool
}

func (runtime *releaseRuntime) Close() {
	for index := len(runtime.pools) - 1; index >= 0; index-- {
		runtime.pools[index].Close()
	}
}

func buildReleaseRuntime(ctx context.Context, config releaseRuntimeConfig) (*releaseRuntime, error) {
	runtime := &releaseRuntime{}
	publishers := make(map[string]releaseworkflow.Publisher, len(config.Routes))
	for _, route := range config.Routes {
		publisher, pools, err := buildReleaseRoute(ctx, config.ArchiveRoot, config.TargetSnapshot, route)
		if err != nil {
			runtime.Close()
			return nil, err
		}
		runtime.pools = append(runtime.pools, pools...)
		publishers[route.KnowledgeBaseID] = publisher
	}
	router, err := releaseworkflow.NewRoutingPublisher(publishers)
	if err != nil {
		runtime.Close()
		return nil, errors.New("create release publisher routes")
	}
	runtime.publisher = router
	return runtime, nil
}

func buildReleaseRoute(ctx context.Context, archiveRoot string, snapshot releaseTargetSnapshotConfig, route releaseRouteConfig) (releaseworkflow.Publisher, []*pgxpool.Pool, error) {
	sourceDatabaseURL, err := readRestrictedValue(route.Source.DatabaseURLFile)
	if err != nil {
		return nil, nil, errors.New("read source database credential")
	}
	targetDatabaseURL, err := readRestrictedValue(route.Target.DatabaseURLFile)
	if err != nil {
		return nil, nil, errors.New("read target database credential")
	}
	if err := validateDistinctDatabaseURLs(sourceDatabaseURL, targetDatabaseURL); err != nil {
		return nil, nil, err
	}
	sourcePool, err := openReleasePool(ctx, sourceDatabaseURL)
	if err != nil {
		return nil, nil, errors.New("connect release source database")
	}
	pools := []*pgxpool.Pool{sourcePool}
	targetPool, err := openReleasePool(ctx, targetDatabaseURL)
	if err != nil {
		sourcePool.Close()
		return nil, nil, errors.New("connect release target database")
	}
	pools = append(pools, targetPool)
	cleanup := func(err error) (releaseworkflow.Publisher, []*pgxpool.Pool, error) {
		for index := len(pools) - 1; index >= 0; index-- {
			pools[index].Close()
		}
		return nil, nil, err
	}

	sourceSearchToken, err := readOptionalRestrictedValue(route.Source.Elasticsearch.TokenFile)
	if err != nil {
		return cleanup(errors.New("read source Elasticsearch credential"))
	}
	targetSearchToken, err := readOptionalRestrictedValue(route.Target.Elasticsearch.TokenFile)
	if err != nil {
		return cleanup(errors.New("read target Elasticsearch credential"))
	}
	sourceSearch, err := v071.NewSearchExporter(v071.SearchExporterConfig{BaseURL: route.Source.Elasticsearch.BaseURL, Index: route.Source.Elasticsearch.Index, AuthorizationToken: sourceSearchToken, PageSize: 500})
	if err != nil {
		return cleanup(errors.New("create source Elasticsearch exporter"))
	}
	targetSearchConfig := v071.SearchExporterConfig{BaseURL: route.Target.Elasticsearch.BaseURL, Index: route.Target.Elasticsearch.Index, AuthorizationToken: targetSearchToken, PageSize: 500}
	targetSearchInspector, err := v071.NewSearchExporter(targetSearchConfig)
	if err != nil {
		return cleanup(errors.New("create target Elasticsearch inspector"))
	}
	targetSearch, err := v071.NewSearchImporter(targetSearchConfig)
	if err != nil {
		return cleanup(errors.New("create target Elasticsearch importer"))
	}

	sourceObject, err := buildS3ObjectStore(route.Source.Object)
	if err != nil {
		return cleanup(errors.New("create source object reader"))
	}
	targetObject, err := buildS3ObjectStore(route.Target.Object)
	if err != nil {
		return cleanup(errors.New("create target object store"))
	}
	objectExporter, err := v071.NewObjectExporter(sourceObject, v071.ObjectExporterLimits{})
	if err != nil {
		return cleanup(errors.New("create source object exporter"))
	}
	objectImporter, err := v071.NewObjectImporter(targetObject, v071.TargetObjectConfig{StorageBackendID: route.Target.Object.StorageBackendID, Provider: route.Target.Object.Provider})
	if err != nil {
		return cleanup(errors.New("create target object importer"))
	}

	sourceContract, err := v071.NewTargetContractInspector(v071.TargetContractInspectorConfig{DB: sourcePool, Mapping: sourceSearch, KnowledgeBaseID: route.KnowledgeBaseID, TenantID: route.TenantID, IndexName: route.Source.Elasticsearch.Index, ReadTestedBase: true})
	if err != nil {
		return cleanup(errors.New("create source release contract inspector"))
	}
	targetContract, err := v071.NewTargetContractInspector(v071.TargetContractInspectorConfig{DB: targetPool, Mapping: targetSearchInspector, KnowledgeBaseID: route.KnowledgeBaseID, TenantID: route.TenantID, IndexName: route.Target.Elasticsearch.Index})
	if err != nil {
		return cleanup(errors.New("create target release contract inspector"))
	}
	postgresExporter, err := v071.NewPostgresExporter(sourcePool)
	if err != nil {
		return cleanup(errors.New("create source PostgreSQL exporter"))
	}
	postgresImporter, err := v071.NewPostgresImporter(targetPool, v071.PostgresImporterConfig{TargetStorageBackendID: route.Target.Object.StorageBackendID, TargetProvider: route.Target.Object.Provider, TargetBucket: route.Target.Object.Bucket})
	if err != nil {
		return cleanup(errors.New("create target PostgreSQL importer"))
	}
	documentExporter, err := v071.NewDocumentExporter(v071.DocumentExporterConfig{Contract: sourceContract, Postgres: postgresExporter, Search: sourceSearch, Objects: objectExporter, TenantID: route.TenantID, EmbeddingModel: route.EmbeddingModel, Now: time.Now})
	if err != nil {
		return cleanup(errors.New("create v0.7.1 document exporter"))
	}

	archive, err := v071.NewBundleArchive(filepath.Join(archiveRoot, route.KnowledgeBaseID), 1024*1024*1024)
	if err != nil {
		return cleanup(errors.New("open release bundle archive"))
	}
	apiKey, err := readRestrictedValue(route.Target.WeKnora.APIKeyFile)
	if err != nil {
		return cleanup(errors.New("read target WeKnora credential"))
	}
	weknoraClient, err := weknora.NewClient(weknora.ClientConfig{BaseURL: route.Target.WeKnora.BaseURL, APIKey: apiKey, Timeout: 60 * time.Second, AllowDockerHTTP: route.Target.WeKnora.AllowDockerHTTP})
	if err != nil {
		return cleanup(errors.New("create target WeKnora verification client"))
	}
	verifier, err := v071.NewWeKnoraVerifier(weknoraClient)
	if err != nil {
		return cleanup(errors.New("create target WeKnora verifier"))
	}
	documentTarget, err := v071.NewDocumentTarget(v071.DocumentTargetConfig{Archive: archive, Contract: targetContract, Postgres: postgresImporter, Search: targetSearch, Objects: objectImporter, Verifier: verifier})
	if err != nil {
		return cleanup(errors.New("create v0.7.1 document target"))
	}
	var publisherOptions []releaseworkflow.PublisherOption
	if snapshot.Root != "" {
		// 写目标库前必须有覆盖该库的、未过期的快照留档（见 releasesnapshot.Guard）。
		maxAge := time.Duration(snapshot.MaxAgeHours) * time.Hour
		if snapshot.MaxAgeHours == 0 {
			maxAge = 24 * time.Hour
		}
		policy := releaseworkflow.SnapshotPolicyMapped
		if snapshot.Policy == "always" {
			policy = releaseworkflow.SnapshotPolicyAlways
		}
		publisherOptions = append(publisherOptions, releaseworkflow.WithTargetSnapshots(
			releasesnapshot.Guard{Root: snapshot.Root, MaxAge: maxAge, VerifyLocalFiles: true}, policy))
	}
	publisher, err := releaseworkflow.NewSidecarPublisher(documentExporter, documentTarget, publisherOptions...)
	if err != nil {
		return cleanup(errors.New("create v0.7.1 release publisher"))
	}
	return publisher, pools, nil
}

func buildS3ObjectStore(config releaseObjectConfig) (*v071.S3ObjectStore, error) {
	accessKey, err := readRestrictedValue(config.AccessKeyFile)
	if err != nil {
		return nil, err
	}
	secretKey, err := readRestrictedValue(config.SecretKeyFile)
	if err != nil {
		return nil, err
	}
	parsed, err := url.Parse(config.Endpoint)
	if err != nil {
		return nil, errors.New("invalid object endpoint")
	}
	return v071.NewS3ObjectStore(v071.S3ObjectStoreConfig{Provider: config.Provider, Endpoint: config.Endpoint, Region: config.Region, AccessKeyID: accessKey, SecretAccessKey: secretKey, Bucket: config.Bucket, Secure: parsed.Scheme == "https", ForcePathStyle: config.Provider == "minio"})
}

func readRestrictedValue(path string) (string, error) {
	return weknora.ReadAPIKeyFile(path)
}

func readOptionalRestrictedValue(path string) (string, error) {
	if path == "" {
		return "", nil
	}
	return readRestrictedValue(path)
}

func openReleasePool(ctx context.Context, databaseURL string) (*pgxpool.Pool, error) {
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		return nil, err
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		return nil, err
	}
	return pool, nil
}

func validateDistinctDatabaseURLs(sourceURL, targetURL string) error {
	source, err := pgxpool.ParseConfig(sourceURL)
	if err != nil {
		return errors.New("invalid source database credential")
	}
	target, err := pgxpool.ParseConfig(targetURL)
	if err != nil {
		return errors.New("invalid target database credential")
	}
	sourceIdentity := strings.ToLower(source.ConnConfig.Host) + "|" + strconv.FormatUint(uint64(source.ConnConfig.Port), 10) + "|" + source.ConnConfig.Database
	targetIdentity := strings.ToLower(target.ConnConfig.Host) + "|" + strconv.FormatUint(uint64(target.ConnConfig.Port), 10) + "|" + target.ConnConfig.Database
	if sourceIdentity == targetIdentity {
		return errors.New("release source and target database must be distinct")
	}
	return nil
}