From fc504add9a3ad467223ed808c7110e68d97d8972 Mon Sep 17 00:00:00 2001 From: "gary.y" Date: Thu, 28 Nov 2024 00:19:30 +0800 Subject: [PATCH] feat(artifact): support ephemeral catalog --- config/config.yaml | 2 +- go.mod | 2 +- go.sum | 4 +- .../000019_refactor_kb_for_chat.down.sql | 4 + .../000019_refactor_kb_for_chat.up.sql | 10 + pkg/handler/knowledgebase.go | 87 ++-- pkg/mock/repository_i_mock.gen.go | 410 +++++++++++++++++- pkg/repository/knowledgebase.go | 26 +- 8 files changed, 486 insertions(+), 59 deletions(-) create mode 100644 pkg/db/migration/000019_refactor_kb_for_chat.down.sql create mode 100644 pkg/db/migration/000019_refactor_kb_for_chat.up.sql diff --git a/config/config.yaml b/config/config.yaml index e6cddb4..39762fb 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -22,7 +22,7 @@ database: host: pg-sql port: 5432 name: artifact - version: 18 + version: 19 timezone: Etc/UTC pool: idleconnections: 5 diff --git a/go.mod b/go.mod index cf347b9..e1ac530 100644 --- a/go.mod +++ b/go.mod @@ -13,7 +13,7 @@ require ( github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1 github.com/influxdata/influxdb-client-go/v2 v2.12.3 - github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241029162707-1398399a24ee + github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241125163328-c29704e47ba4 github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61 github.com/instill-ai/x v0.3.0-alpha.0.20231219052200-6230a89e386c github.com/knadh/koanf v1.5.0 diff --git a/go.sum b/go.sum index 00e78e5..a78749a 100644 --- a/go.sum +++ b/go.sum @@ -344,8 +344,8 @@ github.com/influxdata/influxdb-client-go/v2 v2.12.3 h1:28nRlNMRIV4QbtIUvxhWqaxn0 github.com/influxdata/influxdb-client-go/v2 v2.12.3/go.mod h1:IrrLUbCjjfkmRuaCiGQg4m2GbkaeJDcuWoxiWdQEbA0= github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839 h1:W9WBk7wlPfJLvMCdtV4zPulc4uCPrlywQOmbFOhgQNU= github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839/go.mod h1:xaLFMmpvUxqXtVkUJfg9QmT88cDaCJ3ZKgdZ78oO8Qo= -github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241029162707-1398399a24ee h1:onnzrn5jabO3jDLPo2193Ql6YMRyDWDx9K834Bfi8V0= -github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241029162707-1398399a24ee/go.mod h1:rf0UY7VpEgpaLudYEcjx5rnbuwlBaaLyD4FQmWLtgAY= +github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241125163328-c29704e47ba4 h1:k8X9gMiCwHWShB1FITaWwmlzthFnor1Jj0tSaFG+9x8= +github.com/instill-ai/protogen-go v0.3.3-alpha.0.20241125163328-c29704e47ba4/go.mod h1:rf0UY7VpEgpaLudYEcjx5rnbuwlBaaLyD4FQmWLtgAY= github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61 h1:smPTvmXDhn/QC7y/TPXyMTqbbRd0gvzmFgWBChwTfhE= github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61/go.mod h1:/TAHs4ybuylk5icuy+MQtHRc4XUnIyXzeNKxX9qDFhw= github.com/instill-ai/x v0.3.0-alpha.0.20231219052200-6230a89e386c h1:a2RVkpIV2QcrGnSHAou+t/L+vBsaIfFvk5inVg5Uh4s= diff --git a/pkg/db/migration/000019_refactor_kb_for_chat.down.sql b/pkg/db/migration/000019_refactor_kb_for_chat.down.sql new file mode 100644 index 0000000..744f042 --- /dev/null +++ b/pkg/db/migration/000019_refactor_kb_for_chat.down.sql @@ -0,0 +1,4 @@ +BEGIN; +-- Remove catalog_type column from knowledge_base table +ALTER TABLE knowledge_base DROP COLUMN catalog_type; +COMMIT; diff --git a/pkg/db/migration/000019_refactor_kb_for_chat.up.sql b/pkg/db/migration/000019_refactor_kb_for_chat.up.sql new file mode 100644 index 0000000..981b229 --- /dev/null +++ b/pkg/db/migration/000019_refactor_kb_for_chat.up.sql @@ -0,0 +1,10 @@ +BEGIN; + +-- Add catalog_type column to knowledge_base table +ALTER TABLE knowledge_base +ADD COLUMN catalog_type VARCHAR(255) DEFAULT 'CATALOG_TYPE_PERSISTENT'; + +-- Add comment for the new column +COMMENT ON COLUMN knowledge_base.catalog_type IS 'Type of the knowledge base catalog'; + +COMMIT; diff --git a/pkg/handler/knowledgebase.go b/pkg/handler/knowledgebase.go index 64195e4..dba916f 100644 --- a/pkg/handler/knowledgebase.go +++ b/pkg/handler/knowledgebase.go @@ -59,7 +59,7 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea // check if user has reached the maximum number of catalogs // note: the simple implementation have race condition to bypass the check, // but it is okay for now - kbCount, err := ph.service.Repository.GetKnowledgeBaseCountByOwner(ctx, ns.NsUID.String()) + kbCount, err := ph.service.Repository.GetKnowledgeBaseCountByOwner(ctx, ns.NsUID.String(), artifactpb.CatalogType_CATALOG_TYPE_PERSISTENT) if err != nil { log.Error("failed to get catalog count", zap.Error(err)) return nil, fmt.Errorf(ErrorCreateKnowledgeBaseMsg, err) @@ -117,6 +117,11 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea return nil } + // if catalog type is not set, set it to persistent + if req.GetType() == artifactpb.CatalogType_CATALOG_TYPE_UNSPECIFIED { + req.Type = artifactpb.CatalogType_CATALOG_TYPE_PERSISTENT + } + // create catalog dbData, err := ph.service.Repository.CreateKnowledgeBase(ctx, repository.KnowledgeBase{ @@ -127,6 +132,7 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea Tags: req.Tags, Owner: ns.NsUID.String(), CreatorUID: creatorUUID, + CatalogType: req.GetType().String(), }, callExternalService, ) if err != nil { @@ -135,20 +141,29 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea return &artifactpb.CreateCatalogResponse{ Catalog: &artifactpb.Catalog{ - Name: dbData.Name, - CatalogId: dbData.KbID, - Description: dbData.Description, - Tags: dbData.Tags, - OwnerName: dbData.Owner, - CreateTime: dbData.CreateTime.String(), - UpdateTime: dbData.UpdateTime.String(), - ConvertingPipelines: []string{"preset/indexing-convert-pdf"}, - SplittingPipelines: []string{"preset/indexing-split-text", "preset/indexing-split-markdown"}, - EmbeddingPipelines: []string{"preset/indexing-embed"}, - DownstreamApps: []string{}, - TotalFiles: 0, - TotalTokens: 0, - UsedStorage: 0, + Name: dbData.Name, + CatalogUid: dbData.UID.String(), + CatalogId: dbData.KbID, + Description: dbData.Description, + Tags: dbData.Tags, + OwnerName: dbData.Owner, + CreateTime: dbData.CreateTime.String(), + UpdateTime: dbData.UpdateTime.String(), + ConvertingPipelines: []string{ + service.NamespaceID + "/" + service.ConvertDocToMDPipelineID + "@" + service.DocToMDVersion, + service.NamespaceID + "/" + service.ConvertDocToMDPipelineID2 + "@" + service.DocToMDVersion2, + }, + SplittingPipelines: []string{ + service.NamespaceID + "/" + service.TextChunkPipelineID + "@" + service.TextSplitVersion, + service.NamespaceID + "/" + service.MdChunkPipelineID + "@" + service.MdSplitVersion, + }, + EmbeddingPipelines: []string{ + service.NamespaceID + "/" + service.TextEmbedPipelineID + "@" + service.TextEmbedVersion, + }, + DownstreamApps: []string{}, + TotalFiles: 0, + TotalTokens: 0, + UsedStorage: 0, }, }, nil } @@ -183,7 +198,7 @@ func (ph *PublicHandler) ListCatalogs(ctx context.Context, req *artifactpb.ListC return nil, fmt.Errorf("failed to check namespace permission. err:%w", err) } - dbData, err := ph.service.Repository.ListKnowledgeBases(ctx, ns.NsUID.String()) + dbData, err := ph.service.Repository.ListKnowledgeBasesByCatalogType(ctx, ns.NsUID.String(), artifactpb.CatalogType_CATALOG_TYPE_PERSISTENT) if err != nil { log.Error("failed to get catalogs", zap.Error(err)) return nil, fmt.Errorf(ErrorListKnowledgeBasesMsg, err) @@ -207,23 +222,29 @@ func (ph *PublicHandler) ListCatalogs(ctx context.Context, req *artifactpb.ListC kbs := make([]*artifactpb.Catalog, len(dbData)) for i, kb := range dbData { kbs[i] = &artifactpb.Catalog{ - CatalogUid: kb.UID.String(), - Name: kb.Name, - CatalogId: kb.KbID, - Description: kb.Description, - Tags: kb.Tags, - CreateTime: kb.CreateTime.String(), - UpdateTime: kb.UpdateTime.String(), - OwnerName: kb.Owner, - ConvertingPipelines: []string{service.NamespaceID + "/" + service.ConvertDocToMDPipelineID}, + CatalogUid: kb.UID.String(), + Name: kb.Name, + CatalogId: kb.KbID, + Description: kb.Description, + Tags: kb.Tags, + CreateTime: kb.CreateTime.String(), + UpdateTime: kb.UpdateTime.String(), + OwnerName: kb.Owner, + ConvertingPipelines: []string{ + service.NamespaceID + "/" + service.ConvertDocToMDPipelineID + "@" + service.DocToMDVersion, + service.NamespaceID + "/" + service.ConvertDocToMDPipelineID2 + "@" + service.DocToMDVersion2, + }, SplittingPipelines: []string{ - service.NamespaceID + "/" + service.TextChunkPipelineID, - service.NamespaceID + "/" + service.MdChunkPipelineID}, - EmbeddingPipelines: []string{service.NamespaceID + "/" + service.TextEmbedPipelineID}, - DownstreamApps: []string{}, - TotalFiles: uint32(fileCounts[kb.UID]), - TotalTokens: uint32(tokenCounts[kb.UID]), - UsedStorage: uint64(kb.Usage), + service.NamespaceID + "/" + service.TextChunkPipelineID + "@" + service.TextSplitVersion, + service.NamespaceID + "/" + service.MdChunkPipelineID + "@" + service.MdSplitVersion, + }, + EmbeddingPipelines: []string{ + service.NamespaceID + "/" + service.TextEmbedPipelineID + "@" + service.TextEmbedVersion, + }, + DownstreamApps: []string{}, + TotalFiles: uint32(fileCounts[kb.UID]), + TotalTokens: uint32(tokenCounts[kb.UID]), + UsedStorage: uint64(kb.Usage), } } return &artifactpb.ListCatalogsResponse{ @@ -241,7 +262,7 @@ func (ph *PublicHandler) UpdateCatalog(ctx context.Context, req *artifactpb.Upda if req.CatalogId == "" { log.Error("kb_id is empty", zap.Error(ErrCheckRequiredFields)) return nil, fmt.Errorf("kb_id is empty. err: %w", ErrCheckRequiredFields) - } +} ns, err := ph.service.GetNamespaceByNsID(ctx, req.GetNamespaceId()) if err != nil { diff --git a/pkg/mock/repository_i_mock.gen.go b/pkg/mock/repository_i_mock.gen.go index 64fcaff..f560272 100644 --- a/pkg/mock/repository_i_mock.gen.go +++ b/pkg/mock/repository_i_mock.gen.go @@ -13,6 +13,7 @@ import ( mm_repository "github.com/instill-ai/artifact-backend/pkg/repository" "github.com/instill-ai/artifact-backend/pkg/utils" "gorm.io/gorm" + artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha" pb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha" ) @@ -183,8 +184,8 @@ type RepositoryIMock struct { beforeGetKnowledgeBaseByOwnerAndKbIDCounter uint64 GetKnowledgeBaseByOwnerAndKbIDMock mRepositoryIMockGetKnowledgeBaseByOwnerAndKbID - funcGetKnowledgeBaseCountByOwner func(ctx context.Context, ownerUID string) (i1 int64, err error) - inspectFuncGetKnowledgeBaseCountByOwner func(ctx context.Context, ownerUID string) + funcGetKnowledgeBaseCountByOwner func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (i1 int64, err error) + inspectFuncGetKnowledgeBaseCountByOwner func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) afterGetKnowledgeBaseCountByOwnerCounter uint64 beforeGetKnowledgeBaseCountByOwnerCounter uint64 GetKnowledgeBaseCountByOwnerMock mRepositoryIMockGetKnowledgeBaseCountByOwner @@ -366,6 +367,12 @@ type RepositoryIMock struct { beforeListKnowledgeBasesCounter uint64 ListKnowledgeBasesMock mRepositoryIMockListKnowledgeBases + funcListKnowledgeBasesByCatalogType func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (ka1 []mm_repository.KnowledgeBase, err error) + inspectFuncListKnowledgeBasesByCatalogType func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) + afterListKnowledgeBasesByCatalogTypeCounter uint64 + beforeListKnowledgeBasesByCatalogTypeCounter uint64 + ListKnowledgeBasesByCatalogTypeMock mRepositoryIMockListKnowledgeBasesByCatalogType + funcProcessKnowledgeBaseFiles func(ctx context.Context, fileUIDs []string, requester uuid.UUID) (ka1 []mm_repository.KnowledgeBaseFile, err error) inspectFuncProcessKnowledgeBaseFiles func(ctx context.Context, fileUIDs []string, requester uuid.UUID) afterProcessKnowledgeBaseFilesCounter uint64 @@ -604,6 +611,9 @@ func NewRepositoryIMock(t minimock.Tester) *RepositoryIMock { m.ListKnowledgeBasesMock = mRepositoryIMockListKnowledgeBases{mock: m} m.ListKnowledgeBasesMock.callArgs = []*RepositoryIMockListKnowledgeBasesParams{} + m.ListKnowledgeBasesByCatalogTypeMock = mRepositoryIMockListKnowledgeBasesByCatalogType{mock: m} + m.ListKnowledgeBasesByCatalogTypeMock.callArgs = []*RepositoryIMockListKnowledgeBasesByCatalogTypeParams{} + m.ProcessKnowledgeBaseFilesMock = mRepositoryIMockProcessKnowledgeBaseFiles{mock: m} m.ProcessKnowledgeBaseFilesMock.callArgs = []*RepositoryIMockProcessKnowledgeBaseFilesParams{} @@ -8743,14 +8753,16 @@ type RepositoryIMockGetKnowledgeBaseCountByOwnerExpectation struct { // RepositoryIMockGetKnowledgeBaseCountByOwnerParams contains parameters of the RepositoryI.GetKnowledgeBaseCountByOwner type RepositoryIMockGetKnowledgeBaseCountByOwnerParams struct { - ctx context.Context - ownerUID string + ctx context.Context + ownerUID string + catalogType artifactpb.CatalogType } // RepositoryIMockGetKnowledgeBaseCountByOwnerParamPtrs contains pointers to parameters of the RepositoryI.GetKnowledgeBaseCountByOwner type RepositoryIMockGetKnowledgeBaseCountByOwnerParamPtrs struct { - ctx *context.Context - ownerUID *string + ctx *context.Context + ownerUID *string + catalogType *artifactpb.CatalogType } // RepositoryIMockGetKnowledgeBaseCountByOwnerResults contains results of the RepositoryI.GetKnowledgeBaseCountByOwner @@ -8760,7 +8772,7 @@ type RepositoryIMockGetKnowledgeBaseCountByOwnerResults struct { } // Expect sets up expected params for RepositoryI.GetKnowledgeBaseCountByOwner -func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Expect(ctx context.Context, ownerUID string) *mRepositoryIMockGetKnowledgeBaseCountByOwner { +func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Expect(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) *mRepositoryIMockGetKnowledgeBaseCountByOwner { if mmGetKnowledgeBaseCountByOwner.mock.funcGetKnowledgeBaseCountByOwner != nil { mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("RepositoryIMock.GetKnowledgeBaseCountByOwner mock is already set by Set") } @@ -8773,7 +8785,7 @@ func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwn mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("RepositoryIMock.GetKnowledgeBaseCountByOwner mock is already set by ExpectParams functions") } - mmGetKnowledgeBaseCountByOwner.defaultExpectation.params = &RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID} + mmGetKnowledgeBaseCountByOwner.defaultExpectation.params = &RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID, catalogType} for _, e := range mmGetKnowledgeBaseCountByOwner.expectations { if minimock.Equal(e.params, mmGetKnowledgeBaseCountByOwner.defaultExpectation.params) { mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("Expectation set by When has same params: %#v", *mmGetKnowledgeBaseCountByOwner.defaultExpectation.params) @@ -8827,8 +8839,30 @@ func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwn return mmGetKnowledgeBaseCountByOwner } +// ExpectCatalogTypeParam3 sets up expected param catalogType for RepositoryI.GetKnowledgeBaseCountByOwner +func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) ExpectCatalogTypeParam3(catalogType artifactpb.CatalogType) *mRepositoryIMockGetKnowledgeBaseCountByOwner { + if mmGetKnowledgeBaseCountByOwner.mock.funcGetKnowledgeBaseCountByOwner != nil { + mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("RepositoryIMock.GetKnowledgeBaseCountByOwner mock is already set by Set") + } + + if mmGetKnowledgeBaseCountByOwner.defaultExpectation == nil { + mmGetKnowledgeBaseCountByOwner.defaultExpectation = &RepositoryIMockGetKnowledgeBaseCountByOwnerExpectation{} + } + + if mmGetKnowledgeBaseCountByOwner.defaultExpectation.params != nil { + mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("RepositoryIMock.GetKnowledgeBaseCountByOwner mock is already set by Expect") + } + + if mmGetKnowledgeBaseCountByOwner.defaultExpectation.paramPtrs == nil { + mmGetKnowledgeBaseCountByOwner.defaultExpectation.paramPtrs = &RepositoryIMockGetKnowledgeBaseCountByOwnerParamPtrs{} + } + mmGetKnowledgeBaseCountByOwner.defaultExpectation.paramPtrs.catalogType = &catalogType + + return mmGetKnowledgeBaseCountByOwner +} + // Inspect accepts an inspector function that has same arguments as the RepositoryI.GetKnowledgeBaseCountByOwner -func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Inspect(f func(ctx context.Context, ownerUID string)) *mRepositoryIMockGetKnowledgeBaseCountByOwner { +func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Inspect(f func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType)) *mRepositoryIMockGetKnowledgeBaseCountByOwner { if mmGetKnowledgeBaseCountByOwner.mock.inspectFuncGetKnowledgeBaseCountByOwner != nil { mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("Inspect function is already set for RepositoryIMock.GetKnowledgeBaseCountByOwner") } @@ -8852,7 +8886,7 @@ func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwn } // Set uses given function f to mock the RepositoryI.GetKnowledgeBaseCountByOwner method -func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Set(f func(ctx context.Context, ownerUID string) (i1 int64, err error)) *RepositoryIMock { +func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) Set(f func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (i1 int64, err error)) *RepositoryIMock { if mmGetKnowledgeBaseCountByOwner.defaultExpectation != nil { mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("Default expectation is already set for the RepositoryI.GetKnowledgeBaseCountByOwner method") } @@ -8867,14 +8901,14 @@ func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwn // When sets expectation for the RepositoryI.GetKnowledgeBaseCountByOwner which will trigger the result defined by the following // Then helper -func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) When(ctx context.Context, ownerUID string) *RepositoryIMockGetKnowledgeBaseCountByOwnerExpectation { +func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwner) When(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) *RepositoryIMockGetKnowledgeBaseCountByOwnerExpectation { if mmGetKnowledgeBaseCountByOwner.mock.funcGetKnowledgeBaseCountByOwner != nil { mmGetKnowledgeBaseCountByOwner.mock.t.Fatalf("RepositoryIMock.GetKnowledgeBaseCountByOwner mock is already set by Set") } expectation := &RepositoryIMockGetKnowledgeBaseCountByOwnerExpectation{ mock: mmGetKnowledgeBaseCountByOwner.mock, - params: &RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID}, + params: &RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID, catalogType}, } mmGetKnowledgeBaseCountByOwner.expectations = append(mmGetKnowledgeBaseCountByOwner.expectations, expectation) return expectation @@ -8907,15 +8941,15 @@ func (mmGetKnowledgeBaseCountByOwner *mRepositoryIMockGetKnowledgeBaseCountByOwn } // GetKnowledgeBaseCountByOwner implements repository.RepositoryI -func (mmGetKnowledgeBaseCountByOwner *RepositoryIMock) GetKnowledgeBaseCountByOwner(ctx context.Context, ownerUID string) (i1 int64, err error) { +func (mmGetKnowledgeBaseCountByOwner *RepositoryIMock) GetKnowledgeBaseCountByOwner(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (i1 int64, err error) { mm_atomic.AddUint64(&mmGetKnowledgeBaseCountByOwner.beforeGetKnowledgeBaseCountByOwnerCounter, 1) defer mm_atomic.AddUint64(&mmGetKnowledgeBaseCountByOwner.afterGetKnowledgeBaseCountByOwnerCounter, 1) if mmGetKnowledgeBaseCountByOwner.inspectFuncGetKnowledgeBaseCountByOwner != nil { - mmGetKnowledgeBaseCountByOwner.inspectFuncGetKnowledgeBaseCountByOwner(ctx, ownerUID) + mmGetKnowledgeBaseCountByOwner.inspectFuncGetKnowledgeBaseCountByOwner(ctx, ownerUID, catalogType) } - mm_params := RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID} + mm_params := RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID, catalogType} // Record call args mmGetKnowledgeBaseCountByOwner.GetKnowledgeBaseCountByOwnerMock.mutex.Lock() @@ -8934,7 +8968,7 @@ func (mmGetKnowledgeBaseCountByOwner *RepositoryIMock) GetKnowledgeBaseCountByOw mm_want := mmGetKnowledgeBaseCountByOwner.GetKnowledgeBaseCountByOwnerMock.defaultExpectation.params mm_want_ptrs := mmGetKnowledgeBaseCountByOwner.GetKnowledgeBaseCountByOwnerMock.defaultExpectation.paramPtrs - mm_got := RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID} + mm_got := RepositoryIMockGetKnowledgeBaseCountByOwnerParams{ctx, ownerUID, catalogType} if mm_want_ptrs != nil { @@ -8946,6 +8980,10 @@ func (mmGetKnowledgeBaseCountByOwner *RepositoryIMock) GetKnowledgeBaseCountByOw mmGetKnowledgeBaseCountByOwner.t.Errorf("RepositoryIMock.GetKnowledgeBaseCountByOwner got unexpected parameter ownerUID, want: %#v, got: %#v%s\n", *mm_want_ptrs.ownerUID, mm_got.ownerUID, minimock.Diff(*mm_want_ptrs.ownerUID, mm_got.ownerUID)) } + if mm_want_ptrs.catalogType != nil && !minimock.Equal(*mm_want_ptrs.catalogType, mm_got.catalogType) { + mmGetKnowledgeBaseCountByOwner.t.Errorf("RepositoryIMock.GetKnowledgeBaseCountByOwner got unexpected parameter catalogType, want: %#v, got: %#v%s\n", *mm_want_ptrs.catalogType, mm_got.catalogType, minimock.Diff(*mm_want_ptrs.catalogType, mm_got.catalogType)) + } + } else if mm_want != nil && !minimock.Equal(*mm_want, mm_got) { mmGetKnowledgeBaseCountByOwner.t.Errorf("RepositoryIMock.GetKnowledgeBaseCountByOwner got unexpected parameters, want: %#v, got: %#v%s\n", *mm_want, mm_got, minimock.Diff(*mm_want, mm_got)) } @@ -8957,9 +8995,9 @@ func (mmGetKnowledgeBaseCountByOwner *RepositoryIMock) GetKnowledgeBaseCountByOw return (*mm_results).i1, (*mm_results).err } if mmGetKnowledgeBaseCountByOwner.funcGetKnowledgeBaseCountByOwner != nil { - return mmGetKnowledgeBaseCountByOwner.funcGetKnowledgeBaseCountByOwner(ctx, ownerUID) + return mmGetKnowledgeBaseCountByOwner.funcGetKnowledgeBaseCountByOwner(ctx, ownerUID, catalogType) } - mmGetKnowledgeBaseCountByOwner.t.Fatalf("Unexpected call to RepositoryIMock.GetKnowledgeBaseCountByOwner. %v %v", ctx, ownerUID) + mmGetKnowledgeBaseCountByOwner.t.Fatalf("Unexpected call to RepositoryIMock.GetKnowledgeBaseCountByOwner. %v %v %v", ctx, ownerUID, catalogType) return } @@ -17766,6 +17804,339 @@ func (m *RepositoryIMock) MinimockListKnowledgeBasesInspect() { } } +type mRepositoryIMockListKnowledgeBasesByCatalogType struct { + mock *RepositoryIMock + defaultExpectation *RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation + expectations []*RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation + + callArgs []*RepositoryIMockListKnowledgeBasesByCatalogTypeParams + mutex sync.RWMutex + + expectedInvocations uint64 +} + +// RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation specifies expectation struct of the RepositoryI.ListKnowledgeBasesByCatalogType +type RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation struct { + mock *RepositoryIMock + params *RepositoryIMockListKnowledgeBasesByCatalogTypeParams + paramPtrs *RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs + results *RepositoryIMockListKnowledgeBasesByCatalogTypeResults + Counter uint64 +} + +// RepositoryIMockListKnowledgeBasesByCatalogTypeParams contains parameters of the RepositoryI.ListKnowledgeBasesByCatalogType +type RepositoryIMockListKnowledgeBasesByCatalogTypeParams struct { + ctx context.Context + ownerUID string + catalogType artifactpb.CatalogType +} + +// RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs contains pointers to parameters of the RepositoryI.ListKnowledgeBasesByCatalogType +type RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs struct { + ctx *context.Context + ownerUID *string + catalogType *artifactpb.CatalogType +} + +// RepositoryIMockListKnowledgeBasesByCatalogTypeResults contains results of the RepositoryI.ListKnowledgeBasesByCatalogType +type RepositoryIMockListKnowledgeBasesByCatalogTypeResults struct { + ka1 []mm_repository.KnowledgeBase + err error +} + +// Expect sets up expected params for RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Expect(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation = &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{} + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by ExpectParams functions") + } + + mmListKnowledgeBasesByCatalogType.defaultExpectation.params = &RepositoryIMockListKnowledgeBasesByCatalogTypeParams{ctx, ownerUID, catalogType} + for _, e := range mmListKnowledgeBasesByCatalogType.expectations { + if minimock.Equal(e.params, mmListKnowledgeBasesByCatalogType.defaultExpectation.params) { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("Expectation set by When has same params: %#v", *mmListKnowledgeBasesByCatalogType.defaultExpectation.params) + } + } + + return mmListKnowledgeBasesByCatalogType +} + +// ExpectCtxParam1 sets up expected param ctx for RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) ExpectCtxParam1(ctx context.Context) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation = &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{} + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.params != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Expect") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs = &RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs{} + } + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs.ctx = &ctx + + return mmListKnowledgeBasesByCatalogType +} + +// ExpectOwnerUIDParam2 sets up expected param ownerUID for RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) ExpectOwnerUIDParam2(ownerUID string) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation = &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{} + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.params != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Expect") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs = &RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs{} + } + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs.ownerUID = &ownerUID + + return mmListKnowledgeBasesByCatalogType +} + +// ExpectCatalogTypeParam3 sets up expected param catalogType for RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) ExpectCatalogTypeParam3(catalogType artifactpb.CatalogType) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation = &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{} + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.params != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Expect") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs = &RepositoryIMockListKnowledgeBasesByCatalogTypeParamPtrs{} + } + mmListKnowledgeBasesByCatalogType.defaultExpectation.paramPtrs.catalogType = &catalogType + + return mmListKnowledgeBasesByCatalogType +} + +// Inspect accepts an inspector function that has same arguments as the RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Inspect(f func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType)) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if mmListKnowledgeBasesByCatalogType.mock.inspectFuncListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("Inspect function is already set for RepositoryIMock.ListKnowledgeBasesByCatalogType") + } + + mmListKnowledgeBasesByCatalogType.mock.inspectFuncListKnowledgeBasesByCatalogType = f + + return mmListKnowledgeBasesByCatalogType +} + +// Return sets up results that will be returned by RepositoryI.ListKnowledgeBasesByCatalogType +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Return(ka1 []mm_repository.KnowledgeBase, err error) *RepositoryIMock { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + if mmListKnowledgeBasesByCatalogType.defaultExpectation == nil { + mmListKnowledgeBasesByCatalogType.defaultExpectation = &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{mock: mmListKnowledgeBasesByCatalogType.mock} + } + mmListKnowledgeBasesByCatalogType.defaultExpectation.results = &RepositoryIMockListKnowledgeBasesByCatalogTypeResults{ka1, err} + return mmListKnowledgeBasesByCatalogType.mock +} + +// Set uses given function f to mock the RepositoryI.ListKnowledgeBasesByCatalogType method +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Set(f func(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (ka1 []mm_repository.KnowledgeBase, err error)) *RepositoryIMock { + if mmListKnowledgeBasesByCatalogType.defaultExpectation != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("Default expectation is already set for the RepositoryI.ListKnowledgeBasesByCatalogType method") + } + + if len(mmListKnowledgeBasesByCatalogType.expectations) > 0 { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("Some expectations are already set for the RepositoryI.ListKnowledgeBasesByCatalogType method") + } + + mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType = f + return mmListKnowledgeBasesByCatalogType.mock +} + +// When sets expectation for the RepositoryI.ListKnowledgeBasesByCatalogType which will trigger the result defined by the following +// Then helper +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) When(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) *RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation { + if mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("RepositoryIMock.ListKnowledgeBasesByCatalogType mock is already set by Set") + } + + expectation := &RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation{ + mock: mmListKnowledgeBasesByCatalogType.mock, + params: &RepositoryIMockListKnowledgeBasesByCatalogTypeParams{ctx, ownerUID, catalogType}, + } + mmListKnowledgeBasesByCatalogType.expectations = append(mmListKnowledgeBasesByCatalogType.expectations, expectation) + return expectation +} + +// Then sets up RepositoryI.ListKnowledgeBasesByCatalogType return parameters for the expectation previously defined by the When method +func (e *RepositoryIMockListKnowledgeBasesByCatalogTypeExpectation) Then(ka1 []mm_repository.KnowledgeBase, err error) *RepositoryIMock { + e.results = &RepositoryIMockListKnowledgeBasesByCatalogTypeResults{ka1, err} + return e.mock +} + +// Times sets number of times RepositoryI.ListKnowledgeBasesByCatalogType should be invoked +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Times(n uint64) *mRepositoryIMockListKnowledgeBasesByCatalogType { + if n == 0 { + mmListKnowledgeBasesByCatalogType.mock.t.Fatalf("Times of RepositoryIMock.ListKnowledgeBasesByCatalogType mock can not be zero") + } + mm_atomic.StoreUint64(&mmListKnowledgeBasesByCatalogType.expectedInvocations, n) + return mmListKnowledgeBasesByCatalogType +} + +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) invocationsDone() bool { + if len(mmListKnowledgeBasesByCatalogType.expectations) == 0 && mmListKnowledgeBasesByCatalogType.defaultExpectation == nil && mmListKnowledgeBasesByCatalogType.mock.funcListKnowledgeBasesByCatalogType == nil { + return true + } + + totalInvocations := mm_atomic.LoadUint64(&mmListKnowledgeBasesByCatalogType.mock.afterListKnowledgeBasesByCatalogTypeCounter) + expectedInvocations := mm_atomic.LoadUint64(&mmListKnowledgeBasesByCatalogType.expectedInvocations) + + return totalInvocations > 0 && (expectedInvocations == 0 || expectedInvocations == totalInvocations) +} + +// ListKnowledgeBasesByCatalogType implements repository.RepositoryI +func (mmListKnowledgeBasesByCatalogType *RepositoryIMock) ListKnowledgeBasesByCatalogType(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (ka1 []mm_repository.KnowledgeBase, err error) { + mm_atomic.AddUint64(&mmListKnowledgeBasesByCatalogType.beforeListKnowledgeBasesByCatalogTypeCounter, 1) + defer mm_atomic.AddUint64(&mmListKnowledgeBasesByCatalogType.afterListKnowledgeBasesByCatalogTypeCounter, 1) + + if mmListKnowledgeBasesByCatalogType.inspectFuncListKnowledgeBasesByCatalogType != nil { + mmListKnowledgeBasesByCatalogType.inspectFuncListKnowledgeBasesByCatalogType(ctx, ownerUID, catalogType) + } + + mm_params := RepositoryIMockListKnowledgeBasesByCatalogTypeParams{ctx, ownerUID, catalogType} + + // Record call args + mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.mutex.Lock() + mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.callArgs = append(mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.callArgs, &mm_params) + mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.mutex.Unlock() + + for _, e := range mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.expectations { + if minimock.Equal(*e.params, mm_params) { + mm_atomic.AddUint64(&e.Counter, 1) + return e.results.ka1, e.results.err + } + } + + if mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation != nil { + mm_atomic.AddUint64(&mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.Counter, 1) + mm_want := mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.params + mm_want_ptrs := mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.paramPtrs + + mm_got := RepositoryIMockListKnowledgeBasesByCatalogTypeParams{ctx, ownerUID, catalogType} + + if mm_want_ptrs != nil { + + if mm_want_ptrs.ctx != nil && !minimock.Equal(*mm_want_ptrs.ctx, mm_got.ctx) { + mmListKnowledgeBasesByCatalogType.t.Errorf("RepositoryIMock.ListKnowledgeBasesByCatalogType got unexpected parameter ctx, want: %#v, got: %#v%s\n", *mm_want_ptrs.ctx, mm_got.ctx, minimock.Diff(*mm_want_ptrs.ctx, mm_got.ctx)) + } + + if mm_want_ptrs.ownerUID != nil && !minimock.Equal(*mm_want_ptrs.ownerUID, mm_got.ownerUID) { + mmListKnowledgeBasesByCatalogType.t.Errorf("RepositoryIMock.ListKnowledgeBasesByCatalogType got unexpected parameter ownerUID, want: %#v, got: %#v%s\n", *mm_want_ptrs.ownerUID, mm_got.ownerUID, minimock.Diff(*mm_want_ptrs.ownerUID, mm_got.ownerUID)) + } + + if mm_want_ptrs.catalogType != nil && !minimock.Equal(*mm_want_ptrs.catalogType, mm_got.catalogType) { + mmListKnowledgeBasesByCatalogType.t.Errorf("RepositoryIMock.ListKnowledgeBasesByCatalogType got unexpected parameter catalogType, want: %#v, got: %#v%s\n", *mm_want_ptrs.catalogType, mm_got.catalogType, minimock.Diff(*mm_want_ptrs.catalogType, mm_got.catalogType)) + } + + } else if mm_want != nil && !minimock.Equal(*mm_want, mm_got) { + mmListKnowledgeBasesByCatalogType.t.Errorf("RepositoryIMock.ListKnowledgeBasesByCatalogType got unexpected parameters, want: %#v, got: %#v%s\n", *mm_want, mm_got, minimock.Diff(*mm_want, mm_got)) + } + + mm_results := mmListKnowledgeBasesByCatalogType.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.results + if mm_results == nil { + mmListKnowledgeBasesByCatalogType.t.Fatal("No results are set for the RepositoryIMock.ListKnowledgeBasesByCatalogType") + } + return (*mm_results).ka1, (*mm_results).err + } + if mmListKnowledgeBasesByCatalogType.funcListKnowledgeBasesByCatalogType != nil { + return mmListKnowledgeBasesByCatalogType.funcListKnowledgeBasesByCatalogType(ctx, ownerUID, catalogType) + } + mmListKnowledgeBasesByCatalogType.t.Fatalf("Unexpected call to RepositoryIMock.ListKnowledgeBasesByCatalogType. %v %v %v", ctx, ownerUID, catalogType) + return +} + +// ListKnowledgeBasesByCatalogTypeAfterCounter returns a count of finished RepositoryIMock.ListKnowledgeBasesByCatalogType invocations +func (mmListKnowledgeBasesByCatalogType *RepositoryIMock) ListKnowledgeBasesByCatalogTypeAfterCounter() uint64 { + return mm_atomic.LoadUint64(&mmListKnowledgeBasesByCatalogType.afterListKnowledgeBasesByCatalogTypeCounter) +} + +// ListKnowledgeBasesByCatalogTypeBeforeCounter returns a count of RepositoryIMock.ListKnowledgeBasesByCatalogType invocations +func (mmListKnowledgeBasesByCatalogType *RepositoryIMock) ListKnowledgeBasesByCatalogTypeBeforeCounter() uint64 { + return mm_atomic.LoadUint64(&mmListKnowledgeBasesByCatalogType.beforeListKnowledgeBasesByCatalogTypeCounter) +} + +// Calls returns a list of arguments used in each call to RepositoryIMock.ListKnowledgeBasesByCatalogType. +// The list is in the same order as the calls were made (i.e. recent calls have a higher index) +func (mmListKnowledgeBasesByCatalogType *mRepositoryIMockListKnowledgeBasesByCatalogType) Calls() []*RepositoryIMockListKnowledgeBasesByCatalogTypeParams { + mmListKnowledgeBasesByCatalogType.mutex.RLock() + + argCopy := make([]*RepositoryIMockListKnowledgeBasesByCatalogTypeParams, len(mmListKnowledgeBasesByCatalogType.callArgs)) + copy(argCopy, mmListKnowledgeBasesByCatalogType.callArgs) + + mmListKnowledgeBasesByCatalogType.mutex.RUnlock() + + return argCopy +} + +// MinimockListKnowledgeBasesByCatalogTypeDone returns true if the count of the ListKnowledgeBasesByCatalogType invocations corresponds +// the number of defined expectations +func (m *RepositoryIMock) MinimockListKnowledgeBasesByCatalogTypeDone() bool { + for _, e := range m.ListKnowledgeBasesByCatalogTypeMock.expectations { + if mm_atomic.LoadUint64(&e.Counter) < 1 { + return false + } + } + + return m.ListKnowledgeBasesByCatalogTypeMock.invocationsDone() +} + +// MinimockListKnowledgeBasesByCatalogTypeInspect logs each unmet expectation +func (m *RepositoryIMock) MinimockListKnowledgeBasesByCatalogTypeInspect() { + for _, e := range m.ListKnowledgeBasesByCatalogTypeMock.expectations { + if mm_atomic.LoadUint64(&e.Counter) < 1 { + m.t.Errorf("Expected call to RepositoryIMock.ListKnowledgeBasesByCatalogType with params: %#v", *e.params) + } + } + + afterListKnowledgeBasesByCatalogTypeCounter := mm_atomic.LoadUint64(&m.afterListKnowledgeBasesByCatalogTypeCounter) + // if default expectation was set then invocations count should be greater than zero + if m.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation != nil && afterListKnowledgeBasesByCatalogTypeCounter < 1 { + if m.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.params == nil { + m.t.Error("Expected call to RepositoryIMock.ListKnowledgeBasesByCatalogType") + } else { + m.t.Errorf("Expected call to RepositoryIMock.ListKnowledgeBasesByCatalogType with params: %#v", *m.ListKnowledgeBasesByCatalogTypeMock.defaultExpectation.params) + } + } + // if func was set then invocations count should be greater than zero + if m.funcListKnowledgeBasesByCatalogType != nil && afterListKnowledgeBasesByCatalogTypeCounter < 1 { + m.t.Error("Expected call to RepositoryIMock.ListKnowledgeBasesByCatalogType") + } + + if !m.ListKnowledgeBasesByCatalogTypeMock.invocationsDone() && afterListKnowledgeBasesByCatalogTypeCounter > 0 { + m.t.Errorf("Expected %d calls to RepositoryIMock.ListKnowledgeBasesByCatalogType but found %d calls", + mm_atomic.LoadUint64(&m.ListKnowledgeBasesByCatalogTypeMock.expectedInvocations), afterListKnowledgeBasesByCatalogTypeCounter) + } +} + type mRepositoryIMockProcessKnowledgeBaseFiles struct { mock *RepositoryIMock defaultExpectation *RepositoryIMockProcessKnowledgeBaseFilesExpectation @@ -21484,6 +21855,8 @@ func (m *RepositoryIMock) MinimockFinish() { m.MinimockListKnowledgeBasesInspect() + m.MinimockListKnowledgeBasesByCatalogTypeInspect() + m.MinimockProcessKnowledgeBaseFilesInspect() m.MinimockTextChunkTableNameInspect() @@ -21584,6 +21957,7 @@ func (m *RepositoryIMock) minimockDone() bool { m.MinimockListEmbeddingsByKbFileUIDDone() && m.MinimockListKnowledgeBaseFilesDone() && m.MinimockListKnowledgeBasesDone() && + m.MinimockListKnowledgeBasesByCatalogTypeDone() && m.MinimockProcessKnowledgeBaseFilesDone() && m.MinimockTextChunkTableNameDone() && m.MinimockUpdateChunkDone() && diff --git a/pkg/repository/knowledgebase.go b/pkg/repository/knowledgebase.go index 3dacb9b..db710ed 100644 --- a/pkg/repository/knowledgebase.go +++ b/pkg/repository/knowledgebase.go @@ -10,16 +10,18 @@ import ( "github.com/gofrs/uuid" "github.com/instill-ai/artifact-backend/pkg/customerror" + artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha" "gorm.io/gorm" ) type KnowledgeBaseI interface { CreateKnowledgeBase(ctx context.Context, kb KnowledgeBase, externalService func(kbUID string) error) (*KnowledgeBase, error) ListKnowledgeBases(ctx context.Context, ownerUID string) ([]KnowledgeBase, error) + ListKnowledgeBasesByCatalogType(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) ([]KnowledgeBase, error) UpdateKnowledgeBase(ctx context.Context, ownerUID string, kb KnowledgeBase) (*KnowledgeBase, error) DeleteKnowledgeBase(ctx context.Context, ownerUID, kbID string) (*KnowledgeBase, error) GetKnowledgeBaseByOwnerAndKbID(ctx context.Context, ownerUID uuid.UUID, kbID string) (*KnowledgeBase, error) - GetKnowledgeBaseCountByOwner(ctx context.Context, ownerUID string) (int64, error) + GetKnowledgeBaseCountByOwner(ctx context.Context, ownerUID string, catalogType artifactpb.CatalogType) (int64, error) IncreaseKnowledgeBaseUsage(ctx context.Context, tx *gorm.DB, kbUID string, amount int) error } @@ -37,6 +39,8 @@ type KnowledgeBase struct { // creator CreatorUID uuid.UUID `gorm:"column:creator_uid;type:uuid;not null" json:"creator_uid"` Usage int64 `gorm:"column:usage;not null;default:0" json:"usage"` + // this type is defined in artifact/artifact/v1alpha/catalog.proto + CatalogType string `gorm:"column:catalog_type;size:255" json:"catalog_type"` } // table columns map @@ -51,6 +55,7 @@ type KnowledgeBaseColumns struct { UpdateTime string DeleteTime string Usage string + CatalogType string } var KnowledgeBaseColumn = KnowledgeBaseColumns{ @@ -64,6 +69,7 @@ var KnowledgeBaseColumn = KnowledgeBaseColumns{ UpdateTime: "update_time", DeleteTime: "delete_time", Usage: "usage", + CatalogType: "catalog_type", } // TagsArray is a custom type to handle PostgreSQL VARCHAR(255)[] arrays. @@ -156,6 +162,18 @@ func (r *Repository) ListKnowledgeBases(ctx context.Context, owner string) ([]Kn return knowledgeBases, nil } +// ListKnowledgeBasesByCatalogType fetches all KnowledgeBase records from the database, excluding soft-deleted ones. +func (r *Repository) ListKnowledgeBasesByCatalogType(ctx context.Context, owner string, catalogType artifactpb.CatalogType) ([]KnowledgeBase, error) { + var knowledgeBases []KnowledgeBase + // Exclude records where DeleteTime is not null and filter by owner + whereString := fmt.Sprintf("%v IS NULL AND %v = ? AND %v = ?", KnowledgeBaseColumn.DeleteTime, KnowledgeBaseColumn.Owner, KnowledgeBaseColumn.CatalogType) + if err := r.db.WithContext(ctx).Where(whereString, owner, catalogType.String()).Find(&knowledgeBases).Error; err != nil { + return nil, err + } + + return knowledgeBases, nil +} + // UpdateKnowledgeBase updates a KnowledgeBase record in the database except for CreateTime and DeleteTime. func (r *Repository) UpdateKnowledgeBase(ctx context.Context, owner string, kb KnowledgeBase) (*KnowledgeBase, error) { // Fetch the existing record to ensure it exists and to get the CreateTime and DeleteTime fields @@ -257,10 +275,10 @@ func (r *Repository) GetKnowledgeBaseByOwnerAndKbID(ctx context.Context, owner u } // get the count of knowledge bases by owner -func (r *Repository) GetKnowledgeBaseCountByOwner(ctx context.Context, owner string) (int64, error) { +func (r *Repository) GetKnowledgeBaseCountByOwner(ctx context.Context, owner string, catalogType artifactpb.CatalogType) (int64, error) { var count int64 - whereString := fmt.Sprintf("%v = ? AND %v is NULL", KnowledgeBaseColumn.Owner, KnowledgeBaseColumn.DeleteTime) - if err := r.db.WithContext(ctx).Model(&KnowledgeBase{}).Where(whereString, owner).Count(&count).Error; err != nil { + whereString := fmt.Sprintf("%v = ? AND %v is NULL AND %v = ?", KnowledgeBaseColumn.Owner, KnowledgeBaseColumn.DeleteTime, KnowledgeBaseColumn.CatalogType) + if err := r.db.WithContext(ctx).Model(&KnowledgeBase{}).Where(whereString, owner, catalogType.String()).Count(&count).Error; err != nil { return 0, err } return count, nil