feat: implement a store for stat_traffic_data (#6190)

## About the changes

Implements a new store for collected traffic data usage that connects to
the new table `stat_traffic_data` primary key'd on [day, trafficGroup,
status_code_series].

Day being a date
Traffic group being which endpoint is being counted for, ie /api/admin,
/api/frontend etc
Status code series grouping 2xx status responses and 304 into their
respective 200 / 300 series.

No service here, this is for pro/enterprise

src/lib/db/index.ts

@@ -41,6 +41,7 @@ import { DependentFeaturesStore } from '../features/dependent-features/dependent
 import LastSeenStore from '../features/metrics/last-seen/last-seen-store';
 import FeatureSearchStore from '../features/feature-search/feature-search-store';
 import { InactiveUsersStore } from '../users/inactive/inactive-users-store';
+import { TrafficDataUsageStore } from '../features/traffic-data-usage/traffic-data-usage-store';
 
 export const createStores = (
     config: IUnleashConfig,
@@ -143,6 +144,7 @@ export const createStores = (
         lastSeenStore: new LastSeenStore(db, eventBus, getLogger),
         featureSearchStore: new FeatureSearchStore(db, eventBus, getLogger),
         inactiveUsersStore: new InactiveUsersStore(db, eventBus, getLogger),
+        trafficDataUsageStore: new TrafficDataUsageStore(db, getLogger),
     };
 };
 

src/lib/features/traffic-data-usage/fake-traffic-data-usage-store.ts

@@ -0,0 +1,29 @@
+import {
+    IStatTrafficUsageKey,
+    IStatTrafficUsage,
+} from './traffic-data-usage-store-type';
+import { ITrafficDataUsageStore } from '../../types';
+
+export class FakeTrafficDataUsageStore implements ITrafficDataUsageStore {
+    get(key: IStatTrafficUsageKey): Promise<IStatTrafficUsage> {
+        throw new Error('Method not implemented.');
+    }
+    getAll(query?: Object | undefined): Promise<IStatTrafficUsage[]> {
+        throw new Error('Method not implemented.');
+    }
+    exists(key: IStatTrafficUsageKey): Promise<boolean> {
+        throw new Error('Method not implemented.');
+    }
+    delete(key: IStatTrafficUsageKey): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+    deleteAll(): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+    destroy(): void {
+        throw new Error('Method not implemented.');
+    }
+    upsert(trafficDataUsage: IStatTrafficUsage): Promise<void> {
+        throw new Error('Method not implemented.');
+    }
+}

src/lib/features/traffic-data-usage/traffic-data-usage-store-type.ts

@@ -0,0 +1,19 @@
+import { Store } from '../../types/stores/store';
+
+export type IStatTrafficUsage = {
+    day: Date;
+    trafficGroup: string;
+    statusCodeSeries: number;
+    count: number | string;
+};
+
+export interface IStatTrafficUsageKey {
+    day: Date;
+    trafficGroup: string;
+    statusCodeSeries: number;
+}
+
+export interface ITrafficDataUsageStore
+    extends Store<IStatTrafficUsage, IStatTrafficUsageKey> {
+    upsert(trafficDataUsage: IStatTrafficUsage): Promise<void>;
+}

src/lib/features/traffic-data-usage/traffic-data-usage-store.test.ts

