Skip to content

Commit

Permalink
fix(mssql): simplify insert logic
Browse files Browse the repository at this point in the history
  • Loading branch information
@lohart13
lohart13 committed Feb 1, 2024
1 parent d89a6f3 commit 4fe4828
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 163 deletions.
240 changes: 79 additions & 161 deletions packages/core/src/dialects/mssql/query-generator-typescript.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ export class MsSqlQueryGeneratorTypeScript extends AbstractQueryGenerator {
throw new Error('Cannot use "returning" option with no attributes');
}

let removedAutoIncrement = false;
const rowValues = valueHashes.map(row => {
if (typeof row !== 'object' || row == null || Array.isArray(row)) {
throw new Error(`Invalid row: ${inspect(row)}. Expected an object.`);
Expand All @@ -438,9 +439,19 @@ export class MsSqlQueryGeneratorTypeScript extends AbstractQueryGenerator {
const valueMap = new Map<string, string>();
for (const column of allColumns) {
const rowValue = row[column];
if (attributeMap.get(column)?.autoIncrement && rowValue == null) {
if (attributeMap.get(column)?.autoIncrement) {
// MS SQL Server does not support inserting null values into autoIncrement columns
continue;
if (rowValue == null) {
removedAutoIncrement = true;
continue;
} else if (removedAutoIncrement) {
throw new Error(`Cannot insert a mixture of null and non-null values into an autoIncrement column (${column}).`);
} else {
valueMap.set(column, this.escape(rowValue, {
...bulkInsertOptions,
type: attributeMap.get(column)?.type,
}));
}
} else if (rowValue === undefined) {
// Treat undefined values as DEFAULT (where supported) or NULL (where not supported)
valueMap.set(column, 'DEFAULT');
Expand All @@ -455,129 +466,70 @@ export class MsSqlQueryGeneratorTypeScript extends AbstractQueryGenerator {
return valueMap;
});

let hasInsertStatement = false;
let hasIdentityInsertStatement = false;
if (removedAutoIncrement) {
// Remove autoIncrement columns from the list of columns
allColumns.delete([...allColumns].find(column => attributeMap.get(column)?.autoIncrement)!);
}

const returnFields = this.getReturnFields(bulkInsertOptions, attributeMap);
const autoIncrementColumn = [...allColumns].find(column => attributeMap.get(column)?.autoIncrement);
const outputTableFragment = bulkInsertOptions.returning ? this._generateOutputTableFragment(returnFields, attributeMap) : '';
const returningFragment = bulkInsertOptions.returning
? joinSQLFragments(['OUTPUT', returnFields.map(field => `INSERTED.${field}`).join(', '), 'INTO @output_table'])
: '';

const queries: string[] = [];
if (outputTableFragment) {
queries.push(outputTableFragment);
if (bulkInsertOptions.returning) {
// Due to how the mssql query is built, an output table must be created before the insert statement
// as the query can be split into multiple statements.
queries.push(this._generateOutputTableFragment(returnFields, attributeMap));
}

if (autoIncrementColumn) {
let queryCount = 0;
let identityRemoved = false;
for (const row of rowValues) {
if (row.get(autoIncrementColumn) === undefined) {
if (hasIdentityInsertStatement) {
// If the identity insert statement has been used, close it
hasInsertStatement = false;
hasIdentityInsertStatement = false;
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} OFF`);
}
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} ON`);
}

if (row.size === 0) {
// If the only value is the autoIncrement key, use default values
identityRemoved = false;
hasInsertStatement = false;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
returningFragment,
'DEFAULT VALUES',
]));
} else {
// If the autoIncrement key is null or undefined, omit it from the insert statement
identityRemoved = true;
const rowFragment = [...row.values()].join(',');
const columnFragment = [...row.keys()].map(column => this.quoteIdentifier(column)).join(',');
if (hasInsertStatement && identityRemoved && queryCount < 1000) {
// If the table has already been inserted into, append the row to the insert statement
// as long as there is less than 1000 rows in the statement (SQL Server limit)
queryCount++;
queries.push(`${queries.pop()},(${rowFragment})`);
} else {
// If there is no insert statement, start one
queryCount = 1;
hasInsertStatement = true;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}
}
} else {
identityRemoved = false;
if (!hasIdentityInsertStatement) {
// If the identity insert statement has not been used, open it
hasInsertStatement = false;
hasIdentityInsertStatement = true;
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} ON`);
}
let queryCount = 0;
let hasInsertStatement = false;
const columnFragment = [...allColumns].map(column => this.quoteIdentifier(column)).join(',');

const rowFragment = [...row.values()].join(',');
const columnFragment = [...row.keys()].map(column => this.quoteIdentifier(column)).join(',');
if (hasInsertStatement && queryCount < 1000) {
// If the table has already been inserted into, append the row to the insert statement
// as long as there is less than 1000 rows in the statement (SQL Server limit)
queryCount++;
queries.push(`${queries.pop()},(${rowFragment})`);
} else {
// If there is no insert statement, start one
queryCount = 1;
hasInsertStatement = true;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}
}
for (const row of rowValues) {
if (row.size === 0) {
hasInsertStatement = false;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
returningFragment,
'DEFAULT VALUES',
]));
continue;
}
} else {
let queryCount = 0;
for (const row of rowValues) {
const rowFragment = [...row.values()].join(',');
const columnFragment = [...row.keys()].map(column => this.quoteIdentifier(column)).join(',');
if (hasInsertStatement && queryCount < 1000) {
// If the table has already been inserted into, append the row to the insert statement
// as long as there is less than 1000 rows in the statement (SQL Server limit)
queryCount++;
queries.push(`${queries.pop()},(${rowFragment})`);
} else {
// If there is no insert statement, start one
queryCount = 1;
hasInsertStatement = true;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}

const rowFragment = [...row.values()].join(',');
if (hasInsertStatement && queryCount < 1000) {
// If the table has already been inserted into, append the row to the insert statement
// as long as there is less than 1000 rows in the statement (SQL Server limit)
queryCount++;
queries.push(`${queries.pop()},(${rowFragment})`);
} else {
// If there is no insert statement, start one
queryCount = 1;
hasInsertStatement = true;
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}
}

if (hasIdentityInsertStatement) {
// If the identity insert statement has been used, close it
if (autoIncrementColumn) {
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} OFF`);
}

if (outputTableFragment) {
if (bulkInsertOptions.returning) {
queries.push('SELECT * FROM @output_table');
}

Expand Down Expand Up @@ -657,65 +609,31 @@ export class MsSqlQueryGeneratorTypeScript extends AbstractQueryGenerator {
}

const autoIncrementColumn = [...valueMap.keys()].find(column => attributeMap.get(column)?.autoIncrement);
const outputTableFragment = insertOptions.returning && insertOptions.hasTrigger ? this._generateOutputTableFragment(returnFields, attributeMap) : '';
const queries: string[] = [];
if (outputTableFragment) {
queries.push(outputTableFragment);
if (insertOptions.returning && insertOptions.hasTrigger) {
queries.push(this._generateOutputTableFragment(returnFields, attributeMap));
}

if (autoIncrementColumn) {
if (valueMap.get(autoIncrementColumn) === undefined) {
if (valueMap.size === 0) {
// If the only value is the autoIncrement key, use default values
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
returningFragment,
'DEFAULT VALUES',
]));
} else {
// If the autoIncrement key is null or undefined, omit it from the insert statement
const rowFragment = [...valueMap.values()].join(',');
const columnFragment = [...valueMap.keys()].map(column => this.quoteIdentifier(column)).join(',');
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}
} else {
const rowFragment = [...valueMap.values()].join(',');
const columnFragment = [...valueMap.keys()].map(column => this.quoteIdentifier(column)).join(',');
queries.push(
`SET IDENTITY_INSERT ${this.quoteTable(tableName)} ON`,
joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]),
`SET IDENTITY_INSERT ${this.quoteTable(tableName)} OFF`,
);
}
} else {
const rowFragment = [...valueMap.values()].join(',');
const columnFragment = [...valueMap.keys()].map(column => this.quoteIdentifier(column)).join(',');
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));
}

