@@ -51,6 +51,75 @@ func (i *WorkspaceInstance) TableName() string {
5151 return "d_b_workspace_instance"
5252}
5353
54+ // FindStoppedWorkspaceInstancesInRange finds WorkspaceInstanceForUsage that have been stopped between from (inclusive) and to (exclusive).
55+ func FindStoppedWorkspaceInstancesInRange (ctx context.Context , conn * gorm.DB , from , to time.Time ) ([]WorkspaceInstanceForUsage , error ) {
56+ var instances []WorkspaceInstanceForUsage
57+ var instancesInBatch []WorkspaceInstanceForUsage
58+
59+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
60+ Where ("wsi.stoppingTime >= ?" , TimeToISO8601 (from )).
61+ Where ("wsi.stoppingTime < ?" , TimeToISO8601 (to )).
62+ Where ("wsi.stoppingTime != ?" , "" ).
63+ Where ("wsi.usageAttributionId != ?" , "" ).
64+ FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
65+ instances = append (instances , instancesInBatch ... )
66+ return nil
67+ })
68+ if tx .Error != nil {
69+ return nil , fmt .Errorf ("failed to find workspace instances: %w" , tx .Error )
70+ }
71+
72+ return instances , nil
73+ }
74+
75+ // FindRunningWorkspaceInstances finds WorkspaceInstanceForUsage that are running at the point in time the querty is executed.
76+ func FindRunningWorkspaceInstances (ctx context.Context , conn * gorm.DB ) ([]WorkspaceInstanceForUsage , error ) {
77+ var instances []WorkspaceInstanceForUsage
78+ var instancesInBatch []WorkspaceInstanceForUsage
79+
80+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
81+ Where ("wsi.stoppingTime = ?" , "" ).
82+ Where ("wsi.usageAttributionId != ?" , "" ).
83+ FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
84+ instances = append (instances , instancesInBatch ... )
85+ return nil
86+ })
87+ if tx .Error != nil {
88+ return nil , fmt .Errorf ("failed to find running workspace instances: %w" , tx .Error )
89+ }
90+
91+ return instances , nil
92+ }
93+
94+ // FindWorkspaceInstancesByIds finds WorkspaceInstanceForUsage by Id.
95+ func FindWorkspaceInstancesByIds (ctx context.Context , conn * gorm.DB , workspaceInstanceIds []uuid.UUID ) ([]WorkspaceInstanceForUsage , error ) {
96+ var instances []WorkspaceInstanceForUsage
97+ var instancesInBatch []WorkspaceInstanceForUsage
98+ var idChunks [][]uuid.UUID
99+ chunkSize , totalSize := 1000 , len (workspaceInstanceIds )
100+ // explicit batching to reduce the lengths of the 'in'-part in the SELECT statement below
101+ for i := 0 ; i < totalSize ; i += chunkSize {
102+ end := i + chunkSize
103+ if end > totalSize {
104+ end = totalSize
105+ }
106+ idChunks = append (idChunks , workspaceInstanceIds [i :end ])
107+ }
108+
109+ for _ , idChunk := range idChunks {
110+ err := queryWorkspaceInstanceForUsage (ctx , conn ).
111+ Where ("wsi.id in ?" , idChunk ).
112+ Where ("wsi.usageAttributionId != ?" , "" ).
113+ Find (& instancesInBatch ).Error
114+ if err != nil {
115+ return nil , fmt .Errorf ("failed to find workspace instances by id: %w" , err )
116+ }
117+ instances = append (instances , instancesInBatch ... )
118+ }
119+
120+ return instances , nil
121+ }
122+
54123// ListWorkspaceInstancesInRange lists WorkspaceInstances between from (inclusive) and to (exclusive).
55124// This results in all instances which have existed in the specified period, regardless of their current status, this includes:
56125// - terminated
@@ -61,26 +130,11 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to
61130 var instances []WorkspaceInstanceForUsage
62131 var instancesInBatch []WorkspaceInstanceForUsage
63132
64- tx := conn .WithContext (ctx ).
65- Table (fmt .Sprintf ("%s as wsi" , (& WorkspaceInstance {}).TableName ())).
66- Select ("wsi.id as id, " +
67- "ws.projectId as projectId, " +
68- "ws.type as workspaceType, " +
69- "wsi.workspaceClass as workspaceClass, " +
70- "wsi.usageAttributionId as usageAttributionId, " +
71- "wsi.creationTime as creationTime, " +
72- "wsi.startedTime as startedTime, " +
73- "wsi.stoppingTime as stoppingTime, " +
74- "wsi.stoppedTime as stoppedTime, " +
75- "ws.ownerId as ownerId, " +
76- "ws.id as workspaceId" ,
77- ).
78- Joins (fmt .Sprintf ("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id" , (& Workspace {}).TableName ())).
133+ tx := queryWorkspaceInstanceForUsage (ctx , conn ).
79134 Where (
80135 conn .Where ("wsi.stoppingTime >= ?" , TimeToISO8601 (from )).Or ("wsi.stoppingTime = ?" , "" ),
81136 ).
82- Where ("wsi.creationTime < ?" , TimeToISO8601 (to )).
83- Where ("wsi.startedTime != ?" , "" ).
137+ Where ("wsi.startedTime < ?" , TimeToISO8601 (to )).
84138 Where ("wsi.usageAttributionId != ?" , "" ).
85139 FindInBatches (& instancesInBatch , 1000 , func (_ * gorm.DB , _ int ) error {
86140 instances = append (instances , instancesInBatch ... )
@@ -93,6 +147,24 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to
93147 return instances , nil
94148}
95149
150+ func queryWorkspaceInstanceForUsage (ctx context.Context , conn * gorm.DB ) * gorm.DB {
151+ return conn .WithContext (ctx ).
152+ Table (fmt .Sprintf ("%s as wsi" , (& WorkspaceInstance {}).TableName ())).
153+ Select ("wsi.id as id, " +
154+ "ws.projectId as projectId, " +
155+ "ws.type as workspaceType, " +
156+ "wsi.workspaceClass as workspaceClass, " +
157+ "wsi.usageAttributionId as usageAttributionId, " +
158+ "wsi.creationTime as creationTime, " +
159+ "wsi.startedTime as startedTime, " +
160+ "wsi.stoppingTime as stoppingTime, " +
161+ "wsi.stoppedTime as stoppedTime, " +
162+ "ws.ownerId as ownerId, " +
163+ "ws.id as workspaceId" ,
164+ ).
165+ Joins (fmt .Sprintf ("LEFT JOIN %s AS ws ON wsi.workspaceId = ws.id" , (& Workspace {}).TableName ()))
166+ }
167+
96168const (
97169 AttributionEntity_User = "user"
98170 AttributionEntity_Team = "team"
0 commit comments