Backend functions to modify table primary keys

db/columns.py

@@ -121,13 +121,18 @@ def _transform_column_alter_dict(data):
     return alter_def
 
 
-def add_pkey_column_to_table(table_oid, pkey_type, conn, name="id"):
+def add_pkey_column_to_table(
+        table_oid, pkey_type, conn, drop_old_pkey_column=False, name="id"
+):
     """
     Add a primary key column to a table.
 
     See the `msar.add_pkey_column` function for info on the arguments.
     """
-    db_conn.exec_msar_func(conn, 'add_pkey_column', table_oid, pkey_type, name)
+    return db_conn.exec_msar_func(
+        conn, 'add_pkey_column',
+        table_oid, pkey_type, drop_old_pkey_column, name
+    ).fetchone()[0]
 
 
 def add_columns_to_table(table_oid, column_data_list, conn):

db/sql/00_msar_all_objects_table.sql

@@ -570,6 +570,7 @@ INSERT INTO msar.all_mathesar_objects VALUES
   ('msar', 'msar.add_mathesar_table(oid,text,jsonb,jsonb,regrole,text)', 'FUNCTION', NULL),
   ('msar', 'msar.add_mathesar_table(oid,text,jsonb,jsonb,text)', 'FUNCTION', NULL),
   ('msar', 'msar.add_month_to_vector(point,date)', 'FUNCTION', NULL),
+  ('msar', 'msar.add_pkey_column(regclass,msar.pkey_kind,boolean,text)', 'FUNCTION', NULL),
   ('msar', 'msar.add_pkey_column(regclass,msar.pkey_kind,text)', 'FUNCTION', NULL),
   ('msar', 'msar.add_record_to_table(oid,jsonb,boolean)', 'FUNCTION', NULL),
   ('msar', 'msar.add_record_to_table(oid,jsonb,boolean,jsonb)', 'FUNCTION', NULL),
@@ -1171,6 +1172,7 @@ INSERT INTO msar.all_mathesar_objects VALUES
   ('msar', 'msar.search_records_from_table(oid,jsonb,integer,boolean)', 'FUNCTION', NULL),
   ('msar', 'msar.search_records_from_table(oid,jsonb,integer,boolean,jsonb)', 'FUNCTION', NULL),
   ('msar', 'msar.set_members_to_role(regrole,oid[])', 'FUNCTION', NULL),
+  ('msar', 'msar.set_pkey_column(regclass,integer,msar.pkey_kind,boolean)', 'FUNCTION', NULL),
   ('msar', 'msar.set_schema_description(oid,text)', 'FUNCTION', NULL),
   ('msar', 'msar.time_to_degrees(time without time zone)', 'FUNCTION', NULL),
   ('msar', 'msar.top_level_domains', 'TABLE', NULL),

db/sql/05_msar.sql

@@ -2111,16 +2111,18 @@ Args:
   tab_id: OID of the table whose columns we'll drop.
   col_ids: The attnums of the columns to drop.
 */
-DECLARE col_names text[];
 BEGIN
-  EXECUTE format(
-    'ALTER TABLE %I.%I %s',
-    msar.get_relation_schema_name(tab_id),
-    msar.get_relation_name(tab_id),
-    string_agg(format('DROP COLUMN %I', attname), ', ')
-  )
-  FROM pg_catalog.pg_attribute AS pga INNER JOIN unnest(col_ids) AS x(col) ON pga.attnum=x.col
-  WHERE attrelid=tab_id AND NOT attisdropped;
+  col_ids := array_remove(col_ids, null);
+  IF array_length(col_ids, 1) IS NOT NULL THEN
+    EXECUTE format(
+      'ALTER TABLE %I.%I %s',
+      msar.get_relation_schema_name(tab_id),
+      msar.get_relation_name(tab_id),
+      string_agg(format('DROP COLUMN %I', attname), ', ')
+    )
+    FROM pg_catalog.pg_attribute AS pga INNER JOIN unnest(col_ids) AS x(col) ON pga.attnum=x.col
+    WHERE attrelid=tab_id AND NOT attisdropped;
+  END IF;
 END;
 $$ LANGUAGE plpgsql RETURNS NULL ON NULL INPUT;
 