@@ -0,0 +1,154 @@
+import dbInit, { ITestDb } from '../../../test/e2e/helpers/database-init';
+import getLogger from '../../../test/fixtures/no-logger';
+import { ITrafficDataUsageStore, IUnleashStores } from '../../types';
+
+let stores: IUnleashStores;
+let db: ITestDb;
+let trafficDataUsageStore: ITrafficDataUsageStore;
+
+beforeAll(async () => {
+    db = await dbInit('traffic_data_usage_serial', getLogger, {
+        experimental: {
+            flags: {},
+        },
+    });
+    stores = db.stores;
+    trafficDataUsageStore = stores.trafficDataUsageStore;
+});
+
+afterAll(async () => {
+    await db.destroy();
+});
+
+test('upsert stores new entries', async () => {
+    const data = {
+        day: new Date(),
+        trafficGroup: 'default',
+        statusCodeSeries: 200,
+        count: '1',
+    };
+    await trafficDataUsageStore.upsert(data);
+    const data2 = await trafficDataUsageStore.get({
+        day: data.day,
+        trafficGroup: data.trafficGroup,
+        statusCodeSeries: data.statusCodeSeries,
+    });
+    expect(data2).toBeDefined();
+    expect(data2.count).toBe('1');
+});
+
+test('upsert upserts', async () => {
+    const data = {
+        day: new Date(),
+        trafficGroup: 'default2',
+        statusCodeSeries: 200,
+        count: '1',
+    };
+    const dataSecondTime = {
+        day: new Date(),
+        trafficGroup: 'default2',
+        statusCodeSeries: 200,
+        count: '3',
+    };
+    await trafficDataUsageStore.upsert(data);
+    await trafficDataUsageStore.upsert(dataSecondTime);
+    const data2 = await trafficDataUsageStore.get({
+        day: data.day,
+        trafficGroup: data.trafficGroup,
+        statusCodeSeries: data.statusCodeSeries,
+    });
+    expect(data2).toBeDefined();
+    expect(data2.count).toBe('4');
+});
+
+test('getAll returns all', async () => {
+    trafficDataUsageStore.deleteAll();
+    const data1 = {
+        day: new Date(),
+        trafficGroup: 'default3',
+        statusCodeSeries: 200,
+        count: '1',
+    };
+    const data2 = {
+        day: new Date(),
+        trafficGroup: 'default4',
+        statusCodeSeries: 200,
+        count: '3',
+    };
+    await trafficDataUsageStore.upsert(data1);
+    await trafficDataUsageStore.upsert(data2);
+    const results = await trafficDataUsageStore.getAll();
+    expect(results).toBeDefined();
+    expect(results.length).toBe(2);
+});
+
+test('delete deletes the specified item', async () => {
+    trafficDataUsageStore.deleteAll();
+    const data1 = {
+        day: new Date(),
+        trafficGroup: 'default3',
+        statusCodeSeries: 200,
+        count: '1',
+    };
+    const data2 = {
+        day: new Date(),
+        trafficGroup: 'default4',
+        statusCodeSeries: 200,
+        count: '3',
+    };
+    await trafficDataUsageStore.upsert(data1);
+    await trafficDataUsageStore.upsert(data2);
+    await trafficDataUsageStore.delete({
+        day: data1.day,
+        trafficGroup: data1.trafficGroup,
+        statusCodeSeries: data1.statusCodeSeries,
+    });
+    const results = await trafficDataUsageStore.getAll();
+    expect(results).toBeDefined();
+    expect(results.length).toBe(1);
+    expect(results[0].trafficGroup).toBe('default4');
+});
+
+test('can query for specific items', async () => {
+    trafficDataUsageStore.deleteAll();
+    const data1 = {
+        day: new Date(),
+        trafficGroup: 'default3',
+        statusCodeSeries: 200,
+        count: '1',
+    };
+    const data2 = {
+        day: new Date(),
+        trafficGroup: 'default4',
+        statusCodeSeries: 200,
+        count: '3',
+    };
+    const data3 = {
+        day: new Date(),
+        trafficGroup: 'default5',
+        statusCodeSeries: 200,
+        count: '2',
+    };
+    await trafficDataUsageStore.upsert(data1);
+    await trafficDataUsageStore.upsert(data2);
+    await trafficDataUsageStore.upsert(data3);
+
+    const results_traffic_group = await trafficDataUsageStore.getAll({
+        traffic_group: data1.trafficGroup,
+    });
+    expect(results_traffic_group).toBeDefined();
+    expect(results_traffic_group.length).toBe(1);
+    expect(results_traffic_group[0].trafficGroup).toBe('default3');
+
+    const results_day = await trafficDataUsageStore.getAll({
+        day: data2.day,
+    });
+    expect(results_day).toBeDefined();
+    expect(results_day.length).toBe(3);
+
+    const results_status_code = await trafficDataUsageStore.getAll({
+        status_code_series: 200,
+    });
+    expect(results_status_code).toBeDefined();
+    expect(results_status_code.length).toBe(3);
+});

src/lib/features/traffic-data-usage/traffic-data-usage-store.ts