if (outputTableFragment) {
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} ON`);
}

const rowFragment = [...valueMap.values()].join(',');
const columnFragment = [...valueMap.keys()].map(column => this.quoteIdentifier(column)).join(',');
queries.push(joinSQLFragments([
'INSERT INTO',
this.quoteTable(tableName),
`(${columnFragment})`,
returningFragment,
'VALUES',
`(${rowFragment})`,
]));

if (autoIncrementColumn) {
queries.push(`SET IDENTITY_INSERT ${this.quoteTable(tableName)} OFF`);
}

if (insertOptions.returning && insertOptions.hasTrigger) {
queries.push('SELECT * FROM @output_table');
}

Expand Down
4 changes: 2 additions & 2 deletions packages/core/test/unit/query-generator/bulk-insert-query.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,9 +37,9 @@ describe('QueryGenerator#bulkInsertQuery', () => {
});

it('allow bulk insert with a mixture of null and values for primary key', () => {
expectsql(queryGenerator.bulkInsertQuery(User, [{ id: 1 }, { id: null }, { id: 3, fullName: 'foo' }]), {
expectsql(() => queryGenerator.bulkInsertQuery(User, [{ id: 1 }, { id: null }, { id: 3, fullName: 'foo' }]), {
default: `INSERT INTO [Users] ([id],[fullName]) VALUES (1,DEFAULT),(DEFAULT,DEFAULT),(3,'foo')`,
mssql: `SET IDENTITY_INSERT [Users] ON;INSERT INTO [Users] ([id],[fullName]) VALUES (1,DEFAULT);SET IDENTITY_INSERT [Users] OFF;INSERT INTO [Users] ([fullName]) VALUES (DEFAULT);SET IDENTITY_INSERT [Users] ON;INSERT INTO [Users] ([id],[fullName]) VALUES (3,N'foo');SET IDENTITY_INSERT [Users] OFF`,
mssql: new Error(`Cannot insert a mixture of null and non-null values into an autoIncrement column (id).`),
sqlite: `INSERT INTO \`Users\` (\`id\`,\`fullName\`) VALUES (1,NULL),(NULL,NULL),(3,'foo')`,
});
});
Expand Down

0 comments on commit 4fe4828

Please sign in to comment.