@@ -2576,8 +2578,11 @@ $$ LANGUAGE SQL;
 
 CREATE OR REPLACE FUNCTION
 msar.add_pkey_column(
-  tab_id regclass, pkey_type msar.pkey_kind, col_name text DEFAULT 'id'
-) RETURNS void AS $$/*
+  tab_id regclass,
+  pkey_type msar.pkey_kind,
+  drop_old_pkey_col boolean DEFAULT false,
+  col_name text DEFAULT 'id'
+) RETURNS integer AS $$/*
 Add a primary key column with a predefined default to a table.
 
 Any name collisions for the column are resolved automatically, with this column's name deferring to
@@ -2588,22 +2593,77 @@ Args:
   pkey_type: The "type" of the pkey column. 'UUIDv4' means a `uuid` column using the
              `gen_random_uuid()` function for its default values. 'IDENTITY' means an `integer`
              using `GENERATED BY DEFAULT AS IDENTITY` for default values.
+  drop_old_pkey_col: Whether we should drop the current primary key column during this operation.
   col_name: This optional value will set the name of the column.
 */
-DECLARE
-  pkey_text text;
 BEGIN
-  pkey_text := CASE pkey_type
-    WHEN 'IDENTITY' THEN 'integer PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY'
-    WHEN 'UUIDv4' THEN 'uuid PRIMARY KEY DEFAULT gen_random_uuid()'
-  END;
+  IF drop_old_pkey_col THEN
+    PERFORM msar.drop_columns(tab_id, msar.get_pk_column(tab_id));
+  END IF;
+  PERFORM msar.drop_constraint(tab_id, oid)
+    FROM pg_constraint WHERE conrelid=tab_id AND contype='p';
   EXECUTE format(
     'ALTER TABLE %I.%I ADD COLUMN %I %s;',
     msar.get_relation_schema_name(tab_id),
     msar.get_relation_name(tab_id),
     msar.build_unique_column_name(tab_id, col_name),
-    pkey_text
+    CASE pkey_type
+      WHEN 'IDENTITY' THEN 'integer PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY'
+      WHEN 'UUIDv4' THEN 'uuid PRIMARY KEY DEFAULT gen_random_uuid()'
+    END
+  );
+  RETURN msar.get_pk_column(tab_id);
+END;
+$$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION
+msar.set_pkey_column(
+  tab_id regclass,
+  col_id integer,
+  default_type msar.pkey_kind DEFAULT null,
+  drop_old_pkey_col boolean DEFAULT false
+) RETURNS void AS $$/*
+Set a primary key column with an optional predefined default on a table.
+
+Args:
+  tab_id: This is the OID or name of the table for which we'll set the primary key column.
+  col_id: This is the attnum of the column we'll set as the primary key column.
+  pkey_type: The "type" of the pkey column. 'UUIDv4' means a `uuid` column using the
+             `gen_random_uuid()` function for its default values. 'IDENTITY' means an `integer`
+             using `GENERATED BY DEFAULT AS IDENTITY` for default values. If `null` is passed, we
+             do not set any default, or change the column type.
+  drop_old_pkey_col: Whether we should drop the current primary key column during this operation.
+*/
+DECLARE
+  pkey_col_max_val integer;
+BEGIN
+  IF drop_old_pkey_col THEN
+    PERFORM msar.drop_columns(tab_id, msar.get_pk_column(tab_id));
+  END IF;
+  PERFORM msar.drop_constraint(tab_id, oid)
+    FROM pg_constraint WHERE conrelid=tab_id AND contype='p';
+  PERFORM msar.add_constraints(
+    tab_id,
+    jsonb_build_array(jsonb_build_object('type', 'p', 'columns', jsonb_build_array(col_id)))
   );
+  IF default_type IS NOT NULL THEN
+    EXECUTE format(
+      CASE default_type WHEN 'IDENTITY' THEN
+        $s$
+        ALTER TABLE %1$I.%2$I
+          ALTER COLUMN %3$I TYPE integer USING msar.cast_to_integer(%3$I),
+          ALTER COLUMN %3$I ADD GENERATED BY DEFAULT AS IDENTITY;
+        SELECT setval(pg_get_serial_sequence('%1$I.%2$I', '%3$s'), max(%3$I)) FROM %1$I.%2$I;
+        $s$
+      WHEN 'UUIDv4' THEN
+        'ALTER TABLE %1$I.%2$I ALTER COLUMN %3$I SET DEFAULT gen_random_uuid();'
+      END,
+      msar.get_relation_schema_name(tab_id),
+      msar.get_relation_name(tab_id),
+      msar.get_column_name(tab_id, col_id)
+    );
+  END IF;
 END;
 $$ LANGUAGE plpgsql;
 
