refactor(desktop-task): extract recovery select query and scan task status

Pull the recovered-task SELECT into a dedicated helper so the column list and
lease-expiry filter are testable, add a status column to the scan so logs
report the real task state, and log scan/iterate failures with the client and
recovery mode for observability.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-26 10:19:00 +08:00
parent 513e71f9f1
commit 6b710c8919
2 changed files with 91 additions and 20 deletions
@@ -1363,27 +1363,8 @@ func (s *DesktopTaskService) recoverClientTasks(
defer tx.Rollback(ctx)
repo := repository.NewDesktopTaskRepository(tx)
query := `
SELECT
desktop_id,
workspace_id,
kind,
active_attempt_id
FROM desktop_tasks
WHERE target_client_id = $1
AND workspace_id = $2
AND status = 'in_progress'
`
query := desktopTaskRecoverySelectQuery(mode)
queryArgs := []any{client.ID, client.WorkspaceID}
if mode == desktopTaskRecoveryModeLeaseExpiry {
query += `
AND lease_expires_at IS NOT NULL
AND lease_expires_at < NOW()
`
}
query += `
FOR UPDATE
`
rows, err := tx.Query(ctx, query, queryArgs...)
if err != nil {
@@ -1402,11 +1383,13 @@ func (s *DesktopTaskService) recoverClientTasks(
&item.Status,
&item.ActiveAttempt,
); scanErr != nil {
s.logWarn("desktop task recovery scan failed", scanErr, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode)))
return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to parse recovered desktop tasks")
}
recovered = append(recovered, item)
}
if err := rows.Err(); err != nil {
s.logWarn("desktop task recovery rows failed", err, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode)))
return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to iterate recovered desktop tasks")
}
@@ -1509,6 +1492,31 @@ func (s *DesktopTaskService) recoverClientTasks(
return nil
}
func desktopTaskRecoverySelectQuery(mode desktopTaskRecoveryMode) string {
query := `
SELECT
desktop_id,
workspace_id,
kind,
status,
active_attempt_id
FROM desktop_tasks
WHERE target_client_id = $1
AND workspace_id = $2
AND status = 'in_progress'
`
if mode == desktopTaskRecoveryModeLeaseExpiry {
query += `
AND lease_expires_at IS NOT NULL
AND lease_expires_at < NOW()
`
}
query += `
FOR UPDATE
`
return query
}
func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]byte, string, string, error) {
reason := strings.TrimSpace(string(mode))
message := "desktop task was recovered after the client lost the active lease"
@@ -0,0 +1,63 @@
package app
import (
"regexp"
"strings"
"testing"
)
func TestDesktopTaskRecoverySelectQueryMatchesRecoveredLeaseScanOrder(t *testing.T) {
t.Parallel()
query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeStartup)
gotColumns := recoverDesktopTaskSelectColumns(query)
wantColumns := []string{
"desktop_id",
"workspace_id",
"kind",
"status",
"active_attempt_id",
}
if len(gotColumns) != len(wantColumns) {
t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns)
}
for index := range wantColumns {
if gotColumns[index] != wantColumns[index] {
t.Fatalf("recover select columns = %v, want %v", gotColumns, wantColumns)
}
}
}
func TestDesktopTaskRecoverySelectQueryLeaseExpiryFiltersExpiredLeases(t *testing.T) {
t.Parallel()
query := desktopTaskRecoverySelectQuery(desktopTaskRecoveryModeLeaseExpiry)
for _, fragment := range []string{
"AND lease_expires_at IS NOT NULL",
"AND lease_expires_at < NOW()",
"FOR UPDATE",
} {
if !strings.Contains(query, fragment) {
t.Fatalf("recover lease-expiry query missing %q: %s", fragment, query)
}
}
}
func recoverDesktopTaskSelectColumns(query string) []string {
re := regexp.MustCompile(`(?is)select\s+(.*?)\s+from\s+desktop_tasks`)
match := re.FindStringSubmatch(query)
if len(match) != 2 {
return nil
}
rawColumns := strings.Split(match[1], ",")
columns := make([]string, 0, len(rawColumns))
for _, rawColumn := range rawColumns {
column := strings.TrimSpace(rawColumn)
column = strings.Fields(column)[0]
columns = append(columns, column)
}
return columns
}