From 441abf4a5da326a12eb81d5086ae887ac013cf35 Mon Sep 17 00:00:00 2001 From: Harry Brundage Date: Sat, 2 Nov 2024 12:01:06 +0000 Subject: [PATCH] feat: add support for `AddToCell` mutation elements This adds support for the newish atomic aggregation mutation kind that can be submitted within a `CheckAndMutateRow` or `MutateRow` call. --- src/mutation.ts | 82 ++++++++++++++++++++++++++++++++++++ test/mutation.ts | 107 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 189 insertions(+) diff --git a/src/mutation.ts b/src/mutation.ts index 456528c57..6d7a6990a 100644 --- a/src/mutation.ts +++ b/src/mutation.ts @@ -21,6 +21,7 @@ import {google as btTypes} from '../protos/protos'; export type IMutation = btTypes.bigtable.v2.IMutation; export type IMutateRowRequest = btTypes.bigtable.v2.IMutateRowRequest; export type ISetCell = btTypes.bigtable.v2.Mutation.ISetCell; +export type IAddToCell = btTypes.bigtable.v2.Mutation.IAddToCell; export type Bytes = string | Buffer; // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -79,6 +80,10 @@ export interface SetCellObj { [k: string]: string | ISetCell | undefined; setCell?: ISetCell; } +export interface AddToCellObj { + [k: string]: string | IAddToCell | undefined; + addToCell?: IAddToCell; +} export interface ValueObj { [k: string]: Buffer | Value | ValueObj; } @@ -366,6 +371,79 @@ export class Mutation { }); } + /** + * Formats an `add` mutation to what the proto service expects. + * + * @param {object} data - The entity data. + * @returns {object[]} + * + * @example + * ``` + * Mutation.encodeAddToCell({ + * follows: { + * gwashington: 1, + * alincoln: -1 + * } + * }); + * // [ + * // { + * // addToCell: { + * // familyName: 'follows', + * // columnQualifier: 'gwashington', // as buffer + * // timestamp: -1, // -1 means to use the server time + * // input: 1 + * // } + * // }, + * { + * addToCell: { + * familyName: 'follows', + * columnQualifier: 'alincoln', // as buffer + * timestamp: -1, // -1 means to use the server time + * input: -1 + * } + * } + * // ] + * ``` + * @private + */ + static encodeAddToCell(data: Data): AddToCellObj[] { + const mutations: AddToCellObj[] = []; + + Object.keys(data).forEach(familyName => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const family = (data as any)[familyName]; + + Object.keys(family).forEach(cellName => { + let cell = family[cellName]; + + if (!is.object(cell) || cell instanceof Buffer) { + cell = { + value: cell, + }; + } + + let timestamp = cell.timestamp || new Date(); + + if (is.date(timestamp)) { + timestamp = timestamp.getTime() * 1000; + } + + const addToCell: IAddToCell = { + familyName, + columnQualifier: {rawValue: Mutation.convertToBytes(cellName)}, + timestamp: { + rawTimestampMicros: timestamp, + }, + input: Mutation.convertToBytes(cell.value), + }; + + mutations.push({addToCell}); + }); + }); + + return mutations; + } + /** * Creates a new Mutation object and returns the proto JSON form. * @@ -431,6 +509,8 @@ export class Mutation { mutation.mutations = Mutation.encodeSetCell(this.data); } else if (this.method === Mutation.methods.DELETE) { mutation.mutations = Mutation.encodeDelete(this.data); + } else if (this.method === Mutation.methods.ADD) { + mutation.mutations = Mutation.encodeAddToCell(this.data); } return mutation; @@ -443,9 +523,11 @@ export class Mutation { * * INSERT => setCell * DELETE => deleteFrom* + * ADD => addToCell */ static methods = { INSERT: 'insert', DELETE: 'delete', + ADD: 'add', }; } diff --git a/test/mutation.ts b/test/mutation.ts index 7566a7496..d8539f935 100644 --- a/test/mutation.ts +++ b/test/mutation.ts @@ -415,6 +415,113 @@ describe('Bigtable/Mutation', () => { }); }); + describe('encodeAddToCell', () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let convertCalls: any[]; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const fakeTime = new Date('2018-1-1') as any; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const realTimestamp = new Date() as any; + + beforeEach(() => { + sandbox.stub(global, 'Date').returns(fakeTime); + convertCalls = []; + sandbox.stub(Mutation, 'convertToBytes').callsFake(value => { + convertCalls.push(value); + return value; + }); + }); + + it('should encode a addToCell mutation', () => { + const fakeMutation = { + follows: { + gwashington: 1, + alincoln: -1, + }, + }; + + const cells = Mutation.encodeAddToCell(fakeMutation); + + assert.strictEqual(cells.length, 2); + + assert.deepStrictEqual(cells, [ + { + addToCell: { + familyName: 'follows', + columnQualifier: {rawValue: 'gwashington'}, + timestamp: {rawTimestampMicros: fakeTime * 1000}, // Convert ms to μs + input: 1, + }, + }, + { + addToCell: { + familyName: 'follows', + columnQualifier: {rawValue: 'alincoln'}, + timestamp: {rawTimestampMicros: fakeTime * 1000}, // Convert ms to μs + input: -1, + }, + }, + ]); + + assert.strictEqual(convertCalls.length, 4); + assert.deepStrictEqual(convertCalls, ['gwashington', 1, 'alincoln', -1]); + }); + + it('should optionally accept a timestamp', () => { + const fakeMutation = { + follows: { + gwashington: { + value: 1, + timestamp: realTimestamp, + }, + }, + }; + + const cells = Mutation.encodeAddToCell(fakeMutation); + + assert.deepStrictEqual(cells, [ + { + addToCell: { + familyName: 'follows', + columnQualifier: {rawValue: 'gwashington'}, + timestamp: {rawTimestampMicros: realTimestamp * 1000}, // Convert ms to μs + input: 1, + }, + }, + ]); + + assert.strictEqual(convertCalls.length, 2); + assert.deepStrictEqual(convertCalls, ['gwashington', 1]); + }); + + it('should accept buffers', () => { + const val = Buffer.from([42]); // Using number 42 instead of string + const fakeMutation = { + follows: { + gwashington: val, + }, + }; + + const cells = Mutation.encodeAddToCell(fakeMutation); + + assert.deepStrictEqual(cells, [ + { + addToCell: { + familyName: 'follows', + columnQualifier: {rawValue: 'gwashington'}, + timestamp: { + rawTimestampMicros: fakeTime * 1000, + }, // Convert ms to μs + input: val, + }, + }, + ]); + + assert.strictEqual(convertCalls.length, 2); + assert.deepStrictEqual(convertCalls, ['gwashington', val]); + }); + }); + describe('parse', () => { let toProtoCalled = false; const fakeData = {a: 'a'} as IMutateRowRequest;