✨ backend: Adding atomicCompareAndSet operation to DB connector (#525)

tdrive/backend/node/src/core/platform/services/database/services/orm/connectors/abstract-connector.ts

@@ -20,6 +20,16 @@ export abstract class AbstractConnector<T extends ConnectionOptions> implements
 
   abstract upsert(entities: any[], _options: UpsertOptions): Promise<boolean[]>;
 
+  abstract atomicCompareAndSet<Entity, FieldValueType>(
+    entity: Entity,
+    fieldName: keyof Entity,
+    previousValue: FieldValueType,
+    newValue: FieldValueType,
+  ): Promise<{
+    didSet: boolean;
+    currentValue: FieldValueType | null;
+  }>;
+
   abstract remove(entities: any[]): Promise<boolean[]>;
 
   abstract find<EntityType>(

tdrive/backend/node/src/core/platform/services/database/services/orm/connectors/index.ts

@@ -49,6 +49,34 @@ export interface Connector extends Initializable {
    */
   upsert(entities: any[], _options: UpsertOptions): Promise<boolean[]>;
 
+  /**
+   * Updates an entity's field `fieldName` to `newValue`, only if its
+   * current value in the DB is `previousValue`. Does so atomically
+   * (for whatever that means for the specific DB type - mostly a
+   * single query, with no special write concern etc.).
+   *
+   * Note that the `currentValue` returned may not be atomic depending
+   * on the DB type.
+   *
+   * A single entity for that primary key must already exist. This is not
+   * always tested, and a return without error should not be interpreted as
+   * such a row existing.
+   *
+   * @param entity Entity to try to update
+   * @param fieldName Name of the field to compare and possibly set (the node.js field name)
+   * @param previousValue Value of that field expected in the DB
+   * @param newValue New value to assign if the existing value matches `previousValue`
+   */
+  atomicCompareAndSet<Entity, FieldValueType>(
+    entity: Entity,
+    fieldName: keyof Entity,
+    previousValue: FieldValueType | null,
+    newValue: FieldValueType | null,
+  ): Promise<{
+    didSet: boolean;
+    currentValue: FieldValueType | null;
+  }>;
+
   /**
    * Remove
    * returns true if the object was removed, false otherwise

tdrive/backend/node/src/core/platform/services/database/services/orm/connectors/mongodb/mongodb.ts

@@ -154,6 +154,65 @@ export class MongoConnector extends AbstractConnector<MongoConnectionOptions> {
     });
   }
 
+  async atomicCompareAndSet<Entity, FieldValueType>(
+    entity: Entity,
+    fieldName: keyof Entity,
+    previousValue: FieldValueType,
+    newValue: FieldValueType,
+  ): Promise<{
+    didSet: boolean;
+    currentValue: FieldValueType | null;
+  }> {
+    const { columnsDefinition, entityDefinition } = getEntityDefinition(entity);
+    const primaryKey = unwrapPrimarykey(entityDefinition);
+    const db = await this.getDatabase();
+    const collection = db.collection(`${entityDefinition.name}`);
+    const columnName = (Object.entries(columnsDefinition).filter(
+      ([, { nodename }]) => nodename === fieldName,
+    )[0] ?? [])[0];
+    const transformValue = (key, value) =>
+      transformValueToDbString(value, columnsDefinition[key].type, {
+        columns: columnsDefinition[key].options,
+        secret: this.secret,
+        disableSalts: true,
+        column: { key },
+      });
+    const where: any = {};
+    primaryKey.forEach(key => {
+      where[key] = transformValue(key, entity[columnsDefinition[key].nodename]);
+    });
+    where[columnName] = transformValue(columnName, previousValue);
+    const set = { [columnName]: newValue === undefined ? null : newValue };
+
+    const result = await collection.updateOne(where, { $set: set });
+    if (result.modifiedCount > 1)
+      throw new Error(
+        `Unexpected modified count ${JSON.stringify(result)} on mongo update(${JSON.stringify(
+          where,
+        )}, ${JSON.stringify(set)})`,
+      );
+
+    let currentValue: FieldValueType;
+    if (result.modifiedCount > 0) {
+      currentValue = newValue;
+    } else {
+      delete where[columnName];
+      const existingItem = await collection.findOne(where, {});
+      if (!existingItem)
+        throw new Error(
+          `Error setting ${entityDefinition.name}'s ${
+            fieldName as string
+          } atomically, no row matched PK: ${JSON.stringify(where)}`,
+        );
+      currentValue = existingItem[fieldName as string];
+    }
+
+    return {
+      didSet: result.modifiedCount === 1,
+      currentValue,
+    };
+  }
+
   async remove(entities: any[]): Promise<boolean[]> {
     return new Promise(async resolve => {
       const promises: Promise<mongo.DeleteResult>[] = [];

tdrive/backend/node/src/core/platform/services/database/services/orm/connectors/postgres/postgres-query-builder.ts

@@ -1,5 +1,5 @@
 import { comparisonType, FindOptions } from "../../repository/repository";
-import { ObjectType } from "../../types";
+import { ColumnDefinition, ObjectType } from "../../types";
 import { getEntityDefinition, unwrapPrimarykey } from "../../utils";
 import { PostgresDataTransformer } from "./postgres-data-transform";
 import _ from "lodash";
@@ -109,21 +109,33 @@ export class PostgresQueryBuilder {
     return [query(whereClause, orderByClause, limit, offset), values];
   }
 
+  private toValueKeyDBStringPairFromValue(
+    columnsDefinition: { [key: string]: ColumnDefinition },
+    key: string,
+    value: any,
+  ): [string, any] {
+    return [key, this.dataTransformer.toDbString(value, columnsDefinition[key].type)];
+  }
+
+  private toValueKeyDBStringPairFromEntity<Entity>(
+    entity: Entity,
+    columnsDefinition: { [key: string]: ColumnDefinition },
+    key: string,
+  ) {
+    return this.toValueKeyDBStringPairFromValue(
+      columnsDefinition,
+      key,
+      entity[columnsDefinition[key].nodename],
+    );
+  }
+
   buildDelete<Entity>(entity: Entity): Query {
     const { columnsDefinition, entityDefinition } = getEntityDefinition(entity);
     const primaryKey = unwrapPrimarykey(entityDefinition);
 
-    const toValueKeyDBStringPair = (key: string) => {
+    const where = primaryKey.map(key =>
-      return [
+      this.toValueKeyDBStringPairFromEntity(entity, columnsDefinition, key),
-        key,
+    );
-        this.dataTransformer.toDbString(
-          entity[columnsDefinition[key].nodename],
-          columnsDefinition[key].type,
-        ),
-      ];
-    };
-
-    const where = primaryKey.map(key => toValueKeyDBStringPair(key));
     const query = `DELETE FROM "${entityDefinition.name}"
                 WHERE ${where.map((e, idx) => `${e[0]} = $${idx + 1}`).join(" AND ")}`;
 
@@ -132,54 +144,91 @@ export class PostgresQueryBuilder {
 
   buildInsert<Entity>(entity: Entity): Query {
     const { columnsDefinition, entityDefinition } = getEntityDefinition(entity);
-    const toValueKeyDBStringPair = (key: string) => {
-      return [
-        key,
-        this.dataTransformer.toDbString(
-          entity[columnsDefinition[key].nodename],
-          columnsDefinition[key].type,
-        ),
-      ];
-    };
 
     const fields = Object.keys(columnsDefinition)
       .filter(key => entity[columnsDefinition[key].nodename] !== undefined)
-      .map(key => toValueKeyDBStringPair(key));
+      .map(key => this.toValueKeyDBStringPairFromEntity(entity, columnsDefinition, key));
     const query = `INSERT INTO "${entityDefinition.name}" (${fields.map(e => e[0]).join(", ")})
               VALUES (${fields.map((e, idx) => `$${idx + 1}`).join(", ")})`;
     return [query, fields.map(f => f[1])];
   }
 
+  private buildWhereClause(startIndex: number, where: [string, any[]][]): Query {
+    const equalsOrIs = ([key, value]) =>
+      [key, value === null ? "IS" : "=", value === null ? "NULL" : `$${startIndex++}`].join(" ");
+    return [
+      `WHERE ${where.map(equalsOrIs).join(" AND ")} `,
+      where.filter(f => f[1] !== null).map(f => f[1]),
+    ];
+  }
+
+  private buildUpdateQueryString(
+    tableName: string,
+    set: [string, any][],
+    where: [string, any][],
+    returning?: string,
+  ): Query {
+    // if (!set || set.length == 0)
+    //   throw new Error(`Cannot build UPDATE query with empty SET on ${tableName} with WHERE: ${where}`);
+    const [whereClause, whereParameters] = this.buildWhereClause(set.length + 1, where);
+    const query = `UPDATE "${tableName}" 
+              SET ${set.map((e, idx) => `${e[0]} = $${idx + 1}`).join(", ")} 
+              ${whereClause} 
+              ${returning ? "RETURNING " + returning : ""} `;
+    const values = [];
+    values.push(...set.map(f => f[1]));
+    values.push(...whereParameters);
+    return [query, values];
+  }
+
   buildUpdate<Entity>(entity: Entity): Query {
     const { columnsDefinition, entityDefinition } = getEntityDefinition(entity);
     const primaryKey = unwrapPrimarykey(entityDefinition);
 
-    const toValueKeyDBStringPair = (key: string) => {
-      return [
-        key,
-        this.dataTransformer.toDbString(
-          entity[columnsDefinition[key].nodename],
-          columnsDefinition[key].type,
-        ),
-      ];
-    };
-
     // Set updated content
     const set = Object.keys(columnsDefinition)
       .filter(key => primaryKey.indexOf(key) === -1)
       .filter(key => entity[columnsDefinition[key].nodename] !== undefined)
-      .map(key => toValueKeyDBStringPair(key));
+      .map(key => this.toValueKeyDBStringPairFromEntity(entity, columnsDefinition, key));
     //Set primary key
-    const where = primaryKey.map(key => toValueKeyDBStringPair(key));
+    const where = primaryKey.map(key =>
-    //Start index for where clause params
+      this.toValueKeyDBStringPairFromEntity(entity, columnsDefinition, key),
-    const whereIdx = set.length + 1;
+    );
-    const query = `UPDATE "${entityDefinition.name}" 
+    return this.buildUpdateQueryString(entityDefinition.name, set, where);
-              SET ${set.map((e, idx) => `${e[0]} = $${idx + 1}`).join(", ")} 
+  }
-              WHERE ${where.map((e, idx) => `${e[0]} = $${whereIdx + idx}`).join(" AND ")}`;
-    const values = [];
-    values.push(...set.map(f => f[1]));
-    values.push(...where.map(f => f[1]));
 
-    return [query, values];
+  buildatomicCompareAndSet<Entity, FieldValueType>(
+    entity: Entity,
+    fieldName: keyof Entity,
+    previousValue: FieldValueType,
+    newValue: FieldValueType,
+  ) {
+    const { columnsDefinition, entityDefinition } = getEntityDefinition(entity);
+    const primaryKey = unwrapPrimarykey(entityDefinition);
+    const columnName = (Object.entries(columnsDefinition).filter(
+      ([, { nodename }]) => nodename === fieldName,
+    )[0] ?? [])[0];
+    if (!columnName)
+      throw new Error(
+        `Can't find field ${JSON.stringify(fieldName)} in ${JSON.stringify(columnsDefinition)}`,
+      );
+    const where = primaryKey.map(key =>
+      this.toValueKeyDBStringPairFromEntity(entity, columnsDefinition, key),
+    );
+    const whereQueryWithoutField = this.buildWhereClause(0 + 1, where);
+    where.push(
+      this.toValueKeyDBStringPairFromValue(columnsDefinition, columnName as string, previousValue),
+    );
+    return {
+      updateQuery: this.buildUpdateQueryString(
+        entityDefinition.name,
+        [this.toValueKeyDBStringPairFromValue(columnsDefinition, columnName as string, newValue)],
+        where,
+      ),
+      getValueQuery: [
+        `SELECT "${columnName}" FROM "${entityDefinition.name}" ${whereQueryWithoutField[0]}`,
+        whereQueryWithoutField[1],
+      ] as Query,
+    };
   }
 }

tdrive/backend/node/src/core/platform/services/database/services/orm/connectors/postgres/postgres.ts

@@ -270,19 +270,70 @@ export class PostgresConnector extends AbstractConnector<PostgresConnectionOptio
         return this.execute(this.queryBuilder.buildInsert(entity));
       }
     } else {
-      return false;
+      throw new Error("Missing or unknown UpsertOptions.action");
+    }
+  }
+
+  async atomicCompareAndSet<Entity, FieldValueType>(
+    entity: Entity,
+    fieldName: keyof Entity,
+    previousValue: FieldValueType,
+    newValue: FieldValueType,
+  ): Promise<{
+    didSet: boolean;
+    currentValue: FieldValueType | null;
+  }> {
+    const { updateQuery, getValueQuery } = this.queryBuilder.buildatomicCompareAndSet(
+      entity,
+      fieldName,
+      previousValue,
+      newValue,
+    );
+    const result = await this.executeRaw(updateQuery);
+    if (!result)
+      throw new Error(
+        `Error updating ${JSON.stringify(fieldName)} field of ${JSON.stringify(
+          entity["id"],
+        )} from ${JSON.stringify(previousValue)} to ${JSON.stringify(
+          newValue,
+        )}. Search logs for 'services.database.orm.postgres - Error with SQL query'`,
+      );
+    if (result.rowCount > 1)
+      throw new Error(
+        `Unexpected modified count ${JSON.stringify(result)} on postgres update(${JSON.stringify(
+          updateQuery,
+        )})`,
+      );
+    let currentValue = newValue;
+    if (result.rowCount == 0) {
+      const readValue = await this.executeRaw(getValueQuery);
+      if (readValue.rows.length == 0)
+        throw new Error(
+          `Error setting ${fieldName as string} atomically, no row matched PK: ${JSON.stringify(
+            getValueQuery,
+          )}`,
+        );
+      currentValue = readValue.rows[0][readValue.fields[0].name];
+    }
+    return {
+      didSet: result.rowCount > 0,
+      currentValue,
+    };
+  }
+
+  private async executeRaw(query: Query): Promise<QueryResult | null> {
+    logger.debug(`service.database.orm.postgres - Query: "${query[0]}"`);
+    try {
+      return await this.client.query(query[0] as string, query[1] as never[]);
+    } catch (err) {
+      logger.error({ err }, `services.database.orm.postgres - Error with SQL query: ${query[0]}`);
+      return null;
     }
   }
 
   private async execute(query: Query) {
-    logger.debug(`service.database.orm.postgres - Query: "${query[0]}"`);
+    const result = await this.executeRaw(query);
-    try {
+    return result && result.rowCount > 0;
-      const result: QueryResult = await this.client.query(query[0] as string, query[1] as never[]);
-      return result.rowCount > 0;
-    } catch (err) {
-      logger.error({ err }, `services.database.orm.postgres - Error with SQL query: ${query[0]}`);
-      return false;
-    }
   }
 
   async getTableDefinition(name: string): Promise<string[]> {

tdrive/backend/node/src/core/platform/services/database/services/orm/repository/repository.ts

@@ -96,6 +96,42 @@ export default class Repository<EntityType> {
     return (await this.find(filters, options, context)).getEntities()[0] || null;
   }
 
+  /**
+   * Updates an entity's field `fieldName` to `newValue`, only if its
+   * current value in the DB is `previousValue`. Does so atomically
+   * (for whatever that means for the specific DB type - mostly a
+   * single query, with no special write concern etc.).
+   *
+   * Note that the `currentValue` returned may not be atomic depending
+   * on the DB type.
+   *
+   * Unlike the identical on the `Connector`, the manager is flushed
+   * and reset before the operation is ran.
+   *
+   * A single entity for that primary key must already exist. This is not
+   * always tested, and a return without error should not be interpreted as
+   * such a row existing.
+   *
+   * @param entity Entity to try to update
+   * @param fieldName Name of the field to compare and possibly set (the node.js field name)
+   * @param previousValue Value of that field expected in the DB
+   * @param newValue New value to assign if the existing value matches `previousValue`
+   */
+  async atomicCompareAndSet<FieldValueType>(
+    entity: EntityType,
+    fieldName: keyof EntityType,
+    previousValue: FieldValueType | null,
+    newValue: FieldValueType | null,
+  ): Promise<{
+    didSet: boolean;
+    currentValue: FieldValueType | null;
+  }> {
+    if (previousValue === newValue)
+      throw new Error(`Previous and new values are identical: ${JSON.stringify(previousValue)}`);
+    await this.manager.flush().then(manager => manager.reset());
+    return this.connector.atomicCompareAndSet(entity, fieldName, previousValue, newValue);
+  }
+
   async save(entity: EntityType, _context?: ExecutionContext): Promise<void> {
     this.manager.persist(entity);
     return this.manager.flush().then(manager => manager.reset());

tdrive/backend/node/test/unit/core/services/database/services/orm/connectors/postgres/postgres-query-builder.test.ts

@@ -1,7 +1,7 @@
 import {
   PostgresQueryBuilder
 } from "../../../../../../../../../src/core/platform/services/database/services/orm/connectors/postgres/postgres-query-builder";
-import { expect, jest } from "@jest/globals";
+import { expect, jest, test, describe, beforeEach, afterEach } from "@jest/globals";
 import { randomInt, randomUUID } from "crypto";
 import { newTestDbEntity, normalizeWhitespace, TestDbEntity } from "./utils";
 import {
@@ -127,6 +127,32 @@ describe('The PostgresQueryBuilder', () => {
     assertUpdateQueryParams(entity, query[1] as any[])
   });
 
+  test('buildatomicCompareAndSet to value', async () => {
+    const entity = newTestDbEntity();
+
+    const newValue = "new-tag-value";
+    const [queryText, params] = subj.buildatomicCompareAndSet(entity, "tags", null, newValue);
+
+    expect(normalizeWhitespace(queryText as string)).toBe(`UPDATE "test_table" SET tags = $1 WHERE company_id = $2 AND id = $3 AND tags = $4 RETURNING tags`);
+    expect(params[0]).toBe(JSON.stringify(newValue));
+    expect(params[1]).toBe(entity.company_id);
+    expect(params[2]).toBe(entity.id);
+    expect(params[3]).toBe(null);
+  });
+
+  test('buildatomicCompareAndSet to null', async () => {
+    const entity = newTestDbEntity();
+
+    const previousValue = "new-tag-value";
+    const [queryText, params] = subj.buildatomicCompareAndSet(entity, "tags", previousValue, null);
+
+    expect(normalizeWhitespace(queryText as string)).toBe(`UPDATE "test_table" SET tags = $1 WHERE company_id = $2 AND id = $3 AND tags = $4 RETURNING tags`);
+    expect(params[0]).toBe(null);
+    expect(params[1]).toBe(entity.company_id);
+    expect(params[2]).toBe(entity.id);
+    expect(params[3]).toBe(JSON.stringify(previousValue));
+  });
+
   const assertInsertQueryParams = (actual: TestDbEntity, expected: any[]) => {
     expect(expected[0]).toBe(actual.company_id);
     expect(expected[1]).toBe(actual.id);
@@ -145,5 +171,4 @@ describe('The PostgresQueryBuilder', () => {
     expect(expected[5]).toBe(actual.id);
   }
 
-
 });