feat: process table columns with GENERATED ALWAYS for partition_data_id() and partition_data_time()

sql/functions/partition_data_id.sql

@@ -36,6 +36,8 @@ v_source_tablename          text;
 v_sql                       text;
 v_start_control             bigint;
 v_total_rows                bigint := 0;
+v_has_identity_column       boolean;
+v_override_clause           text := '';
 
 BEGIN
     /*
@@ -193,6 +195,11 @@ END IF;
 
 v_current_partition_name := @extschema@.check_name_length(COALESCE(v_parent_tablename), v_min_partition_id::text, TRUE);
 
+v_has_identity_column := table_has_identity_columns(v_parent_schema || '.' || v_parent_tablename, p_ignored_columns);
+IF v_has_identity_column THEN
+    v_override_clause := 'OVERRIDING SYSTEM VALUE';
+END IF;
+
 IF v_default_exists THEN
 
     -- Child tables cannot be created if data that belongs to it exists in the default
@@ -201,23 +208,25 @@ IF v_default_exists THEN
     -- Temp table created above to avoid excessive temp creation in loop
     EXECUTE format('WITH partition_data AS (
             DELETE FROM %1$I.%2$I WHERE %3$I >= %4$s AND %3$I < %5$s RETURNING *)
-        INSERT INTO partman_temp_data_storage (%6$s) SELECT %6$s FROM partition_data'
+        INSERT INTO partman_temp_data_storage (%6$s) %7$s SELECT %6$s FROM partition_data'
         , v_source_schemaname
         , v_source_tablename
         , v_control
         , v_min_partition_id
         , v_max_partition_id
-        , v_column_list);
+        , v_column_list
+        , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
 
     -- Set analyze to true if a table is created
     v_analyze := @extschema@.create_partition_id(p_parent_table, v_partition_id);
 
     EXECUTE format('WITH partition_data AS (
             DELETE FROM partman_temp_data_storage RETURNING *)
-        INSERT INTO %1$I.%2$I (%3$s) SELECT %3$s FROM partition_data'
+        INSERT INTO %1$I.%2$I (%3$s) %4$s SELECT %3$s FROM partition_data'
         , v_parent_schema
         , v_current_partition_name
-        , v_column_list);
+        , v_column_list
+        , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
 
 ELSE
 
@@ -226,15 +235,16 @@ ELSE
 
     EXECUTE format('WITH partition_data AS (
             DELETE FROM ONLY %1$I.%2$I WHERE %3$I >= %4$s AND %3$I < %5$s RETURNING *)
-        INSERT INTO %6$I.%7$I (%8$s) SELECT %8$s FROM partition_data'
+        INSERT INTO %6$I.%7$I (%8$s) %9$s SELECT %8$s FROM partition_data'
         , v_source_schemaname
         , v_source_tablename
         , v_control
         , v_min_partition_id
         , v_max_partition_id
         , v_parent_schema
         , v_current_partition_name
-        , v_column_list);
+        , v_column_list
+        , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
 
 END IF;
 

sql/functions/partition_data_time.sql

@@ -39,6 +39,8 @@ v_source_tablename          text;
 v_rowcount                  bigint;
 v_start_control             timestamptz;
 v_total_rows                bigint := 0;
+v_has_identity_column       boolean;
+v_override_clause           text := '';
 
 BEGIN
 /*
@@ -237,30 +239,37 @@ FOR i IN 1..p_batch_count LOOP
     v_partition_suffix := to_char(v_min_partition_timestamp, v_datetime_string);
     v_current_partition_name := @extschema@.check_name_length(v_parent_tablename, v_partition_suffix, TRUE);
 
+    v_has_identity_column := table_has_identity_columns(v_parent_schema || '.' || v_parent_tablename, p_ignored_columns);
+    IF v_has_identity_column THEN
+        v_override_clause := 'OVERRIDING SYSTEM VALUE';
+    END IF;
+
     IF v_default_exists THEN
         -- Child tables cannot be created if data that belongs to it exists in the default
         -- Have to move data out to temporary location, create child table, then move it back
 
         -- Temp table created above to avoid excessive temp creation in loop
         EXECUTE format('WITH partition_data AS (
                 DELETE FROM %1$I.%2$I WHERE %3$s >= %4$L AND %3$s < %5$L RETURNING *)
-            INSERT INTO partman_temp_data_storage (%6$s) SELECT %6$s FROM partition_data'
+            INSERT INTO partman_temp_data_storage (%6$s) %7$s SELECT %6$s FROM partition_data'
             , v_source_schemaname
             , v_source_tablename
             , v_partition_expression
             , v_min_partition_timestamp
             , v_max_partition_timestamp
-            , v_column_list);
+            , v_column_list
+            , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
 
         -- Set analyze to true if a table is created
         v_analyze := @extschema@.create_partition_time(p_parent_table, v_partition_timestamp);
 
         EXECUTE format('WITH partition_data AS (
                 DELETE FROM partman_temp_data_storage RETURNING *)
-            INSERT INTO %I.%I (%3$s) SELECT %3$s FROM partition_data'
+            INSERT INTO %I.%I (%3$s) %4$s SELECT %3$s FROM partition_data'
             , v_parent_schema
             , v_current_partition_name
-            , v_column_list);
+            , v_column_list
+            , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
 
     ELSE
 
@@ -269,15 +278,16 @@ FOR i IN 1..p_batch_count LOOP
 
         EXECUTE format('WITH partition_data AS (
                             DELETE FROM ONLY %I.%I WHERE %s >= %L AND %3$s < %5$L RETURNING *)
-                         INSERT INTO %6$I.%7$I (%8$s) SELECT %8$s FROM partition_data'
+                         INSERT INTO %6$I.%7$I (%8$s) %9$s SELECT %8$s FROM partition_data'
                             , v_source_schemaname
                             , v_source_tablename
                             , v_partition_expression
                             , v_min_partition_timestamp
                             , v_max_partition_timestamp
                             , v_parent_schema
                             , v_current_partition_name
-                            , v_column_list);
+                            , v_column_list
+                            , v_override_clause);  -- insert "OVERRIDING SYSTEM VALUE" or blank
     END IF;
 
     GET DIAGNOSTICS v_rowcount = ROW_COUNT;

sql/functions/table_has_identity_columns.sql

@@ -0,0 +1,40 @@
+CREATE OR REPLACE FUNCTION table_has_identity_columns(
+    table_name text,
+    p_ignored_columns text[] DEFAULT NULL
+) RETURNS boolean AS $$
+DECLARE
+    identity_columns_count int;
+    schema_name text;
+    table_only_name text;
+BEGIN
+    -- if table_name contains schema，split schema and table
+    IF table_name LIKE '%.%' THEN
+        schema_name := split_part(table_name, '.', 1);
+        table_only_name := split_part(table_name, '.', 2);
+    ELSE
+        -- if not，use curent schema
+        schema_name := current_schema();
+        table_only_name := table_name;
+    END IF;
+
+    -- query all columns defined with GENERATED ALWAYS AS IDENTITY ，except columns in p_ignored_columns
+    SELECT COUNT(*)
+    INTO identity_columns_count
+    FROM pg_attribute a
+    JOIN pg_class c ON c.oid = a.attrelid
+    JOIN pg_namespace n ON n.oid = c.relnamespace
+    WHERE c.relname = table_only_name
+    AND n.nspname = schema_name
+    AND a.attidentity = 'a'     -- GENERATED ALWAYS AS IDENTITY
+    AND a.attnum > 0            -- except system column
+    AND NOT a.attisdropped      -- except deleted column
+    AND (p_ignored_columns IS NULL OR a.attname != ANY(p_ignored_columns)); 
+
+    IF identity_columns_count > 0 THEN
+        RETURN true;
+    ELSE
+        RETURN false;
+    END IF;
+END;
+$$ LANGUAGE plpgsql;
+