@@ -0,0 +1,90 @@
+import { Db } from '../../db/db';
+import { Logger, LogProvider } from '../../logger';
+import {
+    IStatTrafficUsage,
+    IStatTrafficUsageKey,
+    ITrafficDataUsageStore,
+} from './traffic-data-usage-store-type';
+
+const TABLE = 'stat_traffic_usage';
+const COLUMNS = ['day', 'traffic_group', 'status_code_series', 'count'];
+
+const toRow = (trafficDataUsage: IStatTrafficUsage) => {
+    return {
+        day: trafficDataUsage.day,
+        traffic_group: trafficDataUsage.trafficGroup,
+        status_code_series: trafficDataUsage.statusCodeSeries,
+        count: trafficDataUsage.count,
+    };
+};
+
+const mapRow = (row: any): IStatTrafficUsage => {
+    return {
+        day: row.day,
+        trafficGroup: row.traffic_group,
+        statusCodeSeries: row.status_code_series,
+        count: row.count,
+    };
+};
+
+export class TrafficDataUsageStore implements ITrafficDataUsageStore {
+    private db: Db;
+
+    private logger: Logger;
+
+    constructor(db: Db, getLogger: LogProvider) {
+        this.db = db;
+        this.logger = getLogger('traffic-data-usage-store.ts');
+    }
+    async get(key: IStatTrafficUsageKey): Promise<IStatTrafficUsage> {
+        const row = await this.db
+            .table(TABLE)
+            .select()
+            .where({
+                day: key.day,
+                traffic_group: key.trafficGroup,
+                status_code_series: key.statusCodeSeries,
+            })
+            .first();
+
+        return mapRow(row);
+    }
+    async getAll(query = {}): Promise<IStatTrafficUsage[]> {
+        const rows = await this.db.select(COLUMNS).where(query).from(TABLE);
+        return rows.map(mapRow);
+    }
+    async exists(key: IStatTrafficUsageKey): Promise<boolean> {
+        const result = await this.db.raw(
+            `SELECT EXISTS (SELECT 1 FROM ${TABLE} WHERE 
+                day = ? AND
+                traffic_group = ? AND
+                status_code_series ?) AS present`,
+            [key.day, key.trafficGroup, key.statusCodeSeries],
+        );
+        const { present } = result.rows[0];
+        return present;
+    }
+    async delete(key: IStatTrafficUsageKey): Promise<void> {
+        await this.db(TABLE)
+            .where({
+                day: key.day,
+                traffic_group: key.trafficGroup,
+                status_code_series: key.statusCodeSeries,
+            })
+            .del();
+    }
+    async deleteAll(): Promise<void> {
+        await this.db(TABLE).del();
+    }
+    destroy(): void {}
+
+    async upsert(trafficDataUsage: IStatTrafficUsage): Promise<void> {
+        const row = toRow(trafficDataUsage);
+        await this.db(TABLE)
+            .insert(row)
+            .onConflict(['day', 'traffic_group', 'status_code_series'])
+            .merge({
+                count: this.db.raw('stat_traffic_usage.count + EXCLUDED.count'),
+            });
+    }
+}

src/lib/types/stores.ts

@@ -38,6 +38,7 @@ import { IDependentFeaturesStore } from '../features/dependent-features/dependen
 import { ILastSeenStore } from '../features/metrics/last-seen/types/last-seen-store-type';
 import { IFeatureSearchStore } from '../features/feature-search/feature-search-store-type';
 import { IInactiveUsersStore } from '../users/inactive/types/inactive-users-store-type';
+import { ITrafficDataUsageStore } from '../features/traffic-data-usage/traffic-data-usage-store-type';
 
 export interface IUnleashStores {
     accessStore: IAccessStore;
@@ -80,6 +81,7 @@ export interface IUnleashStores {
     lastSeenStore: ILastSeenStore;
     featureSearchStore: IFeatureSearchStore;
     inactiveUsersStore: IInactiveUsersStore;
+    trafficDataUsageStore: ITrafficDataUsageStore;
 }
 
 export {
@@ -121,4 +123,5 @@ export {
     IDependentFeaturesStore,
     ILastSeenStore,
     IFeatureSearchStore,
+    ITrafficDataUsageStore,
 };

src/test/fixtures/store.ts

@@ -41,6 +41,7 @@ import { FakeDependentFeaturesStore } from '../../lib/features/dependent-feature
 import { FakeLastSeenStore } from '../../lib/features/metrics/last-seen/fake-last-seen-store';
 import FakeFeatureSearchStore from '../../lib/features/feature-search/fake-feature-search-store';
 import { FakeInactiveUsersStore } from '../../lib/users/inactive/fakes/fake-inactive-users-store';
+import { FakeTrafficDataUsageStore } from '../../lib/features/traffic-data-usage/fake-traffic-data-usage-store';
 
 const db = {
     select: () => ({
@@ -91,6 +92,7 @@ const createStores: () => IUnleashStores = () => {
         lastSeenStore: new FakeLastSeenStore(),
         featureSearchStore: new FakeFeatureSearchStore(),
         inactiveUsersStore: new FakeInactiveUsersStore(),
+        trafficDataUsageStore: new FakeTrafficDataUsageStore(),
     };
 };