feat(cockroachdb): add support upsert (#4767)

Co-authored-by: Igor Savin <iselwin@gmail.com>

lib/dialects/cockroachdb/crdb-querybuilder.js

@@ -0,0 +1,11 @@
+const QueryBuilder = require('../../query/querybuilder');
+const isEmpty = require('lodash/isEmpty');
+
+module.exports = class QueryBuilder_CockroachDB extends QueryBuilder {
+  upsert(values, returning, options) {
+    this._method = 'upsert';
+    if (!isEmpty(returning)) this.returning(returning, options);
+    this._single.upsert = values;
+    return this;
+  }
+};

lib/dialects/cockroachdb/crdb-querycompiler.js

@@ -1,9 +1,63 @@
 const QueryCompiler_PG = require('../postgres/query/pg-querycompiler');
+const isEmpty = require('lodash/isEmpty');
+const { columnize: columnize_ } = require('../../formatter/wrappingFormatter');
 
 class QueryCompiler_CRDB extends QueryCompiler_PG {
   truncate() {
     return `truncate ${this.tableName}`;
   }
+
+  upsert() {
+    let sql = this._upsert();
+    if (sql === '') return sql;
+    const { returning } = this.single;
+    if (returning) sql += this._returning(returning);
+    return {
+      sql: sql,
+      returning,
+    };
+  }
+
+  _upsert() {
+    const upsertValues = this.single.upsert || [];
+    let sql = this.with() + `upsert into ${this.tableName} `;
+    if (Array.isArray(upsertValues)) {
+      if (upsertValues.length === 0) return '';
+    } else if (typeof upsertValues === 'object' && isEmpty(upsertValues)) {
+      return sql + this._emptyInsertValue;
+    }
+
+    const upsertData = this._prepInsert(upsertValues);
+    if (typeof upsertData === 'string') {
+      sql += upsertData;
+    } else {
+      if (upsertData.columns.length) {
+        sql += `(${columnize_(
+          upsertData.columns,
+          this.builder,
+          this.client,
+          this.bindingsHolder
+        )}`;
+        sql += ') values (';
+        let i = -1;
+        while (++i < upsertData.values.length) {
+          if (i !== 0) sql += '), (';
+          sql += this.client.parameterize(
+            upsertData.values[i],
+            this.client.valueForUndefined,
+            this.builder,
+            this.bindingsHolder
+          );
+        }
+        sql += ')';
+      } else if (upsertValues.length === 1 && upsertValues[0]) {
+        sql += this._emptyInsertValue;
+      } else {
+        sql = '';
+      }
+    }
+    return sql;
+  }
 }
 
 module.exports = QueryCompiler_CRDB;

lib/dialects/cockroachdb/index.js

@@ -5,6 +5,7 @@ const Transaction = require('../postgres/execution/pg-transaction');
 const QueryCompiler = require('./crdb-querycompiler');
 const TableCompiler = require('./crdb-tablecompiler');
 const ViewCompiler = require('./crdb-viewcompiler');
+const QueryBuilder = require('./crdb-querybuilder');
 
 // Always initialize with the "QueryBuilder" and "QueryCompiler"
 // objects, which extend the base 'lib/query/builder' and
@@ -26,6 +27,10 @@ class Client_CockroachDB extends Client_PostgreSQL {
     return new ViewCompiler(this, ...arguments);
   }
 
+  queryBuilder() {
+    return new QueryBuilder(this);
+  }
+
   _parseVersion(versionString) {
     return versionString.split(' ')[2];
   }

lib/query/querybuilder.js

@@ -1273,6 +1273,12 @@ class Builder extends EventEmitter {
     return this;
   }
 
+  upsert(values, returning, options) {
+    throw new Error(
+      `Upsert is not yet supported for dialect ${this.client.dialect}`
+    );
+  }
+
   _analytic(alias, second, third) {
     let analytic;
     const { schema } = this._single;

package.json

@@ -102,9 +102,9 @@
     "coveralls": "^3.1.1",
     "cross-env": "^7.0.3",
     "dtslint": "4.1.6",
-    "eslint": "^8.0.1",
+    "eslint": "^8.1.0",
     "eslint-config-prettier": "^8.3.0",
-    "eslint-plugin-import": "^2.24.2",
+    "eslint-plugin-import": "^2.25.2",
     "husky": "^4.3.8",
     "jake": "^8.1.1",
     "JSONStream": "^1.3.5",

test/integration2/dialects/cockroachdb.spec.js

@@ -0,0 +1,62 @@
+'use strict';
+
+const { expect } = require('chai');
+const {
+  Db,
+  getKnexForDb,
+  getAllDbs,
+} = require('../util/knex-instance-provider');
+
+describe('CockroachDB dialect', () => {
+  getAllDbs()
+    .filter((db) => db === Db.CockroachDB)
+    .forEach((db) => {
+      describe(db, () => {
+        let knex;
+        before(async () => {
+          knex = getKnexForDb(db);
+
+          await knex.schema.createTable('test', function () {
+            this.increments('id').primary();
+            this.decimal('balance');
+            this.string('name');
+          });
+        });
+
+        after(async () => {
+          await knex.schema.dropTable('test');
+          await knex.destroy();
+        });
+
+        describe('Upsert into query with unique', () => {
+          it('select empty table', async () => {
+            const results = await knex('test').select();
+            expect(results).is.empty;
+          });
+
+          it('upsert id=1 with result insert', async () => {
+            const results = await knex('test').upsert({ id: 1, balance: 10 });
+            expect(results.command).to.eql('INSERT');
+            expect(results.rowCount).to.eql(1);
+          });
+
+          it('select rows', async () => {
+            const results = await knex('test').select();
+            expect(results).to.eql([{ id: '1', balance: '10.00', name: null }]);
+          });
+
+          it('upsert id=1 with result update balance', async () => {
+            const results = await knex('test').upsert({ id: 1, balance: 5 });
+            expect(results.command).to.eql('INSERT');
+            expect(results.rowCount).to.eql(1);
+          });
+
+          it('select rows', async () => {
+            const results = await knex('test').select();
+            expect(results.length).to.eql(1);
+            expect(results).to.eql([{ id: '1', balance: '5.00', name: null }]);
+          });
+        });
+      });
+    });
+});

test/unit/dialects/cockroachdb.js

@@ -0,0 +1,76 @@
+'use strict';
+const expect = require('chai').expect;
+const knex = require('../../../knex');
+
+describe('CockroachDB unit tests', () => {
+  const knexInstance = knex({
+    client: 'cockroachdb',
+    connection: {
+      user: 'user',
+      password: 'password',
+      connectString: 'connect-string',
+      externalAuth: true,
+      host: 'host',
+      database: 'database',
+    },
+  });
+
+  it('correctly builds single value upsert', async () => {
+    const qb = knexInstance.client.queryBuilder();
+    qb.upsert({ value: 1 }).into('fakeTable');
+    const compiler = knexInstance.client.queryCompiler(qb);
+    const sql = compiler.upsert();
+    expect(sql).to.eql({
+      sql: 'upsert into "fakeTable" ("value") values (?)',
+      returning: undefined,
+    });
+  });
+
+  it('correctly builds default values only upsert', async () => {
+    const qb = knexInstance.client.queryBuilder();
+    qb.upsert({}).into('fakeTable');
+    const compiler = knexInstance.client.queryCompiler(qb);
+    const sql = compiler.upsert();
+    expect(sql).to.eql({
+      sql: 'upsert into "fakeTable" default values',
+      returning: undefined,
+    });
+  });
+
+  it('correctly builds default values only upsert and returning', async () => {
+    const qb = knexInstance.client.queryBuilder();
+    qb.upsert({}).into('fakeTable').returning('id');
+    const compiler = knexInstance.client.queryCompiler(qb);
+    const sql = compiler.upsert();
+    expect(sql).to.eql({
+      sql: 'upsert into "fakeTable" default values returning "id"',
+      returning: 'id',
+    });
+  });
+
+  it('correctly builds bulk values only upsert and returning', async () => {
+    const qb = knexInstance.client.queryBuilder();
+    qb.upsert([{ a: 1 }, { a: 2 }])
+      .into('fakeTable')
+      .returning('a');
+    const compiler = knexInstance.client.queryCompiler(qb);
+    const sql = compiler.upsert();
+    expect(sql).to.eql({
+      sql: 'upsert into "fakeTable" ("a") values (?), (?) returning "a"',
+      returning: 'a',
+    });
+  });
+
+  it('correctly builds bulk values different columns upsert and returning', async () => {
+    const qb = knexInstance.client.queryBuilder();
+    qb.upsert([{ a: 1 }, { b: 2 }])
+      .into('fakeTable')
+      .returning(['a', 'b']);
+    const compiler = knexInstance.client.queryCompiler(qb);
+    const sql = compiler.upsert();
+    expect(sql).to.eql({
+      sql: 'upsert into "fakeTable" ("a", "b") values (?, DEFAULT), (DEFAULT, ?) returning "a", "b"',
+      returning: ['a', 'b'],
+    });
+  });
+});

types/index.d.ts

@@ -456,10 +456,11 @@ export declare namespace Knex {
 
   type DbRecordArr<TRecord> = Readonly<MaybeArray<DbRecord<TRecord>>>;
 
-  export type CompositeTableType<TBase, TInsert = TBase, TUpdate = Partial<TInsert>> = {
+  export type CompositeTableType<TBase, TInsert = TBase, TUpdate = Partial<TInsert>, TUpsert = Partial<TInsert>> = {
     base: TBase,
     insert: TInsert,
     update: TUpdate,
+    upsert: TUpsert,
   };
 
   type TableNames = keyof Tables;
@@ -714,6 +715,75 @@ export declare namespace Knex {
         : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>
     ): QueryBuilder<TRecord, TResult2>;
 
+    upsert(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>,
+      returning: '*',
+      options?: DMLOptions
+    ): QueryBuilder<TRecord, DeferredKeySelection<TRecord, never>[]>;
+    upsert<
+      TKey extends StrKey<ResolveTableType<TRecord>>,
+      TResult2 = DeferredIndex.Augment<
+        UnwrapArrayMember<TResult>,
+        ResolveTableType<TRecord>,
+        TKey
+        >[]
+      >(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>,
+      returning: TKey,
+      options?: DMLOptions
+    ): QueryBuilder<TRecord, TResult2>;
+    upsert<
+      TKey extends StrKey<ResolveTableType<TRecord>>,
+      TResult2 = DeferredKeySelection.Augment<
+        UnwrapArrayMember<TResult>,
+        ResolveTableType<TRecord>,
+        TKey
+        >[]
+      >(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>,
+      returning: readonly TKey[],
+      options?: DMLOptions
+    ): QueryBuilder<TRecord, TResult2>;
+    upsert<
+      TKey extends string,
+      TResult2 = DeferredIndex.Augment<
+        UnwrapArrayMember<TResult>,
+        TRecord,
+        TKey
+        >[]
+      >(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>,
+      returning: TKey,
+      options?: DMLOptions
+    ): QueryBuilder<TRecord, TResult2>;
+    upsert<
+      TKey extends string,
+      TResult2 = DeferredIndex.Augment<
+        UnwrapArrayMember<TResult>,
+        TRecord,
+        TKey
+        >[]
+      >(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>,
+      returning: readonly TKey[],
+      options?: DMLOptions
+    ): QueryBuilder<TRecord, TResult2>;
+    upsert<TResult2 = number[]>(
+      data: TRecord extends CompositeTableType<unknown>
+        ? ResolveTableType<TRecord, 'upsert'> | ReadonlyArray<ResolveTableType<TRecord, 'upsert'>>
+        : DbRecordArr<TRecord> | ReadonlyArray<DbRecordArr<TRecord>>
+    ): QueryBuilder<TRecord, TResult2>;
+
     modify<TRecord2 extends {} = any, TResult2 extends {} = any>(
       callback: QueryCallbackWithArgs<TRecord, any>,
       ...args: any[]