feat(artifact): add admin APIs for KB consolidation migration (#331)

**Because**

- The KB consolidation migration (convert000070) needs to move files
from legacy/duplicate KBs to the primary default KB
- Files can belong to multiple KBs (many-to-many relationship via
junction table)
- After consolidating files, legacy KBs need to be hard-deleted
(bypassing soft delete) with CASCADE removal of file-KB associations
- The previous `easimon/maximize-build-space` action was deprecated

**This commit**

- Adds `AddFilesToKnowledgeBaseAdmin` handler that accepts file resource
names (`namespaces/{namespace}/files/{file}`), resolves file IDs to UIDs
internally, and inserts into junction table (skipping duplicates via `ON
CONFLICT DO NOTHING`)
- Adds `DeleteKnowledgeBaseAdmin` handler that drops Milvus collection,
purges ACL entry, and hard-deletes KB record (CASCADE removes
file_knowledge_base associations)
- Adds `HardDeleteKnowledgeBase` repository method using
`Unscoped().Delete()` to bypass soft delete
- Adds `AddFilesToKnowledgeBase` repository method with raw SQL to
lookup file UIDs by IDs and insert associations
- Updates integration-test workflow to use `jlumbroso/free-disk-space`
action
- Regenerates repository mocks with new methods

Author: pinglin

.github/workflows/integration-test.yml

@@ -13,16 +13,16 @@ jobs:
     name: Integration test
     runs-on: ubuntu-latest
     steps:
-      - name: Maximize build space
-        uses: easimon/maximize-build-space@master
+      - name: Free disk space
+        uses: jlumbroso/free-disk-space@main
         with:
-          root-reserve-mb: 5120
-          remove-dotnet: "true"
-          remove-android: "true"
-          remove-haskell: "true"
-          remove-codeql: "true"
-          remove-docker-images: "true"
-          build-mount-path: "/var/lib/docker"
+          tool-cache: false
+          android: true
+          dotnet: true
+          haskell: true
+          large-packages: true
+          docker-images: true
+          swap-storage: false
 
       - name: Restart docker
         run: sudo service docker restart

go.mod

@@ -13,7 +13,7 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware v1.4.0
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.1
 	github.com/iancoleman/strcase v0.3.0
-	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201062717-96c489b7a2ba
+	github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201213435-984ba4ebc92e
 	github.com/instill-ai/x v0.10.1-alpha.0.20260129195415-09e1680f104d
 	github.com/knadh/koanf v1.5.0
 	github.com/mennanov/fieldmask-utils v1.1.2

go.sum

@@ -419,8 +419,8 @@ github.com/iancoleman/strcase v0.3.0 h1:nTXanmYxhfFAMjZL34Ov6gkzEsSJZ5DbhxWjvSAS
 github.com/iancoleman/strcase v0.3.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47ZCWhYzw7ho=
 github.com/imkira/go-interpol v1.1.0/go.mod h1:z0h2/2T3XF8kyEPpRgJ3kmNv+C43p+I/CoI+jC3w2iA=
 github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201062717-96c489b7a2ba h1:H6YsrxOBckOZD8gp7yzXt4DrQi5b3XZBcmrt7wj4DMI=
-github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201062717-96c489b7a2ba/go.mod h1:bCnBosofpaUxKBuTTJM3/I3thAK37kvfBnKByjnLsl4=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201213435-984ba4ebc92e h1:xXUG15PsU2sU1BFJmAvuY25HNCA7Xa0gs4CjG1oPSzQ=
+github.com/instill-ai/protogen-go v0.3.3-alpha.0.20260201213435-984ba4ebc92e/go.mod h1:bCnBosofpaUxKBuTTJM3/I3thAK37kvfBnKByjnLsl4=
 github.com/instill-ai/x v0.10.1-alpha.0.20260129195415-09e1680f104d h1:FLDqMpALmq2e/j9inTGWrY4mhRFhe5/xXvjjxwQDqDg=
 github.com/instill-ai/x v0.10.1-alpha.0.20260129195415-09e1680f104d/go.mod h1:/xG0LMh43kyyK9jXGtV/JGd5p9GVcreSZ6GRGVbALEU=
 github.com/iris-contrib/blackfriday v2.0.0+incompatible/go.mod h1:UzZ2bDEoaSGPbkg6SAB4att1aAwTmVIx/5gCVqeyUdI=

pkg/handler/private.go

@@ -1049,3 +1049,119 @@ func (h *PrivateHandler) ResetKnowledgeBaseEmbeddingsAdmin(ctx context.Context,
 		FilesToReembed: int32(fileCount),
 	}, nil
 }
