1+ import { Collection , Db , ObjectId } from 'mongodb' ;
2+ import { ProjectDBScheme , WorkspaceDBScheme } from '@hawk.so/types' ;
3+ import { WorkspaceWithTariffPlan } from '../types' ;
4+ import HawkCatcher from '@hawk.so/nodejs' ;
5+ import { CriticalError } from '../../../lib/workerErrors' ;
6+
7+ /**
8+ * Class that implements methods used for interaction between limiter and db
9+ */
10+ export class DbHelper {
11+ /**
12+ * Connection to events DB
13+ */
14+ private eventsDbConnection : Db ;
15+
16+ /**
17+ * Collection with projects
18+ */
19+ private projectsCollection : Collection < ProjectDBScheme > ;
20+
21+ /**
22+ * Collection with workspaces
23+ */
24+ private workspacesCollection : Collection < WorkspaceDBScheme > ;
25+
26+ /**
27+ * @param projects - projects collection
28+ * @param workspaces - workspaces collection
29+ * @param eventsDbConnection - connection to events DB
30+ */
31+ constructor ( projects : Collection < ProjectDBScheme > , workspaces : Collection < WorkspaceDBScheme > , eventsDbConnection : Db ) {
32+ this . eventsDbConnection = eventsDbConnection ;
33+ this . projectsCollection = projects ;
34+ this . workspacesCollection = workspaces ;
35+ }
36+
37+ /**
38+ * Method that returns all workspaces with their tariff plans
39+ */
40+ public async getWorkspacesWithTariffPlans ( ) :Promise < WorkspaceWithTariffPlan [ ] > ;
41+ /**
42+ * Method that returns workspace with its tariff plan by its id
43+ *
44+ * @param id - id of the workspace to fetch
45+ */
46+ public async getWorkspacesWithTariffPlans ( id : string ) :Promise < WorkspaceWithTariffPlan > ;
47+ /**
48+ * Returns workspace with its tariff plan by its id
49+ *
50+ * @param id - workspace id
51+ */
52+ public async getWorkspacesWithTariffPlans ( id ?: string ) :Promise < WorkspaceWithTariffPlan [ ] | WorkspaceWithTariffPlan > {
53+ /* eslint-disable-next-line */
54+ const queue : any [ ] = [
55+ {
56+ $lookup : {
57+ from : 'plans' ,
58+ localField : 'tariffPlanId' ,
59+ foreignField : '_id' ,
60+ as : 'tariffPlan' ,
61+ } ,
62+ } ,
63+ {
64+ $unwind : {
65+ path : '$tariffPlan' ,
66+ } ,
67+ } ,
68+ ] ;
69+
70+ if ( id !== undefined ) {
71+ queue . unshift ( {
72+ $match : {
73+ _id : new ObjectId ( id ) ,
74+ } ,
75+ } ) ;
76+ }
77+
78+ const workspacesArray = await this . workspacesCollection . aggregate < WorkspaceWithTariffPlan > ( queue ) . toArray ( ) ;
79+
80+ return ( id !== undefined ) ? workspacesArray [ 0 ] : workspacesArray ;
81+ }
82+
83+ /**
84+ * Updates workspaces data in Database
85+ *
86+ * @param workspacesToUpdate - array of workspaces to be updated
87+ */
88+ public async updateWorkspacesEventsCountAndIsBlocked ( workspacesToUpdate : WorkspaceWithTariffPlan [ ] ) : Promise < void > {
89+ if ( workspacesToUpdate . length === 0 ) {
90+ return ;
91+ }
92+
93+ const operations = workspacesToUpdate . map ( workspace => {
94+ return {
95+ updateOne : {
96+ filter : {
97+ _id : workspace . _id ,
98+ } ,
99+ update : {
100+ $set : {
101+ billingPeriodEventsCount : workspace . billingPeriodEventsCount ,
102+ isBlocked : workspace . isBlocked ,
103+ } ,
104+ } ,
105+ } ,
106+ } ;
107+ } ) ;
108+
109+ await this . workspacesCollection . bulkWrite ( operations ) ;
110+ }
111+
112+ /**
113+ * Method to change workspace isBlocked state
114+ *
115+ * @param workspaceId - id of the workspace to be changed
116+ * @param isBlocked - new isBlocked state of the workspace
117+ */
118+ public async changeWorkspaceBlockedState ( workspaceId : string , isBlocked : boolean ) : Promise < void > {
119+ await this . workspacesCollection . updateOne (
120+ { _id : new ObjectId ( workspaceId ) } ,
121+ {
122+ $set : {
123+ isBlocked,
124+ } ,
125+ }
126+ ) ;
127+ }
128+
129+ /**
130+ * Returns total event counts for last billing period
131+ *
132+ * @param project - project to check
133+ * @param since - timestamp of the time from which we count the events
134+ */
135+ public async getEventsCountByProject (
136+ project : ProjectDBScheme ,
137+ since : number
138+ ) : Promise < number > {
139+ try {
140+ const repetitionsCollection = this . eventsDbConnection . collection ( 'repetitions:' + project . _id . toString ( ) ) ;
141+ const eventsCollection = this . eventsDbConnection . collection ( 'events:' + project . _id . toString ( ) ) ;
142+
143+ const query = {
144+ $or : [ {
145+ timestamp : {
146+ $gt : since ,
147+ } ,
148+ } ,
149+ {
150+ 'payload.timestamp' : {
151+ $gt : since ,
152+ } ,
153+ } ] ,
154+ } ;
155+
156+ const repetitionsCount = await repetitionsCollection . countDocuments ( query ) ;
157+ const originalEventCount = await eventsCollection . countDocuments ( query ) ;
158+
159+ return repetitionsCount + originalEventCount ;
160+ } catch ( e ) {
161+ HawkCatcher . send ( e ) ;
162+ throw new CriticalError ( e ) ;
163+ }
164+ }
165+
166+ /**
167+ * Calculates total events count for all provided projects since the specific date
168+ *
169+ * @param projects - projects to calculate for
170+ * @param since - timestamp of the time from which we count the events
171+ */
172+ public async getEventsCountByProjects ( projects : ProjectDBScheme [ ] , since : number ) : Promise < number > {
173+ const sum = ( array : number [ ] ) : number => array . reduce ( ( acc , val ) => acc + val , 0 ) ;
174+
175+ return Promise . all ( projects . map (
176+ project => this . getEventsCountByProject ( project , since )
177+ ) )
178+ . then ( sum ) ;
179+ }
180+
181+ /**
182+ * Returns all projects from Database or projects of the specified workspace
183+ *
184+ * @param [workspaceId] - workspace ids to fetch projects that belongs that workspace
185+ */
186+ public getProjects ( workspaceId ?: string ) : Promise < ProjectDBScheme [ ] > {
187+ const query = workspaceId
188+ ? {
189+ $or : [
190+ { workspaceId : workspaceId } ,
191+ { workspaceId : new ObjectId ( workspaceId ) } ,
192+ ] ,
193+ }
194+ : { } ;
195+
196+ return this . projectsCollection . find ( query ) . toArray ( ) ;
197+ }
198+ }
0 commit comments