refactor(api): other count queries

kouprlabs · Aug 26, 2024 · 1c7a2f4 · 1c7a2f4
1 parent c4740f4
commit 1c7a2f4
Show file tree

Hide file tree

Showing 8 changed files with 48 additions and 84 deletions.
diff --git a/api/repo/file_repo.go b/api/repo/file_repo.go
@@ -348,17 +348,12 @@ func (repo *fileRepo) DeleteChunk(ids []string) error {
 }
 
 func (repo *fileRepo) Count() (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw("SELECT count(*) as result FROM file").
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&fileEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *fileRepo) GetIDsByWorkspace(workspaceID string) ([]string, error) {
@@ -480,7 +475,7 @@ func (repo *fileRepo) GetItemCount(id string) (int64, error) {
 			id).
 		Scan(&res)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
 	return res.Result - 1, nil
 }

diff --git a/api/repo/group_repo.go b/api/repo/group_repo.go
@@ -181,7 +181,7 @@ func (repo *groupRepo) Count() (int64, error) {
 	var count int64
 	db := repo.db.Model(&groupEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
 	return count, nil
 }
@@ -280,19 +280,14 @@ func (repo *groupRepo) GetMembers(id string) ([]model.User, error) {
 }
 
 func (repo *groupRepo) GetOwnerCount(id string) (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw(`SELECT count(*) as result FROM userpermission
-             WHERE resource_id = ? and permission = ?`,
-			id, model.PermissionOwner).
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&userPermissionEntity{}).
+		Where("resource_id = ? and permission = ?").
+		Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *groupRepo) GrantUserPermission(id string, userID string, permission string) error {

diff --git a/api/repo/invitation_repo.go b/api/repo/invitation_repo.go
@@ -164,10 +164,11 @@ func (repo *invitationRepo) GetIncomingCount(email string) (int64, error) {
 	var count int64
 	db := repo.db.
 		Model(&invitationEntity{}).
-		Where("email = ? and status = 'pending'", email).
+		Where("email = ?", email).
+		Where("status = 'pending'").
 		Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
 	return count, nil
 }

diff --git a/api/repo/organization_repo.go b/api/repo/organization_repo.go
@@ -167,17 +167,12 @@ func (repo *organizationRepo) Find(id string) (model.Organization, error) {
 }
 
 func (repo *organizationRepo) Count() (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw("SELECT count(*) as result FROM organization").
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&organizationEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *organizationRepo) Save(org model.Organization) error {
@@ -256,19 +251,15 @@ func (repo *organizationRepo) GetGroups(id string) ([]model.Group, error) {
 }
 
 func (repo *organizationRepo) GetOwnerCount(id string) (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw(`SELECT count(*) as result FROM userpermission
-             WHERE resource_id = ? and permission = ?`,
-			id, model.PermissionOwner).
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&userPermissionEntity{}).
+		Where("resource_id = ?", id).
+		Where("permission = ?", model.PermissionOwner).
+		Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *organizationRepo) GrantUserPermission(id string, userID string, permission string) error {

diff --git a/api/repo/snapshot_repo.go b/api/repo/snapshot_repo.go
@@ -43,7 +43,7 @@ type SnapshotRepo interface {
 	DeleteMappingsForTree(fileID string) error
 	DeleteAllDangling() error
 	GetLatestVersionForFile(fileID string) (int64, error)
-	CountAssociations(id string) (int, error)
+	CountAssociations(id string) (int64, error)
 	Attach(sourceFileID string, targetFileID string) error
 	Detach(id string, fileID string) error
 }
@@ -687,22 +687,19 @@ func (repo *snapshotRepo) GetLatestVersionForFile(fileID string) (int64, error)
              WHERE map.file_id = ?`,
 			fileID).
 		Scan(&res); db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
 	return res.Result, nil
 }
 
-func (repo *snapshotRepo) CountAssociations(id string) (int, error) {
-	type Result struct {
-		Count int
-	}
-	var res Result
-	if db := repo.db.
-		Raw("SELECT COUNT(*) count FROM snapshot_file WHERE snapshot_id = ?", id).
-		Scan(&res); db.Error != nil {
-		return 0, db.Error
+func (repo *snapshotRepo) CountAssociations(id string) (int64, error) {
+	var count int64
+	if db := repo.db.Model(&SnapshotFileEntity{}).
+		Where("snapshot_id = ?", id).
+		Count(&count); db.Error != nil {
+		return -1, db.Error
 	}
-	return res.Count, nil
+	return count, nil
 }
 
 func (repo *snapshotRepo) Attach(sourceFileID string, targetFileID string) error {

diff --git a/api/repo/task_repo.go b/api/repo/task_repo.go
@@ -229,17 +229,12 @@ func (repo *taskRepo) Find(id string) (model.Task, error) {
 }
 
 func (repo *taskRepo) Count() (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw("SELECT count(*) as result FROM task").
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&taskEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *taskRepo) GetIDs(userID string) ([]string, error) {
@@ -267,7 +262,7 @@ func (repo *taskRepo) GetCountByEmail(userID string) (int64, error) {
 		Where("user_id = ?", userID).
 		Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
 	return count, nil
 }

diff --git a/api/repo/user_repo.go b/api/repo/user_repo.go
@@ -137,15 +137,10 @@ func (repo *userRepo) FindAll() ([]model.User, error) {
 }
 
 func (repo *userRepo) Count() (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw(`SELECT count(*) as result FROM "user"`).
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&userEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
diff --git a/api/repo/workspace_repo.go b/api/repo/workspace_repo.go
@@ -202,17 +202,12 @@ func (repo *workspaceRepo) Find(id string) (model.Workspace, error) {
 }
 
 func (repo *workspaceRepo) Count() (int64, error) {
-	type Result struct {
-		Result int64
-	}
-	var res Result
-	db := repo.db.
-		Raw("SELECT count(*) as result FROM workspace").
-		Scan(&res)
+	var count int64
+	db := repo.db.Model(&workspaceEntity{}).Count(&count)
 	if db.Error != nil {
-		return 0, db.Error
+		return -1, db.Error
 	}
-	return res.Result, nil
+	return count, nil
 }
 
 func (repo *workspaceRepo) UpdateName(id string, name string) (model.Workspace, error) {