@@ -3232,7 +3292,8 @@ BEGIN
   RETURN jsonb_build_object(
     'oid', created_table_id::bigint,
     'name', relname,
-    'renamed_columns', renamed_columns::jsonb
+    'renamed_columns', renamed_columns::jsonb,
+    'pkey_column_attnum', msar.get_pk_column(created_table_id)
   ) FROM pg_catalog.pg_class WHERE oid = created_table_id;
 END;
 $$ LANGUAGE plpgsql;
@@ -3292,7 +3353,8 @@ BEGIN
     'copy_sql', copy_sql,
     'table_oid', rel_id::bigint,
     'table_name', relname,
-    'renamed_columns', (mathesar_table ->> 'renamed_columns')::jsonb
+    'renamed_columns', mathesar_table -> 'renamed_columns',
+    'pkey_column_attnum', mathesar_table -> 'pkey_column_attnum'
   ) FROM pg_catalog.pg_class WHERE oid = rel_id;
 END;
 $$ LANGUAGE plpgsql;

db/sql/test_sql_functions.sql

@@ -192,10 +192,13 @@ $$ LANGUAGE plpgsql;
 CREATE OR REPLACE FUNCTION test_add_pkey_column_uuid() RETURNS SETOF TEXT AS $f$
 BEGIN
   PERFORM __setup_add_pkey_col();
-  PERFORM msar.add_pkey_column(
-    tab_id => 'add_pkey_col_testable'::regclass,
-    pkey_type => 'UUIDv4',
-    col_name => 'User Id'
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'UUIDv4',
+      drop_old_pkey_col => true,
+      col_name => 'User Id'
+    ), 3, 'Should return correct attnum'
   );
   RETURN NEXT col_is_pk('add_pkey_col_testable', 'User Id');
   RETURN NEXT col_type_is('add_pkey_col_testable', 'User Id', 'uuid');
@@ -206,31 +209,150 @@ $f$ LANGUAGE plpgsql;
 CREATE OR REPLACE FUNCTION test_add_pkey_column_identity() RETURNS SETOF TEXT AS $f$
 BEGIN
   PERFORM __setup_add_pkey_col();
-  PERFORM msar.add_pkey_column(
-    tab_id => 'add_pkey_col_testable'::regclass,
-    pkey_type => 'IDENTITY',
-    col_name => 'Identity'
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY',
+      drop_old_pkey_col => true,
+      col_name => 'Identity'
+    ), 3, 'Should return correct attnum'
   );
   RETURN NEXT col_is_pk('add_pkey_col_testable', 'Identity');
   RETURN NEXT col_type_is('add_pkey_col_testable', 'Identity', 'integer');
 END;
 $f$ LANGUAGE plpgsql;
 
 
+CREATE OR REPLACE FUNCTION test_add_pkey_column_collision() RETURNS SETOF TEXT AS $f$
+BEGIN
+  PERFORM __setup_add_pkey_col();
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY',
+      drop_old_pkey_col => true,
+      col_name => 'col1'
+    ), 3, 'Should return correct attnum'
+  );
+  RETURN NEXT col_is_pk('add_pkey_col_testable', 'col1 1', 'rename when collision');
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY',
+      drop_old_pkey_col => true,
+      col_name => 'col1'
+    ), 4, 'Should return correct attnum'
+  );
+  RETURN NEXT col_is_pk(
+    'add_pkey_col_testable', 'col1 1', 'do not rename when collision with previous dropped pkey'
+  );
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY',
+      drop_old_pkey_col => false,
+      col_name => 'col1'
+    ), 5, 'Should return correct attnum'
+  );
+  RETURN NEXT col_is_pk(
+    'add_pkey_col_testable', 'col1 2', 'rename when collision with previous undropped pkey'
+  );
+  RETURN NEXT columns_are('add_pkey_col_testable', ARRAY['col1', 'col2', 'col1 1', 'col1 2']);
+END;
+$f$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION test_add_pkey_column_defaults_collide() RETURNS SETOF TEXT AS $f$
+BEGIN
+  PERFORM __setup_add_pkey_col();
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY'
+    ), 3, 'Should return correct attnum'
+  );
+  RETURN NEXT col_is_pk('add_pkey_col_testable', 'id');
+  RETURN NEXT is(
+    msar.add_pkey_column(
+      tab_id => 'add_pkey_col_testable'::regclass,
+      pkey_type => 'IDENTITY'
+    ), 4, 'Should return correct attnum'
+  );
+  RETURN NEXT col_is_pk('add_pkey_col_testable', 'id 1');
+END;
+$f$ LANGUAGE plpgsql;
+
+
 CREATE OR REPLACE FUNCTION test_add_pkey_column_malformed() RETURNS SETOF TEXT AS $f$
 BEGIN
   PERFORM __setup_add_pkey_col();
   RETURN NEXT throws_like(
     $s$SELECT msar.add_pkey_column(
         tab_id => 'add_pkey_col_testable'::regclass,
-        pkey_type => 'ident'
+        pkey_type => 'ident',
+        drop_old_pkey_col => false
     );$s$,
     'invalid input value for enum%'
   );
 END;
 $f$ LANGUAGE plpgsql;
 
 