+
+// AddFilesToKnowledgeBaseAdmin adds file associations to a target KB by file resource names (admin only).
+// Files can belong to multiple KBs (many-to-many relationship).
+func (h *PrivateHandler) AddFilesToKnowledgeBaseAdmin(ctx context.Context, req *artifactpb.AddFilesToKnowledgeBaseAdminRequest) (*artifactpb.AddFilesToKnowledgeBaseAdminResponse, error) {
+	logger, _ := logx.GetZapLogger(ctx)
+
+	// Parse target KB from name (format: namespaces/{namespace}/knowledge-bases/{kb})
+	targetNamespaceID, targetKBID, err := parseKnowledgeBaseFromName(req.GetTargetKnowledgeBase())
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid target_knowledge_base format: %v", err)
+	}
+
+	logger.Info("AddFilesToKnowledgeBaseAdmin called",
+		zap.String("target_namespace_id", targetNamespaceID),
+		zap.String("target_kb_id", targetKBID),
+		zap.Int("file_count", len(req.GetFiles())))
+
+	// Get target namespace
+	targetNs, err := h.service.GetNamespaceByNsID(ctx, targetNamespaceID)
+	if err != nil {
+		logger.Error("failed to get target namespace", zap.Error(err))
+		return nil, status.Errorf(codes.NotFound, "target namespace not found: %v", err)
+	}
+
+	// Get target knowledge base
+	targetKB, err := h.service.Repository().GetKnowledgeBaseByOwnerAndKbID(ctx, targetNs.NsUID, targetKBID)
+	if err != nil {
+		logger.Error("failed to get target knowledge base", zap.Error(err))
+		return nil, status.Errorf(codes.NotFound, "target knowledge base not found: %v", err)
+	}
+
+	// Parse file resource names to extract file IDs
+	// Format: namespaces/{namespace}/files/{file}
+	fileIDs := make([]string, 0, len(req.GetFiles()))
+	for _, fileName := range req.GetFiles() {
+		_, _, fileID, err := parseFileFromName(fileName)
+		if err != nil {
+			return nil, status.Errorf(codes.InvalidArgument, "invalid file resource name format: %s", fileName)
+		}
+		fileIDs = append(fileIDs, fileID)
+	}
+
+	// Add file associations
+	filesAdded, err := h.service.Repository().AddFilesToKnowledgeBase(ctx, types.KnowledgeBaseUIDType(targetKB.UID), fileIDs)
+	if err != nil {
+		logger.Error("failed to add files to knowledge base", zap.Error(err))
+		return nil, status.Errorf(codes.Internal, "failed to add files: %v", err)
+	}
+
+	logger.Info("AddFilesToKnowledgeBaseAdmin completed",
+		zap.String("target_kb_uid", targetKB.UID.String()),
+		zap.Int64("files_added", filesAdded))
+
+	return &artifactpb.AddFilesToKnowledgeBaseAdminResponse{
+		FilesAdded: int32(filesAdded),
+	}, nil
+}
+
+// DeleteKnowledgeBaseAdmin force-deletes a knowledge base (admin only).
+// CASCADE removes file-KB associations from junction table. Files remain orphaned.
+// Also cleans up the Milvus collection.
+func (h *PrivateHandler) DeleteKnowledgeBaseAdmin(ctx context.Context, req *artifactpb.DeleteKnowledgeBaseAdminRequest) (*artifactpb.DeleteKnowledgeBaseAdminResponse, error) {
+	logger, _ := logx.GetZapLogger(ctx)
+
+	// Parse namespace and KB ID from name (format: namespaces/{namespace}/knowledge-bases/{kb})
+	namespaceID, knowledgeBaseID, err := parseKnowledgeBaseFromName(req.GetName())
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid name format: %v", err)
+	}
+
+	logger.Info("DeleteKnowledgeBaseAdmin called",
+		zap.String("namespace_id", namespaceID),
+		zap.String("knowledge_base_id", knowledgeBaseID))
+
+	// Get namespace
+	ns, err := h.service.GetNamespaceByNsID(ctx, namespaceID)
+	if err != nil {
+		logger.Error("failed to get namespace", zap.Error(err))
+		return nil, status.Errorf(codes.NotFound, "namespace not found: %v", err)
+	}
+
+	// Get knowledge base
+	kb, err := h.service.Repository().GetKnowledgeBaseByOwnerAndKbID(ctx, ns.NsUID, knowledgeBaseID)
+	if err != nil {
+		logger.Error("failed to get knowledge base", zap.Error(err))
+		return nil, status.Errorf(codes.NotFound, "knowledge base not found: %v", err)
+	}
+
+	collectionName := constant.KBCollectionName(kb.ActiveCollectionUID)
+
+	// Step 1: Drop Milvus collection if it exists
+	if err := h.service.Repository().DropCollection(ctx, collectionName); err != nil {
+		// Log but continue - collection might not exist
+		logger.Warn("Failed to drop Milvus collection (may not exist)", zap.Error(err), zap.String("collection", collectionName))
+	} else {
+		logger.Info("Dropped Milvus collection", zap.String("collection", collectionName))
+	}
+
+	// Step 2: Remove ACL entry
+	if err := h.service.ACLClient().Purge(ctx, "knowledgebase", kb.UID); err != nil {
+		// Log but continue - ACL entry might not exist
+		logger.Warn("Failed to purge ACL entry", zap.Error(err))
+	}
+
+	// Step 3: Hard delete the knowledge base (CASCADE will remove file_knowledge_base associations)
+	if err := h.service.Repository().HardDeleteKnowledgeBase(ctx, kb.UID.String()); err != nil {
+		logger.Error("failed to delete knowledge base", zap.Error(err))
+		return nil, status.Errorf(codes.Internal, "failed to delete knowledge base: %v", err)
+	}
+
+	logger.Info("DeleteKnowledgeBaseAdmin completed",
+		zap.String("kb_uid", kb.UID.String()),
+		zap.String("kb_id", kb.ID))
+
+	return &artifactpb.DeleteKnowledgeBaseAdminResponse{}, nil
+}

