@@ -1120,6 +1120,14 @@ func (q *FakeQuerier) getWorkspaceAgentScriptsByAgentIDsNoLock(ids []uuid.UUID)
11201120 return scripts , nil
11211121}
11221122
1123+ // getOwnerFromTags returns the lowercase owner from tags, matching SQL's COALESCE(tags ->> 'owner', ”)
1124+ func getOwnerFromTags (tags map [string ]string ) string {
1125+ if owner , ok := tags ["owner" ]; ok {
1126+ return strings .ToLower (owner )
1127+ }
1128+ return ""
1129+ }
1130+
11231131func (* FakeQuerier ) AcquireLock (_ context.Context , _ int64 ) error {
11241132 return xerrors .New ("AcquireLock must only be called within a transaction" )
11251133}
@@ -2773,6 +2781,63 @@ func (q *FakeQuerier) GetDeploymentWorkspaceStats(ctx context.Context) (database
27732781 return stat , nil
27742782}
27752783
2784+ func (q * FakeQuerier ) GetEligibleProvisionerDaemonsByProvisionerJobIDs (_ context.Context , provisionerJobIds []uuid.UUID ) ([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow , error ) {
2785+ q .mutex .RLock ()
2786+ defer q .mutex .RUnlock ()
2787+
2788+ results := make ([]database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow , 0 )
2789+ seen := make (map [string ]struct {}) // Track unique combinations
2790+
2791+ for _ , jobID := range provisionerJobIds {
2792+ var job database.ProvisionerJob
2793+ found := false
2794+ for _ , j := range q .provisionerJobs {
2795+ if j .ID == jobID {
2796+ job = j
2797+ found = true
2798+ break
2799+ }
2800+ }
2801+ if ! found {
2802+ continue
2803+ }
2804+
2805+ for _ , daemon := range q .provisionerDaemons {
2806+ if daemon .OrganizationID != job .OrganizationID {
2807+ continue
2808+ }
2809+
2810+ if ! tagsSubset (job .Tags , daemon .Tags ) {
2811+ continue
2812+ }
2813+
2814+ provisionerMatches := false
2815+ for _ , p := range daemon .Provisioners {
2816+ if p == job .Provisioner {
2817+ provisionerMatches = true
2818+ break
2819+ }
2820+ }
2821+ if ! provisionerMatches {
2822+ continue
2823+ }
2824+
2825+ key := jobID .String () + "-" + daemon .ID .String ()
2826+ if _ , exists := seen [key ]; exists {
2827+ continue
2828+ }
2829+ seen [key ] = struct {}{}
2830+
2831+ results = append (results , database.GetEligibleProvisionerDaemonsByProvisionerJobIDsRow {
2832+ JobID : jobID ,
2833+ ProvisionerDaemon : daemon ,
2834+ })
2835+ }
2836+ }
2837+
2838+ return results , nil
2839+ }
2840+
27762841func (q * FakeQuerier ) GetExternalAuthLink (_ context.Context , arg database.GetExternalAuthLinkParams ) (database.ExternalAuthLink , error ) {
27772842 if err := validateDatabaseType (arg ); err != nil {
27782843 return database.ExternalAuthLink {}, err
@@ -10344,25 +10409,26 @@ func (q *FakeQuerier) UpsertOAuthSigningKey(_ context.Context, value string) err
1034410409}
1034510410
1034610411func (q * FakeQuerier ) UpsertProvisionerDaemon (_ context.Context , arg database.UpsertProvisionerDaemonParams ) (database.ProvisionerDaemon , error ) {
10347- err := validateDatabaseType (arg )
10348- if err != nil {
10412+ if err := validateDatabaseType (arg ); err != nil {
1034910413 return database.ProvisionerDaemon {}, err
1035010414 }
1035110415
1035210416 q .mutex .Lock ()
1035310417 defer q .mutex .Unlock ()
10354- for _ , d := range q .provisionerDaemons {
10355- if d .Name == arg .Name {
10356- if d .Tags [provisionersdk .TagScope ] == provisionersdk .ScopeOrganization && arg .Tags [provisionersdk .TagOwner ] != "" {
10357- continue
10358- }
10359- if d .Tags [provisionersdk .TagScope ] == provisionersdk .ScopeUser && arg .Tags [provisionersdk .TagOwner ] != d .Tags [provisionersdk .TagOwner ] {
10360- continue
10361- }
10418+
10419+ // Look for existing daemon using the same composite key as SQL
10420+ for i , d := range q .provisionerDaemons {
10421+ if d .OrganizationID == arg .OrganizationID &&
10422+ d .Name == arg .Name &&
10423+ getOwnerFromTags (d .Tags ) == getOwnerFromTags (arg .Tags ) {
1036210424 d .Provisioners = arg .Provisioners
1036310425 d .Tags = maps .Clone (arg .Tags )
10364- d .Version = arg .Version
1036510426 d .LastSeenAt = arg .LastSeenAt
10427+ d .Version = arg .Version
10428+ d .APIVersion = arg .APIVersion
10429+ d .OrganizationID = arg .OrganizationID
10430+ d .KeyID = arg .KeyID
10431+ q .provisionerDaemons [i ] = d
1036610432 return d , nil
1036710433 }
1036810434 }
@@ -10372,7 +10438,6 @@ func (q *FakeQuerier) UpsertProvisionerDaemon(_ context.Context, arg database.Up
1037210438 Name : arg .Name ,
1037310439 Provisioners : arg .Provisioners ,
1037410440 Tags : maps .Clone (arg .Tags ),
10375- ReplicaID : uuid.NullUUID {},
1037610441 LastSeenAt : arg .LastSeenAt ,
1037710442 Version : arg .Version ,
1037810443 APIVersion : arg .APIVersion ,
0 commit comments