Export YugabyteDB CREATE TABLE DDL Using Pure SQL (Including SPLIT AT and SPLIT INTO)

				
					CREATE OR REPLACE FUNCTION public.export_table(
    p_rel regclass,
    include_if_not_exists boolean DEFAULT false
)
RETURNS text
LANGUAGE plpgsql
SET client_min_messages TO 'warning'
AS $function$
DECLARE
  v_schema             text;
  v_table              text;
  v_relkind            "char";
  v_persistence        "char";
  v_reloptions         text[];

  v_cols               text;
  v_constraints        text;
  v_pk_clause          text;
  v_part_clause        text;
  v_tablespace         text;
  v_stmt               text;
  v_split_clause       text;
  v_with_items         text[] := ARRAY[]::text[];

  v_num_tablets        int;
  v_num_hash_cols      int := 0;
  v_is_colocated       boolean;
  v_colocation_id      oid;

  v_prefix             text := 'CREATE TABLE';

  v_conindid           oid;
  v_indexdef           text;
  v_allow_split_into   boolean := false;
BEGIN
  SELECT n.nspname,
         c.relname,
         c.relkind,
         c.relpersistence,
         c.reloptions
  INTO v_schema, v_table, v_relkind, v_persistence, v_reloptions
  FROM pg_class c
  JOIN pg_namespace n
    ON n.oid = c.relnamespace
  WHERE c.oid = p_rel;

  IF NOT FOUND THEN
    RAISE EXCEPTION 'relation % not found', p_rel;
  END IF;

  IF v_persistence = 't' THEN
    v_prefix := 'CREATE TEMP TABLE';
  ELSIF v_persistence = 'u' THEN
    v_prefix := 'CREATE UNLOGGED TABLE';
  END IF;

  IF include_if_not_exists THEN
    v_prefix := v_prefix || ' IF NOT EXISTS';
  END IF;

  BEGIN
    SELECT num_tablets,
           num_hash_key_columns,
           is_colocated,
           colocation_id
    INTO v_num_tablets,
         v_num_hash_cols,
         v_is_colocated,
         v_colocation_id
    FROM yb_table_properties(p_rel);
  EXCEPTION
    WHEN OTHERS THEN
      v_num_tablets   := NULL;
      v_num_hash_cols := 0;
      v_is_colocated  := NULL;
      v_colocation_id := NULL;
  END;

  /*
   * Column definitions
   */
  SELECT string_agg(def, E',\n  ' ORDER BY attnum)
  INTO v_cols
  FROM (
    SELECT
      a.attnum,
      format(
        '%I %s%s%s%s',
        a.attname,
        format_type(a.atttypid, a.atttypmod),
        CASE
          WHEN ad.adbin IS NOT NULL
          THEN ' DEFAULT ' || pg_get_expr(ad.adbin, ad.adrelid)
          ELSE ''
        END,
        CASE WHEN a.attnotnull THEN ' NOT NULL' ELSE '' END,
        CASE
          WHEN a.attidentity = 'a' THEN ' GENERATED ALWAYS AS IDENTITY'
          WHEN a.attidentity = 'd' THEN ' GENERATED BY DEFAULT AS IDENTITY'
          ELSE ''
        END
      ) AS def
    FROM pg_attribute a
    LEFT JOIN pg_attrdef ad
      ON ad.adrelid = a.attrelid
     AND ad.adnum   = a.attnum
    WHERE a.attrelid = p_rel
      AND a.attnum > 0
      AND NOT a.attisdropped
  ) s;

  /*
   * Reconstruct PK from pg_get_indexdef()
   */
  v_pk_clause := NULL;

  SELECT conindid
  INTO v_conindid
  FROM pg_constraint
  WHERE conrelid = p_rel
    AND contype = 'p';

  IF v_conindid IS NOT NULL THEN
    SELECT pg_get_indexdef(v_conindid)
    INTO v_indexdef;

    v_pk_clause :=
      regexp_replace(
        v_indexdef,
        '^.*USING\s+\S+\s+(\(.*\))$',
        'PRIMARY KEY \1'
      );

    IF v_pk_clause = v_indexdef THEN
      v_pk_clause := NULL;
    END IF;
  END IF;

  /*
   * Non-PK constraints
   */
  SELECT string_agg(pg_get_constraintdef(oid, true), E',\n  ' ORDER BY contype, conname)
  INTO v_constraints
  FROM pg_constraint
  WHERE conrelid = p_rel
    AND contype IN ('u','f','c','x');

  IF v_pk_clause IS NOT NULL THEN
    v_cols := v_cols || E',\n  ' || v_pk_clause;
  END IF;

  IF v_constraints IS NOT NULL THEN
    v_cols := v_cols || E',\n  ' || v_constraints;
  END IF;

  /*
   * PARTITION BY
   */
  IF v_relkind = 'p' THEN
    SELECT 'PARTITION BY ' || pg_get_partkeydef(p_rel)
    INTO v_part_clause;
  END IF;

  /*
   * WITH (...)
   */
  IF v_reloptions IS NOT NULL THEN
    v_with_items := v_with_items || v_reloptions;
  END IF;

  IF v_is_colocated IS TRUE THEN
    v_with_items := array_append(v_with_items, 'COLOCATION = true');
  END IF;

  IF v_colocation_id IS NOT NULL THEN
    v_with_items := array_append(v_with_items, 'COLOCATION_ID = ' || v_colocation_id);
  END IF;

  /*
   * Base CREATE TABLE
   */
  v_stmt :=
      v_prefix || ' '
   || quote_ident(v_schema) || '.' || quote_ident(v_table)
   || ' (' || E'\n  '
   || v_cols
   || E'\n)';

  IF array_length(v_with_items, 1) IS NOT NULL THEN
    v_stmt := v_stmt || E'\nWITH (' || array_to_string(v_with_items, ', ') || ')';
  END IF;

  IF v_part_clause IS NOT NULL THEN
    v_stmt := v_stmt || E'\n' || v_part_clause;
  END IF;

  /*
   * TABLESPACE
   */
  SELECT ts.spcname
  INTO v_tablespace
  FROM pg_class c
  JOIN pg_tablespace ts
    ON ts.oid = c.reltablespace
  WHERE c.oid = p_rel
    AND c.reltablespace <> 0;

  IF v_tablespace IS NOT NULL THEN
    v_stmt := v_stmt || E'\nTABLESPACE ' || quote_ident(v_tablespace);
  END IF;

  /*
   * SPLIT clause
   *
   * Use the exact same call pattern that test_split_clause() proved works.
   * Then only allow SPLIT INTO fallback for clearly hash-style / non-range cases.
   */
  v_split_clause := NULL;

  BEGIN
    SELECT yb_get_range_split_clause(p_rel)
    INTO v_split_clause;
  EXCEPTION
    WHEN OTHERS THEN
      v_split_clause := NULL;
  END;

  IF v_split_clause IS NOT NULL THEN
    v_split_clause := btrim(v_split_clause);
    IF v_split_clause = '' THEN
      v_split_clause := NULL;
    END IF;
  END IF;

  /*
   * Conservative fallback:
   * - allow SPLIT INTO when there is no PK clause at all
   * - or when the PK clause explicitly contains HASH
   * - or when YB says there are hash key columns
   */
  v_allow_split_into :=
       v_pk_clause IS NULL
    OR position(' HASH' IN coalesce(v_pk_clause, '')) > 0
    OR COALESCE(v_num_hash_cols, 0) > 0;

  IF v_split_clause IS NOT NULL THEN
    v_stmt := v_stmt || E'\n' || v_split_clause;
  ELSIF v_allow_split_into
     AND COALESCE(v_num_tablets, 0) > 1
     AND COALESCE(v_is_colocated, false) = false
  THEN
    v_stmt := v_stmt || E'\nSPLIT INTO ' || v_num_tablets || ' TABLETS';
  END IF;

  RETURN v_stmt || ';';
END;
$function$;