pkg/repository/file.go

@@ -96,6 +96,11 @@ type File interface {
 	// Returns the number of files affected.
 	ResetFileStatusesByKBUID(ctx context.Context, kbUID types.KBUIDType) (int64, error)
 
+	// AddFilesToKnowledgeBase adds file associations to a target KB by file IDs.
+	// Files can belong to multiple KBs (many-to-many relationship).
+	// Returns the number of new associations added (skips duplicates).
+	AddFilesToKnowledgeBase(ctx context.Context, targetKBUID types.KnowledgeBaseUIDType, fileIDs []string) (int64, error)
+
 	// Deprecated methods
 
 	// GetFileByKBUIDAndFileID returns the file by
@@ -643,6 +648,33 @@ func (r *repository) ResetFileStatusesByKBUID(ctx context.Context, kbUID types.K
 	return result.RowsAffected, nil
 }
 
+// AddFilesToKnowledgeBase adds file associations to a target KB by file IDs.
+// Files can belong to multiple KBs (many-to-many relationship).
+// Returns the number of new associations added (skips duplicates).
+func (r *repository) AddFilesToKnowledgeBase(ctx context.Context, targetKBUID types.KnowledgeBaseUIDType, fileIDs []string) (int64, error) {
+	if len(fileIDs) == 0 {
+		return 0, nil
+	}
+
+	// Use raw SQL to:
+	// 1. Look up file UIDs from the file table by their IDs
+	// 2. Insert into junction table, skipping duplicates
+	sql := `
+		INSERT INTO file_knowledge_base (file_uid, kb_uid, created_at)
+		SELECT f.uid, ?, NOW()
+		FROM file f
+		WHERE f.id = ANY(?)
+		ON CONFLICT (file_uid, kb_uid) DO NOTHING
+	`
+
+	result := r.db.WithContext(ctx).Exec(sql, targetKBUID, fileIDs)
+	if result.Error != nil {
+		return 0, fmt.Errorf("failed to add files to knowledge base: %w", result.Error)
+	}
+
+	return result.RowsAffected, nil
+}
+
 // ProcessFiles updates the process status of the files
 func (r *repository) ProcessFiles(
 	ctx context.Context,

pkg/repository/knowledge_base.go

@@ -100,6 +100,10 @@ type KnowledgeBase interface {
 	// DeleteKnowledgeBaseTx soft-deletes a KB within a transaction
 	DeleteKnowledgeBaseTx(ctx context.Context, tx *gorm.DB, owner, kbID string) error
 
+	// HardDeleteKnowledgeBase permanently deletes a KB and CASCADE removes file_knowledge_base associations
+	// Used by admin consolidation operations to remove duplicate KBs after moving files
+	HardDeleteKnowledgeBase(ctx context.Context, kbUID string) error
+
 	// UpdateKnowledgeBaseResources updates kb_uid references in all resource tables
 	// This is critical for atomic swap to ensure resources follow their knowledge bases
 	UpdateKnowledgeBaseResources(ctx context.Context, fromKBUID, toKBUID types.KBUIDType) error
@@ -589,6 +593,20 @@ func (r *repository) DeleteKnowledgeBaseTx(ctx context.Context, tx *gorm.DB, own
 	return nil
 }
 
+// HardDeleteKnowledgeBase permanently deletes a KB and CASCADE removes file_knowledge_base associations
+// Used by admin consolidation operations to remove duplicate KBs after moving files
+func (r *repository) HardDeleteKnowledgeBase(ctx context.Context, kbUID string) error {
+	// Use Unscoped to bypass soft delete and perform actual DELETE
+	result := r.db.WithContext(ctx).Unscoped().Delete(&KnowledgeBaseModel{}, "uid = ?", kbUID)
+	if result.Error != nil {
+		return fmt.Errorf("hard deleting knowledge base: %w", result.Error)
+	}
+	if result.RowsAffected == 0 {
+		return fmt.Errorf("knowledge base not found: %s", kbUID)
+	}
+	return nil
+}
+
 // UpdateKBUIDInResources updates kb_uid in all resource tables (files, chunks, embeddings, converted_files)
 // This is CRITICAL for atomic swap: when KBs are swapped, all resources must follow their KBs.
 // Without this, queries will fail because resources point to old KB UIDs.