From 472ffd5b5c45127a6ea2c5d5e9d16b240a667bb6 Mon Sep 17 00:00:00 2001
From: Ali BARIN <ali.barin53@gmail.com>
Date: Tue, 7 Mar 2023 21:26:12 +0000
Subject: [PATCH] feat: increase consumed task count upon tasks

---
 packages/backend/src/models/execution-step.ts | 16 +++++++++++++++-
 packages/backend/src/models/execution.ts      |  1 +
 packages/backend/src/models/flow.ts           | 10 ++++++++++
 packages/backend/src/models/usage-data.ee.ts  |  4 ++++
 4 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/packages/backend/src/models/execution-step.ts b/packages/backend/src/models/execution-step.ts
index 2e4d2a8f..035c61c8 100644
--- a/packages/backend/src/models/execution-step.ts
+++ b/packages/backend/src/models/execution-step.ts
@@ -1,9 +1,10 @@
 import type { QueryContext } from 'objection';
+import { IJSONObject } from '@automatisch/types';
+import appConfig from '../config/app';
 import Base from './base';
 import Execution from './execution';
 import Step from './step';
 import Telemetry from '../helpers/telemetry';
-import { IJSONObject } from '@automatisch/types';
 
 class ExecutionStep extends Base {
   id!: string;
@@ -14,6 +15,7 @@ class ExecutionStep extends Base {
   errorDetails: IJSONObject;
   status: 'success' | 'failure';
   step: Step;
+  execution?: Execution;
 
   static tableName = 'execution_steps';
 
@@ -57,6 +59,18 @@ class ExecutionStep extends Base {
   async $afterInsert(queryContext: QueryContext) {
     await super.$afterInsert(queryContext);
     Telemetry.executionStepCreated(this);
+
+    if (appConfig.isCloud) {
+      const execution = await this.$relatedQuery('execution');
+
+      if (!execution.testRun && !this.isFailed) {
+        const flow = await execution.$relatedQuery('flow');
+        const user = await flow.$relatedQuery('user');
+        const usageData = await user.$relatedQuery('usageData');
+
+        await usageData.increaseConsumedTaskCountByOne();
+      }
+    }
   }
 }
 
diff --git a/packages/backend/src/models/execution.ts b/packages/backend/src/models/execution.ts
index c4847b17..e5bd82f2 100644
--- a/packages/backend/src/models/execution.ts
+++ b/packages/backend/src/models/execution.ts
@@ -10,6 +10,7 @@ class Execution extends Base {
   testRun: boolean;
   internalId: string;
   executionSteps: ExecutionStep[];
+  flow?: Flow;
 
   static tableName = 'executions';
 
diff --git a/packages/backend/src/models/flow.ts b/packages/backend/src/models/flow.ts
index 5577dd35..e0868b97 100644
--- a/packages/backend/src/models/flow.ts
+++ b/packages/backend/src/models/flow.ts
@@ -3,6 +3,7 @@ import type { ModelOptions, QueryContext } from 'objection';
 import ExtendedQueryBuilder from './query-builder';
 import Base from './base';
 import Step from './step';
+import User from './user';
 import Execution from './execution';
 import Telemetry from '../helpers/telemetry';
 
@@ -15,6 +16,7 @@ class Flow extends Base {
   published_at: string;
   remoteWebhookId: string;
   executions?: Execution[];
+  user?: User;
 
   static tableName = 'flows';
 
@@ -51,6 +53,14 @@ class Flow extends Base {
         to: 'executions.flow_id',
       },
     },
+    user: {
+      relation: Base.HasOneRelation,
+      modelClass: User,
+      join: {
+        from: 'flows.user_id',
+        to: 'users.id',
+      },
+    },
   });
 
   async lastInternalId() {
diff --git a/packages/backend/src/models/usage-data.ee.ts b/packages/backend/src/models/usage-data.ee.ts
index 50e6c84f..f5898b82 100644
--- a/packages/backend/src/models/usage-data.ee.ts
+++ b/packages/backend/src/models/usage-data.ee.ts
@@ -31,6 +31,10 @@ class UsageData extends Base {
       },
     },
   });
+
+  async increaseConsumedTaskCountByOne() {
+    return await this.$query().patch({ consumedTaskCount: this.consumedTaskCount + 1 });
+  }
 }
 
 export default UsageData;