A simple file and memory based database (#9649)

* File and memory database

* Add read and write hook tests

extensions/azurecore/src/account-provider/utils/fileDatabase.ts

@@ -0,0 +1,189 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { promises as fs, constants as fsConstants } from 'fs';
+
+export type ReadWriteHook = (contents: string) => Promise<string>;
+const noOpHook: ReadWriteHook = async (contents): Promise<string> => {
+	return contents;
+};
+
+export class AlreadyInitializedError extends Error {
+
+}
+
+export class FileDatabase {
+	private db: { [key: string]: string };
+	private isDirty = false;
+	private isSaving = false;
+	private isInitialized = false;
+	private saveInterval: NodeJS.Timer;
+
+	constructor(
+		private readonly dbPath: string,
+		private readHook: ReadWriteHook = noOpHook,
+		private writeHook: ReadWriteHook = noOpHook
+	) {
+
+	}
+
+	/**
+	 * Sets a new read hook. Throws AlreadyInitializedError if the database has already started.
+	 * @param hook
+	 */
+	public setReadHook(hook: ReadWriteHook): void {
+		if (this.isInitialized) {
+			throw new AlreadyInitializedError();
+		}
+		this.readHook = hook;
+	}
+
+	/**
+	 * Sets a new write hook.
+	 * @param hook
+	 */
+	public setWriteHook(hook: ReadWriteHook): void {
+		this.writeHook = hook;
+	}
+
+	public async set(key: string, value: string): Promise<void> {
+		await this.waitForFileSave();
+		this.db[key] = value;
+		this.isDirty = true;
+	}
+
+	public get(key: string): string {
+		return this.db[key];
+	}
+
+	public async delete(key: string): Promise<void> {
+		await this.waitForFileSave();
+		delete this.db[key];
+		this.isDirty = true;
+	}
+
+	public async clear(): Promise<void> {
+		await this.waitForFileSave();
+		this.db = {};
+		this.isDirty = true;
+	}
+
+	public getPrefix(keyPrefix: string): { key: string, value: string }[] {
+		return Object.entries(this.db).filter(([key]) => {
+			return key.startsWith(keyPrefix);
+		}).map(([key, value]) => {
+			return { key, value };
+		});
+	}
+
+	public async deletePrefix(keyPrefix: string): Promise<void> {
+		await this.waitForFileSave();
+		Object.keys(this.db).forEach(s => {
+			if (s.startsWith(keyPrefix)) {
+				delete this.db[s];
+			}
+		});
+		this.isDirty = true;
+	}
+
+
+	public async initialize(): Promise<void> {
+		this.isInitialized = true;
+		this.setupSaveTask();
+		let fileContents: string;
+		try {
+			await fs.access(this.dbPath, fsConstants.R_OK | fsConstants.R_OK);
+			fileContents = await fs.readFile(this.dbPath, { encoding: 'utf8' });
+			fileContents = await this.readHook(fileContents);
+		} catch (ex) {
+			console.log(`file db does not exist ${ex}`);
+			await this.createFile();
+			this.db = {};
+			this.isDirty = true;
+			return;
+		}
+
+		try {
+			this.db = JSON.parse(fileContents);
+		} catch (ex) {
+			console.log(`DB was corrupted, resetting it ${ex}`);
+			await this.createFile();
+			this.db = {};
+		}
+	}
+
+	private setupSaveTask(): NodeJS.Timer {
+		return setInterval(() => this.save(), 20 * 1000);
+	}
+
+	public async shutdown(): Promise<void> {
+		await this.waitForFileSave();
+		clearInterval((this.saveInterval));
+		await this.save();
+	}
+
+	/**
+	 * This doesn't need to be called as a timer will automatically call it.
+	 */
+	public async save(): Promise<void> {
+		try {
+			await this.waitForFileSave();
+			if (this.isDirty === false) {
+				return;
+			}
+
+			this.isSaving = true;
+			let contents = JSON.stringify(this.db);
+			contents = await this.writeHook(contents);
+
+			await fs.writeFile(this.dbPath, contents, { encoding: 'utf8' });
+
+			this.isDirty = false;
+		} catch (ex) {
+			console.log(`File saving is erroring! ${ex}`);
+		} finally {
+			this.isSaving = false;
+		}
+	}
+
+	private async waitForFileSave(): Promise<void> {
+		const cleanupCrew: NodeJS.Timer[] = [];
+
+		const sleepToFail = (time: number): Promise<void> => {
+			return new Promise((_, reject) => {
+				const timeout = setTimeout(reject, time);
+				cleanupCrew.push(timeout);
+			});
+		};
+
+		const poll = (func: () => boolean): Promise<void> => {
+			return new Promise(resolve => {
+				const interval = setInterval(() => {
+					if (func() === true) {
+						resolve();
+					}
+				}, 100);
+				cleanupCrew.push(interval);
+			});
+		};
+
+		if (this.isSaving) {
+			const timeout = sleepToFail(5 * 1000);
+			const check = poll(() => !this.isSaving);
+
+			try {
+				return await Promise.race([timeout, check]);
+			} catch (ex) {
+				throw new Error('Save timed out');
+			} finally {
+				cleanupCrew.forEach(clearInterval);
+			}
+		}
+	}
+
+	private async createFile(): Promise<void> {
+		return fs.writeFile(this.dbPath, '', { encoding: 'utf8' });
+	}
+}

extensions/azurecore/src/account-provider/utils/memoryDatabase.ts

@@ -0,0 +1,23 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+export class MemoryDatabase<T> {
+	db: { [key: string]: T } = {};
+
+	constructor() {
+	}
+
+	public set(key: string, value: T): void {
+		this.db[key] = value;
+	}
+
+	public delete(key: string): void {
+		delete this.db[key];
+	}
+
+	public get(key: string): T {
+		return this.db[key];
+	}
+}

extensions/azurecore/src/test/account-provider/utils/fileDatabase.test.ts

@@ -0,0 +1,99 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as should from 'should';
+import * as os from 'os';
+import * as path from 'path';
+import * as crypto from 'crypto';
+import { promises as fs } from 'fs';
+
+import 'mocha';
+
+import { FileDatabase } from '../../../account-provider/utils/fileDatabase';
+
+let fileDatabase: FileDatabase;
+let fileName: string;
+
+const k1 = 'k1';
+const v1 = 'v1';
+
+const k2 = 'k2';
+const v2 = 'v2';
+
+const fakeDB = {
+	k1: v1,
+	k2: v2
+};
+
+// These tests don't work on Linux systems because gnome-keyring doesn't like running on headless machines.
+describe('AccountProvider.FileDatabase', function (): void {
+	beforeEach(async function (): Promise<void> {
+		fileName = crypto.randomBytes(4).toString('hex');
+		fileDatabase = new FileDatabase(path.join(os.tmpdir(), fileName));
+	});
+
+	it('set, get, and clear', async function (): Promise<void> {
+		await fileDatabase.initialize();
+		await fileDatabase.set(k1, v1);
+
+		let x = fileDatabase.get(k1);
+		should(x).be.equal(v1);
+
+		await fileDatabase.clear();
+
+		x = fileDatabase.get(k1);
+		should(x).be.undefined();
+	});
+
+	it('read the file contents', async function (): Promise<void> {
+		await fileDatabase.initialize();
+		await fileDatabase.set(k1, v1);
+
+		let x = fileDatabase.get(k1);
+		should(x).be.equal(v1);
+
+		await fileDatabase.shutdown();
+		const data = await fs.readFile(path.join(os.tmpdir(), fileName));
+
+		should(data.toString()).be.equal(JSON.stringify({ k1: v1 }));
+	});
+
+	it('delete prefix', async function (): Promise<void> {
+		await fileDatabase.initialize();
+		await Promise.all([fileDatabase.set(k1, v1), fileDatabase.set(k2, v2)]);
+
+		let x = fileDatabase.get(k1);
+		should(x).be.equal(v1);
+
+		x = fileDatabase.get(k2);
+		should(x).be.equal(v2);
+
+		await fileDatabase.deletePrefix('k');
+
+		x = await fileDatabase.get(k1);
+		should(x).be.undefined();
+	});
+
+	it('Test write hook', async function (): Promise<void> {
+		fileDatabase.setWriteHook(async (contents): Promise<string> => {
+			should(contents).be.equal(JSON.stringify(fakeDB));
+			return contents;
+		});
+
+		await fileDatabase.initialize();
+		await fileDatabase.set(k1, v1);
+		await fileDatabase.set(k2, v2);
+		await fileDatabase.save();
+	});
+
+	it('Test read hook', async function (): Promise<void> {
+		fileDatabase.setReadHook(async (contents): Promise<string> => {
+			should(contents).be.equal(JSON.stringify(fakeDB));
+			return contents;
+		});
+		await fs.writeFile(path.join(os.tmpdir(), fileName), JSON.stringify(fakeDB));
+		await fileDatabase.initialize();
+	});
+});

extensions/azurecore/src/test/account-provider/utils/memoryDatabase.test.ts

@@ -0,0 +1,34 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as should from 'should';
+import 'mocha';
+
+import { MemoryDatabase } from '../../../account-provider/utils/memoryDatabase';
+
+let memoryDatabase: MemoryDatabase<string>;
+
+// These tests don't work on Linux systems because gnome-keyring doesn't like running on headless machines.
+describe('AccountProvider.MemoryDatabase', function (): void {
+	beforeEach(function (): void {
+		memoryDatabase = new MemoryDatabase<string>();
+	});
+	it('set, get, and clear', async function (): Promise<void> {
+		memoryDatabase.set('k1', 'v1');
+
+		let val = memoryDatabase.get('k1');
+		should(val).equal('v1');
+
+		memoryDatabase.set('k1', 'v2');
+		val = memoryDatabase.get('k1');
+		should(val).be.equal('v2');
+
+		memoryDatabase.delete('k1');
+		val = memoryDatabase.get('k1');
+
+		should(val).be.undefined();
+	});
+
+});