#130 make it compatible with sqlite3 and also add an option to disabl…

…e prefetching Signed-off-by: Long Zhang <[email protected]>
electrolux-oss · Jan 21, 2025 · 385ff98 · 385ff98
1 parent e67b8f1
commit 385ff98
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 15 deletions.
diff --git a/plugins/infrawallet-backend/config.d.ts b/plugins/infrawallet-backend/config.d.ts
@@ -1,6 +1,7 @@
 export interface Config {
   backend: {
     infraWallet: {
+      prefetchCostData?: boolean; // true to enable prefetching cost data and saving them into plugin db
       integrations: {
         azure?: {
           name: string;

diff --git a/plugins/infrawallet-backend/src/cost-clients/InfraWalletClient.ts b/plugins/infrawallet-backend/src/cost-clients/InfraWalletClient.ts
@@ -255,6 +255,7 @@ export abstract class InfraWalletClient {
   }
 
   async getCostReports(query: CostQuery): Promise<ClientResponse> {
+    const prefetchCostData = this.config.getOptionalBoolean('backend.infraWallet.prefetchCostData') ?? true;
     const integrationConfigs = this.config.getOptionalConfigArray(
       `backend.infraWallet.integrations.${this.provider.toLowerCase()}`,
     );
@@ -265,8 +266,8 @@ export abstract class InfraWalletClient {
     const results: Report[] = [];
     const errors: CloudProviderError[] = [];
 
-    // for a query without any tags or groups, we get the results from the plugin database
-    if (query.tags === '()' && query.groups === '') {
+    // if prefetchCostData enabled, for a query without any tags or groups, we get the results from the plugin database
+    if (query.tags === '()' && query.groups === '' && prefetchCostData) {
       const reportsFromDatabase = await this.getCostReportsFromDatabase(query);
       reportsFromDatabase.forEach(report => {
         results.push(report);
@@ -393,6 +394,8 @@ export abstract class InfraWalletClient {
         costItems,
         (accumulator: { [key: string]: Report }, row: CostItem) => {
           const key = row.key;
+          const otherColumns =
+            typeof row.other_columns === 'string' ? JSON.parse(row.other_columns) : row.other_columns;
 
           if (!accumulator[key]) {
             accumulator[key] = {
@@ -403,7 +406,7 @@ export abstract class InfraWalletClient {
               provider: row.provider,
               providerType: PROVIDER_TYPE.INTEGRATION,
               reports: {},
-              ...row.other_columns,
+              ...otherColumns,
             };
           }
           accumulator[key].reports[usageDateToPeriodString(row.usage_date)] = parseFloat(row.cost as string);

diff --git a/plugins/infrawallet-backend/src/models/CostItem.ts b/plugins/infrawallet-backend/src/models/CostItem.ts
@@ -11,7 +11,7 @@ export type CostItem = {
   category: string;
   provider: string;
   usage_date: number; // format YYYYMMDD
-  other_columns: Record<string, string>; // example: {"cluster":"value_a", "project":"value_b"}
+  other_columns: Record<string, string> | string; // example: {"cluster":"value_a", "project":"value_b"}
   // If Postgres is used, the column type is decimal but Knex gets the values as strings
   // see https://stackoverflow.com/questions/45569216/knex-postgres-returns-strings-for-numeric-decimal-values
   cost: number | string;
@@ -121,16 +121,18 @@ export async function bulkInsertCostItems(
           category: report.category,
           provider: report.provider,
           usage_date: usageDate,
-          other_columns: otherColumns,
+          other_columns: knex.client.dialect === 'sqlite3' ? JSON.stringify(otherColumns) : otherColumns,
           cost: cost,
         });
       }
     }
   });
 
   // bulk insert the records
+  // for sqlite3, we need a smaller chunk size
+  const chunkSize = knex.client.dialect === 'sqlite3' ? 500 : 1000;
   await knex
-    .batchInsert(`cost_items_${granularity}`, rows)
+    .batchInsert(`cost_items_${granularity}`, rows, chunkSize)
     .then(() => {
       console.log(`${reports.length} ${granularity} records have been inserted`);
     })

diff --git a/plugins/infrawallet-backend/src/service/router.ts b/plugins/infrawallet-backend/src/service/router.ts
@@ -45,15 +45,21 @@ export async function createRouter(options: RouterOptions): Promise<express.Rout
   // do database migrations here to support the legacy backend system
   await setUpDatabase(database);
 
-  // put scheduler here for now to support legacy backends
-  await scheduler.scheduleTask({
-    frequency: { cron: '0 */8 * * *' }, // every 8 hours
-    timeout: { hours: 1 },
-    id: 'infrawallet-fetch-and-save-costs',
-    fn: async () => {
-      await fetchAndSaveCosts(options);
-    },
-  });
+  const prefetchCostData = config.getOptionalBoolean('backend.infraWallet.prefetchCostData') ?? true;
+
+  if (prefetchCostData) {
+    // put scheduler here for now to support legacy backends
+    await scheduler.scheduleTask({
+      frequency: { cron: '0 */8 * * *' }, // every 8 hours
+      timeout: { hours: 1 },
+      id: 'infrawallet-fetch-and-save-costs',
+      fn: async () => {
+        await fetchAndSaveCosts(options);
+      },
+    });
+    // trigger this task when the plugin starts up
+    scheduler.triggerTask('infrawallet-fetch-and-save-costs');
+  }
 
   // init CategoryMappingService
   CategoryMappingService.initInstance(cache, logger);