Automating Local-Read Indexes When a New Region Is Added

				
					CREATE OR REPLACE PROCEDURE yb_tip.queue_missing_dimension_local_indexes()
LANGUAGE plpgsql
AS $$
DECLARE
  v_cfg                   record;
  v_target                record;

  v_region_token          text;
  v_tablespace_name       text;
  v_index_name            text;

  v_placement_blocks      jsonb;
  v_replica_regions       jsonb;
  v_replica_placement     text;

  v_tablespace_exists     boolean;
  v_index_exists          boolean;
  v_table_exists          regclass;

  v_key_columns_sql       text;
  v_include_columns_sql   text;
  v_include_clause        text;

  v_index_verb            text;

  v_create_tablespace_ddl text;
  v_create_index_ddl      text;
BEGIN
  /*
   * Prevent overlapping queue-generation runs.
   */
  IF NOT pg_try_advisory_lock(hashtext('yb_tip.queue_missing_dimension_local_indexes')) THEN
    RAISE NOTICE 'Another queue-generation run is already active. Skipping.';
    RETURN;
  END IF;

  /*
   * Loop through each enabled dimension table configuration.
   */
  FOR v_cfg IN
    SELECT *
    FROM yb_tip.dimension_index_config
    WHERE enabled = true
    ORDER BY config_id
  LOOP
    /*
     * Make sure the configured table exists.
     */
    SELECT to_regclass(format('%I.%I', v_cfg.table_schema, v_cfg.table_name))
    INTO v_table_exists;

    IF v_table_exists IS NULL THEN
      RAISE NOTICE 'Configured table %.% does not exist. Skipping config_id %.',
        v_cfg.table_schema,
        v_cfg.table_name,
        v_cfg.config_id;

      CONTINUE;
    END IF;

    /*
     * Validate that all configured key and INCLUDE columns exist.
     */
    IF EXISTS (
      WITH configured_columns AS (
        SELECT unnest(v_cfg.hash_key_columns || v_cfg.include_columns) AS column_name
      )
      SELECT 1
      FROM configured_columns cc
      WHERE NOT EXISTS (
        SELECT 1
        FROM information_schema.columns ic
        WHERE ic.table_schema = v_cfg.table_schema
          AND ic.table_name = v_cfg.table_name
          AND ic.column_name = cc.column_name
      )
    ) THEN
      RAISE NOTICE 'Configured columns for %.% are invalid. Skipping config_id %.',
        v_cfg.table_schema,
        v_cfg.table_name,
        v_cfg.config_id;

      CONTINUE;
    END IF;

    /*
     * Build the index key column list.
     * Example: ARRAY['store_id'] -> store_id HASH
     */
    SELECT string_agg(format('%I HASH', col), ', ')
    INTO v_key_columns_sql
    FROM unnest(v_cfg.hash_key_columns) AS col;

    /*
     * Build the INCLUDE column list.
     */
    SELECT string_agg(format('%I', col), ', ')
    INTO v_include_columns_sql
    FROM unnest(v_cfg.include_columns) AS col;

    v_include_clause :=
      CASE
        WHEN v_include_columns_sql IS NULL OR length(v_include_columns_sql) = 0 THEN ''
        ELSE format(' INCLUDE (%s)', v_include_columns_sql)
      END;

    v_index_verb :=
      CASE
        WHEN v_cfg.unique_index THEN 'CREATE UNIQUE INDEX'
        ELSE 'CREATE INDEX'
      END;

    /*
     * Loop through each distinct cloud + region currently visible
     * in the cluster topology.
     */
    FOR v_target IN
      SELECT DISTINCT cloud, region
      FROM yb_servers()
      WHERE cloud IS NOT NULL
        AND region IS NOT NULL
      ORDER BY cloud, region
    LOOP
      /*
       * Normalize region names to match the naming convention from the previous tip.
       *
       * Examples:
       *   us-east-1 -> east_1
       *   us-east-2 -> east_2
       *   us-west-2 -> west_2
       *   us-west-1 -> west_1
       */
      v_region_token :=
        regexp_replace(lower(v_target.region), '^us-', '');

      v_region_token :=
        regexp_replace(v_region_token, '[^a-z0-9]+', '_', 'g');

      /*
       * Use one regional tablespace per region. Multiple configured
       * dimension tables can share the same regional tablespace.
       */
      v_tablespace_name := 'ts_dim_' || v_region_token;

      /*
       * Use the configured prefix and suffix to build the index name.
       * Example: store_dimension + west_1 + _idx
       */
      v_index_name :=
        v_cfg.index_name_prefix || '_' || v_region_token || v_cfg.index_name_suffix;

      /*
       * PostgreSQL-compatible identifiers are limited to 63 bytes.
       * If the generated index name is too long, shorten it and add a hash.
       */
      IF length(v_index_name) > 63 THEN
        v_index_name :=
          left(v_cfg.index_name_prefix || '_' || v_region_token, 50)
          || '_'
          || substr(md5(v_cfg.config_id::text || '.' || v_target.cloud || '.' || v_target.region), 1, 8);
      END IF;

      /*
       * Pick the target region first, then enough peer regions to satisfy
       * the configured replication factor.
       *
       * The target region receives leader_preference = 1 so reads from that
       * region can prefer the local index leader.
       */
      WITH discovered_regions AS (
        SELECT DISTINCT cloud, region
        FROM yb_servers()
        WHERE cloud IS NOT NULL
          AND region IS NOT NULL
      ),
      selected_regions AS (
        SELECT cloud, region
        FROM discovered_regions
        ORDER BY
          CASE
            WHEN cloud = v_target.cloud AND region = v_target.region THEN 0
            ELSE 1
          END,
          cloud,
          region
        LIMIT v_cfg.replication_factor
      )
      SELECT
        jsonb_agg(
          CASE
            WHEN cloud = v_target.cloud AND region = v_target.region THEN
              jsonb_build_object(
                'cloud', cloud,
                'region', region,
                'zone', '*',
                'min_num_replicas', 1,
                'leader_preference', 1
              )
            ELSE
              jsonb_build_object(
                'cloud', cloud,
                'region', region,
                'zone', '*',
                'min_num_replicas', 1
              )
          END
          ORDER BY
            CASE
              WHEN cloud = v_target.cloud AND region = v_target.region THEN 0
              ELSE 1
            END,
            cloud,
            region
        ),
        jsonb_agg(
          jsonb_build_object('cloud', cloud, 'region', region)
          ORDER BY
            CASE
              WHEN cloud = v_target.cloud AND region = v_target.region THEN 0
              ELSE 1
            END,
            cloud,
            region
        )
      INTO v_placement_blocks, v_replica_regions
      FROM selected_regions;

      IF jsonb_array_length(v_placement_blocks) < v_cfg.replication_factor THEN
        RAISE NOTICE 'Not enough regions to satisfy RF %. Skipping config_id %, region %.',
          v_cfg.replication_factor,
          v_cfg.config_id,
          v_target.region;

        CONTINUE;
      END IF;

      v_replica_placement :=
        jsonb_build_object(
          'num_replicas', v_cfg.replication_factor,
          'placement_blocks', v_placement_blocks
        )::text;

      /*
       * Check whether the regional tablespace already exists.
       */
      SELECT EXISTS (
        SELECT 1
        FROM pg_tablespace
        WHERE spcname = v_tablespace_name
      )
      INTO v_tablespace_exists;

      /*
       * Check whether the regional duplicate covering index already exists.
       */
      SELECT EXISTS (
        SELECT 1
        FROM pg_class c
        JOIN pg_namespace n
          ON n.oid = c.relnamespace
        WHERE n.nspname = v_cfg.table_schema
          AND c.relname = v_index_name
          AND c.relkind = 'i'
      )
      INTO v_index_exists;

      /*
       * Generate one-line DDL commands.
       * Keeping them one line makes them easier for external automation to consume.
       */
      v_create_tablespace_ddl := format(
        'CREATE TABLESPACE %I WITH (replica_placement = %L);',
        v_tablespace_name,
        v_replica_placement
      );

      v_create_index_ddl := format(
        '%s %I ON %I.%I (%s)%s TABLESPACE %I;',
        v_index_verb,
        v_index_name,
        v_cfg.table_schema,
        v_cfg.table_name,
        v_key_columns_sql,
        v_include_clause,
        v_tablespace_name
      );

      /*
       * Queue CREATE TABLESPACE if needed.
       *
       * The tablespace is per region, not per table. Therefore, only one
       * pending or errored CREATE TABLESPACE command should exist for a region.
       */
      IF NOT v_tablespace_exists THEN
        INSERT INTO yb_tip.dimension_index_ddl_queue (
          config_id,
          cloud,
          region,
          table_schema,
          table_name,
          tablespace_name,
          index_name,
          ddl_type,
          ddl_command
        )
        SELECT
          NULL,
          v_target.cloud,
          v_target.region,
          NULL,
          NULL,
          v_tablespace_name,
          NULL,
          'create_tablespace',
          v_create_tablespace_ddl
        WHERE NOT EXISTS (
          SELECT 1
          FROM yb_tip.dimension_index_ddl_queue
          WHERE cloud = v_target.cloud
            AND region = v_target.region
            AND ddl_type = 'create_tablespace'
            AND status IN ('pending', 'error')
        );
      END IF;

      /*
       * Queue CREATE INDEX if needed.
       */
      IF NOT v_index_exists THEN
        INSERT INTO yb_tip.dimension_index_ddl_queue (
          config_id,
          cloud,
          region,
          table_schema,
          table_name,
          tablespace_name,
          index_name,
          ddl_type,
          ddl_command
        )
        SELECT
          v_cfg.config_id,
          v_target.cloud,
          v_target.region,
          v_cfg.table_schema,
          v_cfg.table_name,
          v_tablespace_name,
          v_index_name,
          'create_index',
          v_create_index_ddl
        WHERE NOT EXISTS (
          SELECT 1
          FROM yb_tip.dimension_index_ddl_queue
          WHERE config_id = v_cfg.config_id
            AND cloud = v_target.cloud
            AND region = v_target.region
            AND ddl_type = 'create_index'
            AND status IN ('pending', 'error')
        );
      END IF;

      /*
       * Record the discovered table + region combination and expected object names.
       */
      INSERT INTO yb_tip.dimension_index_region_registry (
        config_id,
        cloud,
        region,
        table_schema,
        table_name,
        tablespace_name,
        index_name,
        replica_regions,
        created_at,
        last_seen_at
      )
      VALUES (
        v_cfg.config_id,
        v_target.cloud,
        v_target.region,
        v_cfg.table_schema,
        v_cfg.table_name,
        v_tablespace_name,
        v_index_name,
        v_replica_regions,
        now(),
        now()
      )
      ON CONFLICT (config_id, cloud, region)
      DO UPDATE SET
        table_schema = EXCLUDED.table_schema,
        table_name = EXCLUDED.table_name,
        tablespace_name = EXCLUDED.tablespace_name,
        index_name = EXCLUDED.index_name,
        replica_regions = EXCLUDED.replica_regions,
        last_seen_at = now();

    END LOOP;
  END LOOP;

  PERFORM pg_advisory_unlock(hashtext('yb_tip.queue_missing_dimension_local_indexes'));

EXCEPTION
  WHEN OTHERS THEN
    PERFORM pg_advisory_unlock(hashtext('yb_tip.queue_missing_dimension_local_indexes'));
    RAISE;
END;
$$;