+-- msar.set_pkey_column ----------------------------------------------------------------------------
+
+CREATE OR REPLACE FUNCTION __setup_set_pkey_col() RETURNS SETOF TEXT AS $$
+BEGIN
+  CREATE TABLE set_pkey_col_testable (
+    id integer PRIMARY KEY GENERATED BY DEFAULT AS IDENTITY,
+    col1 numeric,
+    col2 varchar,
+    "Column 3" uuid
+  );
+  INSERT INTO set_pkey_col_testable (col1, col2, "Column 3") VALUES
+    (324, 'abc', 'f2539a15-1644-413e-91e1-36008c6700e5'),
+    (567, 'def', 'e4c7d5d9-4eb3-4274-91f8-6af32c013f3a');
+END;
+$$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION test_set_pkey_column_numeric() RETURNS SETOF TEXT AS $f$
+BEGIN
+  PERFORM __setup_set_pkey_col();
+  PERFORM msar.set_pkey_column(
+    tab_id => 'set_pkey_col_testable'::regclass,
+    col_id => 2,
+    default_type => 'IDENTITY',
+    drop_old_pkey_col => true
+  );
+  RETURN NEXT col_is_pk('set_pkey_col_testable', 'col1');
+  RETURN NEXT col_type_is('set_pkey_col_testable', 'col1', 'integer');
+  RETURN NEXT columns_are('set_pkey_col_testable', ARRAY['col1', 'col2', 'Column 3']);
+  INSERT INTO set_pkey_col_testable(col2) VALUES ('ghi');
+  RETURN NEXT results_eq('SELECT max(col1) FROM set_pkey_col_testable', 'VALUES (568)');
+END;
+$f$ LANGUAGE plpgsql;
+
+
+CREATE OR REPLACE FUNCTION test_set_pkey_column_uuid() RETURNS SETOF TEXT AS $f$
+BEGIN
+  PERFORM __setup_set_pkey_col();
+  PERFORM msar.set_pkey_column(
+    tab_id => 'set_pkey_col_testable'::regclass,
+    col_id => 4,
+    default_type => 'UUIDv4',
+    drop_old_pkey_col => false
+  );
+  RETURN NEXT col_is_pk('set_pkey_col_testable', 'Column 3');
+  RETURN NEXT columns_are('set_pkey_col_testable', ARRAY['id', 'col1', 'col2', 'Column 3']);
+  INSERT INTO set_pkey_col_testable(col2) VALUES ('ghi');
+  RETURN NEXT results_eq(
+    'SELECT COUNT(DISTINCT "Column 3") FROM set_pkey_col_testable',
+    'VALUES (3::bigint)'
+  );
+END;
+$f$ LANGUAGE plpgsql;
+
+
 -- msar.add_columns --------------------------------------------------------------------------------
 
 CREATE OR REPLACE FUNCTION __setup_add_columns() RETURNS SETOF TEXT AS $$
@@ -1705,6 +1827,7 @@ BEGIN
     response ->> 'copy_sql', 'COPY tab_create_schema.anewtable () FROM STDIN'
   );
   RETURN NEXT is(response -> 'renamed_columns', '{}'::jsonb);
+  RETURN NEXT is((response ->> 'pkey_column_attnum')::integer, 1);
 END;
 $f$ LANGUAGE plpgsql;
 
@@ -1730,6 +1853,7 @@ BEGIN
     response ->> 'copy_sql', 'COPY tab_create_schema.anewtable () FROM STDIN'
   );
   RETURN NEXT is(response -> 'renamed_columns', '{}'::jsonb);
+  RETURN NEXT is((response ->> 'pkey_column_attnum')::integer, 1);
 END;
 $f$ LANGUAGE plpgsql;
 
@@ -1762,6 +1886,7 @@ BEGIN
     'COPY tab_create_schema.anewtable ("My Col", col2) FROM STDIN'
   );
   RETURN NEXT is(response -> 'renamed_columns', '{}'::jsonb);
+  RETURN NEXT is((response ->> 'pkey_column_attnum')::integer, 1);
 END;
 $f$ LANGUAGE plpgsql;