console/fix: changing primary key removes the existing PK when failing to create PK

https://github.com/hasura/graphql-engine-mono/pull/1509

GitOrigin-RevId: b4fec58f399b21125dcacf8083b01148778f97b0

console/src/components/Services/Data/TableModify/ModifyActions.js

@@ -454,7 +454,8 @@ const savePrimaryKeys = (tableName, schemaName, constraintName) => {
       return dispatch(showSuccessNotification('No changes'));
     }
     const migrationUp = [];
-    if (constraintName) {
+    // droping PK
+    if (constraintName && !numSelectedPkColumns) {
       migrationUp.push(
         getRunSqlQuery(
           dataSource.getDropConstraintSql(
@@ -466,8 +467,22 @@ const savePrimaryKeys = (tableName, schemaName, constraintName) => {
         )
       );
     }
-    // skip creating a new config if no columns were selected
+    // Altering PK
-    if (numSelectedPkColumns) {
+    else if (constraintName && numSelectedPkColumns) {
+      migrationUp.push(
+        getRunSqlQuery(
+          dataSource.getAlterPkSql({
+            schemaName,
+            tableName,
+            selectedPkColumns,
+            constraintName,
+          }),
+          source
+        )
+      );
+    }
+    // Creating a new PK entry
+    else if (!constraintName && numSelectedPkColumns) {
       migrationUp.push(
         getRunSqlQuery(
           dataSource.getCreatePkSql({

console/src/dataSources/index.ts

@@ -285,6 +285,17 @@ export interface DataSourcesAPI {
     selectedPkColumns: string[];
     constraintName?: string;
   }) => string;
+  getAlterPkSql: ({
+    schemaName,
+    tableName,
+    selectedPkColumns,
+    constraintName,
+  }: {
+    schemaName: string;
+    tableName: string;
+    selectedPkColumns: string[];
+    constraintName: string;
+  }) => string;
   getFunctionDefinitionSql:
     | null
     | ((

console/src/dataSources/services/bigquery/index.tsx

@@ -359,6 +359,9 @@ export const bigquery: DataSourcesAPI = {
   getCreatePkSql: () => {
     return '';
   },
+  getAlterPkSql: () => {
+    return '';
+  },
   getFunctionDefinitionSql: null,
   primaryKeysInfoSql: () => {
     return 'select []';

console/src/dataSources/services/mssql/__tests__/__snapshots__/mssql.test.ts.snap

@@ -42,6 +42,24 @@ exports[`mssql datasource tests getAlterFKSql should generate SQL query for alte
     "
 `;
 
+exports[`mssql datasource tests getAlterPkSql should generate alter operation as a single transaction  1`] = `
+"BEGIN TRANSACTION;
+    ALTER TABLE \\"public\\".\\"users\\" DROP CONSTRAINT \\"PK__users__1234\\";
+    ALTER TABLE \\"public\\".\\"users\\"
+      ADD CONSTRAINT \\"PK__users__1234\\" PRIMARY KEY (\\"id\\");
+    
+    COMMIT TRANSACTION;"
+`;
+
+exports[`mssql datasource tests getAlterPkSql should work with multi-column PKs  1`] = `
+"BEGIN TRANSACTION;
+    ALTER TABLE \\"public\\".\\"users\\" DROP CONSTRAINT \\"test_constraint\\";
+    ALTER TABLE \\"public\\".\\"users\\"
+      ADD CONSTRAINT \\"test_constraint\\" PRIMARY KEY (\\"id\\", \\"account\\");
+    
+    COMMIT TRANSACTION;"
+`;
+
 exports[`mssql datasource tests getCreateFKeySql should generate query for create foreign keys with multiple columns 1`] = `
 "
     ALTER TABLE \\"dbo\\".\\"user\\"

console/src/dataSources/services/mssql/__tests__/mssql.test.ts

@@ -138,6 +138,36 @@ describe('mssql datasource tests', () => {
     });
   });
 
+  describe('getAlterPkSql', () => {
+    const { getAlterPkSql } = mssql;
+    it('should generate alter operation as a single transaction ', () => {
+      const query = getAlterPkSql({
+        schemaName: 'public',
+        tableName: 'users',
+        selectedPkColumns: ['id'],
+        constraintName: 'PK__users__1234',
+      });
+      expect(query).toContain('BEGIN TRANSACTION');
+      expect(query).toContain('ALTER TABLE "public"."users"');
+      expect(query).toContain('DROP CONSTRAINT "PK__users__1234"');
+      expect(query).toContain('ADD CONSTRAINT "PK__users__1234"');
+      expect(query).toContain('COMMIT TRANSACTION');
+      expect(query).toMatchSnapshot();
+    });
+    it('should work with multi-column PKs ', () => {
+      const query = getAlterPkSql({
+        schemaName: 'public',
+        tableName: 'users',
+        selectedPkColumns: ['id', 'account'],
+        constraintName: 'test_constraint',
+      });
+      expect(query).toContain(
+        `ADD CONSTRAINT "test_constraint" PRIMARY KEY ("id", "account")`
+      );
+      expect(query).toMatchSnapshot();
+    });
+  });
+
   describe('getAddColumnSql', () => {
     const { getAddColumnSql } = mssql;
     it('should generate SQL query for adding a column with nullable:false, unique:false with default', () => {

console/src/dataSources/services/mssql/index.tsx

@@ -708,6 +708,26 @@ FROM sys.objects as obj
     ADD CONSTRAINT "${constraintName}"
     PRIMARY KEY (${selectedPkColumns.map(pkc => `"${pkc}"`).join(',')})`;
   },
+  getAlterPkSql: ({
+    schemaName,
+    tableName,
+    selectedPkColumns,
+    constraintName,
+  }: {
+    schemaName: string;
+    tableName: string;
+    selectedPkColumns: string[];
+    constraintName: string; // compulsory for PG
+  }) => {
+    return `BEGIN TRANSACTION;
+    ALTER TABLE "${schemaName}"."${tableName}" DROP CONSTRAINT "${constraintName}";
+    ALTER TABLE "${schemaName}"."${tableName}"
+      ADD CONSTRAINT "${constraintName}" PRIMARY KEY (${selectedPkColumns
+      .map(pkc => `"${pkc}"`)
+      .join(', ')});
+    
+    COMMIT TRANSACTION;`;
+  },
   getFunctionDefinitionSql: null,
   primaryKeysInfoSql: ({ schemas }) => {
     let whereClause = '';

console/src/dataSources/services/mysql/index.tsx

@@ -20,6 +20,7 @@ import {
   checkSchemaModification,
   getCreateCheckConstraintSql,
   getCreatePkSql,
+  getAlterPkSql,
   getCreateTriggerSql,
   getDropSql,
   getDropSchemaSql,
@@ -232,6 +233,7 @@ WHERE
   checkSchemaModification,
   getCreateCheckConstraintSql,
   getCreatePkSql,
+  getAlterPkSql,
   frequentlyUsedColumns: [],
   primaryKeysInfoSql,
   uniqueKeysSql,

console/src/dataSources/services/mysql/sqlUtils.ts

@@ -307,6 +307,10 @@ export const getCreatePkSql = ({
   )} add primary key (${selectedPkColumns.join(', ')});
 `;
 
+export const getAlterPkSql = () => {
+  return '';
+};
+
 export const getCreateTableQueries = (
   currentSchema: string,
   tableName: string,

console/src/dataSources/services/postgresql/__tests__/__snapshots__/postgresql.test.ts.snap

@@ -0,0 +1,19 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`postgresql datasource tests getAlterPkSql should generate alter operation as a single transaction  1`] = `
+"BEGIN TRANSACTION;
+ALTER TABLE \\"public\\".\\"users\\" DROP CONSTRAINT \\"PK__users__1234\\";
+
+ALTER TABLE \\"public\\".\\"users\\"
+    ADD CONSTRAINT \\"PK__users__1234\\" PRIMARY KEY (\\"id\\");
+COMMIT TRANSACTION;"
+`;
+
+exports[`postgresql datasource tests getAlterPkSql should work with multi-column PKs  1`] = `
+"BEGIN TRANSACTION;
+ALTER TABLE \\"public\\".\\"users\\" DROP CONSTRAINT \\"test_constraint\\";
+
+ALTER TABLE \\"public\\".\\"users\\"
+    ADD CONSTRAINT \\"test_constraint\\" PRIMARY KEY (\\"id\\", \\"account\\");
+COMMIT TRANSACTION;"
+`;

console/src/dataSources/services/postgresql/__tests__/postgresql.test.ts

@@ -0,0 +1,33 @@
+import { postgres } from '../index';
+
+describe('postgresql datasource tests', () => {
+  describe('getAlterPkSql', () => {
+    const { getAlterPkSql } = postgres;
+    it('should generate alter operation as a single transaction ', () => {
+      const query = getAlterPkSql({
+        schemaName: 'public',
+        tableName: 'users',
+        selectedPkColumns: ['id'],
+        constraintName: 'PK__users__1234',
+      });
+      expect(query).toContain('BEGIN TRANSACTION');
+      expect(query).toContain('ALTER TABLE "public"."users"');
+      expect(query).toContain('DROP CONSTRAINT "PK__users__1234"');
+      expect(query).toContain('ADD CONSTRAINT "PK__users__1234"');
+      expect(query).toContain('COMMIT TRANSACTION');
+      expect(query).toMatchSnapshot();
+    });
+    it('should work with multi-column PKs ', () => {
+      const query = getAlterPkSql({
+        schemaName: 'public',
+        tableName: 'users',
+        selectedPkColumns: ['id', 'account'],
+        constraintName: 'test_constraint',
+      });
+      expect(query).toContain(
+        `ADD CONSTRAINT "test_constraint" PRIMARY KEY ("id", "account")`
+      );
+      expect(query).toMatchSnapshot();
+    });
+  });
+});

console/src/dataSources/services/postgresql/index.tsx

@@ -45,6 +45,7 @@ import {
   checkSchemaModification,
   getCreateCheckConstraintSql,
   getCreatePkSql,
+  getAlterPkSql,
   getFunctionDefinitionSql,
   primaryKeysInfoSql,
   checkConstraintsSql,
@@ -703,6 +704,7 @@ export const postgres: DataSourcesAPI = {
   checkSchemaModification,
   getCreateCheckConstraintSql,
   getCreatePkSql,
+  getAlterPkSql,
   getFunctionDefinitionSql,
   primaryKeysInfoSql,
   checkConstraintsSql,

console/src/dataSources/services/postgresql/sqlUtils.ts

@@ -797,6 +797,26 @@ export const getCreatePkSql = ({
     add constraint "${constraintName}"
     primary key (${selectedPkColumns.map(pkc => `"${pkc}"`).join(', ')});`;
 };
+export const getAlterPkSql = ({
+  schemaName,
+  tableName,
+  selectedPkColumns,
+  constraintName,
+}: {
+  schemaName: string;
+  tableName: string;
+  selectedPkColumns: string[];
+  constraintName: string; // compulsory for PG
+}) => {
+  return `BEGIN TRANSACTION;
+ALTER TABLE "${schemaName}"."${tableName}" DROP CONSTRAINT "${constraintName}";
+
+ALTER TABLE "${schemaName}"."${tableName}"
+    ADD CONSTRAINT "${constraintName}" PRIMARY KEY (${selectedPkColumns
+    .map(pkc => `"${pkc}"`)
+    .join(', ')});
+COMMIT TRANSACTION;`;
+};
 
 const trackableFunctionsWhere = `
 AND has_variadic = FALSE