From 13aeb6875ab32ee523268bc01edcffe79f016d40 Mon Sep 17 00:00:00 2001 From: Emily Morgan Date: Tue, 14 May 2024 10:30:26 +0200 Subject: [PATCH] feat: pgroll support import (#1470) --- cli/package.json | 3 + cli/src/commands/import/csv.ts | 83 ++++++++++++++++----- cli/src/utils/compareSchema.ts | 129 +++++++++++++++++++++++++++++++++ pnpm-lock.yaml | 28 ++++++- 4 files changed, 221 insertions(+), 22 deletions(-) create mode 100644 cli/src/utils/compareSchema.ts diff --git a/cli/package.json b/cli/package.json index 89d2b59d8..83ab461be 100644 --- a/cli/package.json +++ b/cli/package.json @@ -42,6 +42,7 @@ "ini": "^4.1.2", "lodash.compact": "^3.0.1", "lodash.get": "^4.4.2", + "lodash.keyby": "^4.6.0", "lodash.set": "^4.3.2", "node-fetch": "^3.3.2", "open": "^10.1.0", @@ -51,6 +52,7 @@ "text-table": "^0.2.0", "tmp": "^0.2.3", "tslib": "^2.6.2", + "type-fest": "^4.18.1", "which": "^4.0.0", "zod": "^3.23.8" }, @@ -59,6 +61,7 @@ "@types/babel__core": "^7.20.5", "@types/lodash.compact": "^3.0.9", "@types/lodash.get": "^4.4.9", + "@types/lodash.keyby": "^4.6.9", "@types/lodash.set": "^4.3.9", "@types/relaxed-json": "^1.0.4", "@types/text-table": "^0.2.5", diff --git a/cli/src/commands/import/csv.ts b/cli/src/commands/import/csv.ts index 762d78b64..8c1a9e3e5 100644 --- a/cli/src/commands/import/csv.ts +++ b/cli/src/commands/import/csv.ts @@ -5,7 +5,13 @@ import { importColumnTypes } from '@xata.io/importer'; import { open, writeFile } from 'fs/promises'; import { BaseCommand } from '../../base.js'; import { enumFlag } from '../../utils/oclif.js'; -import { getBranchDetailsWithPgRoll } from '../../migrations/pgroll.js'; +import { + getBranchDetailsWithPgRoll, + waitForMigrationToFinish, + xataColumnTypeToPgRollComment +} from '../../migrations/pgroll.js'; +import { compareSchemas } from '../../utils/compareSchema.js'; +import keyBy from 'lodash.keyby'; const ERROR_CONSOLE_LOG_LIMIT = 200; const ERROR_LOG_FILE = 'errors.log'; @@ -23,6 +29,8 @@ const bufferEncodings: BufferEncoding[] = [ 'hex' ]; +const INTERNAL_COLUMNS_PGROLL = ['xata_id', 'xata_createdat', 'xata_updatedat', 'xata_version']; + export default class ImportCSV extends BaseCommand { static description = 'Import a CSV file'; @@ -144,12 +152,26 @@ export default class ImportCSV extends BaseCommand { if (!parseResults.success) { throw new Error('Failed to parse CSV file'); } - const batchRows = parseResults.data.map(({ data }) => data); + const batchRows = parseResults.data.map(({ data }) => { + const formattedRow: { [k: string]: any } = {}; + const keys = Object.keys(data); + for (const key of keys) { + if (INTERNAL_COLUMNS_PGROLL.includes(key) && key !== 'xata_id') continue; + formattedRow[key] = data[key]; + } + return formattedRow; + }); + const importResult = await xata.import.importBatch( { workspace, region, database, branch }, - { columns: parseResults.columns, table, batchRows } + { + columns: parseResults.columns.filter( + ({ name }) => name === 'xata_id' || !INTERNAL_COLUMNS_PGROLL.includes(name) + ), + table, + batchRows + } ); - await xata.import.importFiles( { database, branch, region, workspace: workspace }, { @@ -212,22 +234,39 @@ export default class ImportCSV extends BaseCommand { const xata = await this.getXataClient(); const { workspace, region, database, branch } = await this.parseDatabase(); const { schema: existingSchema } = await getBranchDetailsWithPgRoll(xata, { workspace, region, database, branch }); - const newSchema = { - tables: [ - ...existingSchema.tables.filter((t) => t.name !== table), - { name: table, columns: columns.filter((c) => c.name !== 'id') } - ] - }; - const { edits } = await xata.api.migrations.compareBranchWithUserSchema({ - pathParams: { workspace, region, dbBranchName: `${database}:main` }, - body: { schema: newSchema } - }); - if (edits.operations.length > 0) { - const destructiveOperations = edits.operations + const { edits } = compareSchemas( + {}, + { + tables: { + [table]: { + name: table, + xataCompatible: false, + columns: keyBy( + columns + .filter((c) => !INTERNAL_COLUMNS_PGROLL.includes(c.name as any)) + .map((c) => { + return { + name: c.name, + type: c.type, + nullable: c.notNull !== false, + default: c.defaultValue ?? null, + unique: c.unique, + comment: xataColumnTypeToPgRollComment(c) + }; + }), + 'name' + ) + } + } + } + ); + + if (edits.length > 0) { + const destructiveOperations = edits .map((op) => { - if (!('removeColumn' in op)) return undefined; - return op.removeColumn.column; + if (!('drop_column' in op)) return undefined; + return op.drop_column.column; }) .filter((x) => x !== undefined); @@ -262,10 +301,14 @@ export default class ImportCSV extends BaseCommand { process.exit(1); } - await xata.api.migrations.applyBranchSchemaEdit({ + const { jobID } = await xata.api.migrations.applyMigration({ pathParams: { workspace, region, dbBranchName: `${database}:${branch}` }, - body: { edits } + body: { + adaptTables: true, + operations: edits + } }); + await waitForMigrationToFinish(xata.api, workspace, region, database, branch, jobID); } } } diff --git a/cli/src/utils/compareSchema.ts b/cli/src/utils/compareSchema.ts new file mode 100644 index 000000000..3bcdcabfa --- /dev/null +++ b/cli/src/utils/compareSchema.ts @@ -0,0 +1,129 @@ +import { PgRollOperation } from '@xata.io/pgroll'; +import { PartialDeep } from 'type-fest'; +import { Schemas } from '@xata.io/client'; +import { generateLinkReference, tableNameFromLinkComment, xataColumnTypeToPgRoll } from '../migrations/pgroll.js'; + +export function compareSchemas( + source: PartialDeep, + target: PartialDeep +): { edits: PgRollOperation[] } { + const edits: PgRollOperation[] = []; + + // Compare tables + const sourceTables = Object.keys(source.tables ?? {}); + const targetTables = Object.keys(target.tables ?? {}); + const newTables = targetTables.filter((table) => !sourceTables.includes(table)); + const deletedTables = sourceTables.filter((table) => !targetTables.includes(table)); + + // Compare columns + for (const table of sourceTables) { + const sourceColumns = Object.keys(source.tables?.[table]?.columns ?? {}); + const targetColumns = Object.keys(target.tables?.[table]?.columns ?? {}); + const newColumns = targetColumns.filter((column) => !sourceColumns.includes(column)); + const deletedColumns = sourceColumns.filter((column) => !targetColumns.includes(column)); + + // Add columns + for (const column of newColumns) { + const props = target.tables?.[table]?.columns?.[column] ?? {}; + edits.push({ + add_column: { + table, + column: { + name: column, + type: xataColumnTypeToPgRoll(props?.type as any), + comment: props?.comment, + nullable: !(props?.nullable === false), + unique: props?.unique, + default: props?.default ?? undefined, + references: + props?.type === 'link' && props?.name + ? generateLinkReference({ + column: props.name, + table: tableNameFromLinkComment(props?.comment ?? '') ?? '' + }) + : undefined + } + } + }); + } + + // Delete columns + for (const column of deletedColumns) { + edits.push({ drop_column: { table, column } }); + } + + // Compare column properties + for (const column of targetColumns) { + const sourceProps = source.tables?.[table]?.columns?.[column] ?? {}; + const targetProps = target.tables?.[table]?.columns?.[column] ?? {}; + + if (sourceProps.type !== targetProps.type) { + edits.push({ + alter_column: { + table, + column, + type: targetProps.type, + references: + targetProps?.type === 'link' && targetProps?.name + ? generateLinkReference({ + column: targetProps.name, + table: tableNameFromLinkComment(targetProps?.comment ?? '') ?? '' + }) + : undefined + } + }); + } + + if (sourceProps.nullable !== targetProps.nullable) { + edits.push({ alter_column: { table, column, nullable: targetProps.nullable } }); + } + + if (sourceProps.unique !== targetProps.unique) { + edits.push({ + alter_column: { + table, + column, + unique: { + name: `${table}_${column}_unique` + } + } + }); + } + } + } + + // Delete tables + for (const table of deletedTables) { + edits.push({ drop_table: { name: table } }); + } + + // Add new tables + for (const table of newTables) { + const props = target.tables?.[table] ?? {}; + edits.push({ + create_table: { + name: table, + comment: props.comment, + columns: Object.entries(props.columns ?? {}).map(([name, column]) => { + return { + name, + type: xataColumnTypeToPgRoll(column?.type as any), + comment: column?.comment, + nullable: !(column?.nullable === false), + unique: column?.unique, + default: column?.default ?? undefined, + references: + column?.type === 'link' && column?.name + ? generateLinkReference({ + column: column?.name, + table: tableNameFromLinkComment(column?.comment ?? '') ?? '' + }) + : undefined + }; + }) + } + }); + } + + return { edits }; +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 599cd3e50..21daa8c64 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -228,6 +228,9 @@ importers: lodash.get: specifier: ^4.4.2 version: 4.4.2 + lodash.keyby: + specifier: ^4.6.0 + version: 4.6.0 lodash.set: specifier: ^4.3.2 version: 4.3.2 @@ -255,6 +258,9 @@ importers: tslib: specifier: ^2.6.2 version: 2.6.2 + type-fest: + specifier: ^4.18.1 + version: 4.18.1 which: specifier: ^4.0.0 version: 4.0.0 @@ -274,6 +280,9 @@ importers: '@types/lodash.get': specifier: ^4.4.9 version: 4.4.9 + '@types/lodash.keyby': + specifier: ^4.6.9 + version: 4.6.9 '@types/lodash.set': specifier: ^4.3.9 version: 4.3.9 @@ -6285,6 +6294,12 @@ packages: '@types/lodash': 4.14.199 dev: true + /@types/lodash.keyby@4.6.9: + resolution: {integrity: sha512-N8xfQdZ2ADNPDL72TaLozIL4K1xFCMG1C1T9GN4dOFI+sn1cjl8d4U+POp8PRCAnNxDCMkYAZVD/rOBIWYPT5g==} + dependencies: + '@types/lodash': 4.14.199 + dev: true + /@types/lodash.pick@4.4.9: resolution: {integrity: sha512-hDpr96x9xHClwy1KX4/RXRejqjDFTEGbEMT3t6wYSYeFDzxmMnSKB/xHIbktRlPj8Nii2g8L5dtFDRaNFBEzUQ==} dependencies: @@ -9011,9 +9026,9 @@ packages: eslint: 9.2.0 eslint-module-utils: 2.8.0(@typescript-eslint/parser@6.21.0)(eslint-import-resolver-node@0.3.9)(eslint-import-resolver-typescript@3.6.1)(eslint@9.2.0) eslint-plugin-import: 2.29.1(@typescript-eslint/parser@6.21.0)(eslint-import-resolver-typescript@3.6.1)(eslint@9.2.0) - fast-glob: 3.3.2 + fast-glob: 3.3.1 get-tsconfig: 4.7.2 - is-core-module: 2.13.1 + is-core-module: 2.13.0 is-glob: 4.0.3 transitivePeerDependencies: - '@typescript-eslint/parser' @@ -11144,6 +11159,10 @@ packages: resolution: {integrity: sha512-oSXzaWypCMHkPC3NvBEaPHf0KsA5mvPrOPgQWDsbg8n7orZ290M0BmC/jgRZ4vcJ6DTAhjrsSYgdsW/F+MFOBA==} dev: false + /lodash.keyby@4.6.0: + resolution: {integrity: sha512-PRe4Cn20oJM2Sn6ljcZMeKgyhTHpzvzFmdsp9rK+6K0eJs6Tws0MqgGFpfX/o2HjcoQcBny1Eik9W7BnVTzjIQ==} + dev: false + /lodash.merge@4.6.2: resolution: {integrity: sha512-0KpjqXRVvrYyCsX1swR/XTK0va6VQkQM6MNo7PqW77ByjAhoARA8EfrP1N4+KlKj8YS0ZUCtRT/YUuhyYDujIQ==} @@ -14627,6 +14646,11 @@ packages: resolution: {integrity: sha512-tLq3bSNx+xSpwvAJnzrK0Ep5CLNWjvFTOp71URMaAEWBfRb9nnJiBoUe0tF8bI4ZFO3omgBR6NvnbzVUT3Ly4g==} engines: {node: '>=14.16'} + /type-fest@4.18.1: + resolution: {integrity: sha512-qXhgeNsX15bM63h5aapNFcQid9jRF/l3ojDoDFmekDQEUufZ9U4ErVt6SjDxnHp48Ltrw616R8yNc3giJ3KvVQ==} + engines: {node: '>=16'} + dev: false + /type-fest@4.9.0: resolution: {integrity: sha512-KS/6lh/ynPGiHD/LnAobrEFq3Ad4pBzOlJ1wAnJx9N4EYoqFhMfLIBjUT2UEx4wg5ZE+cC1ob6DCSpppVo+rtg==} engines: {node: '>=16'}