#!/bin/bash

if (( $# != 3 )); then
  echo "Usage: $(basename $0) <username> <db-host> <db-name>"
  exit 1
fi

PSQL_USER=${1}
PSQL_HOST=${2}
PSQL_DB=${3}
PSQL="psql --no-align --tuples-only --log-file $(hostname).log --no-psqlrc -U ${PSQL_USER} -h ${PSQL_HOST} -d ${PSQL_DB}"

#Check the connection and permissions
pg_version=$(${PSQL} -c "select version();" 2>/dev/null)
if (( $? != 0 )); then
  echo "Error connecting to database: ${PSQL_USER}@${PSQL_HOST} (${PSQL_DB})"
  exit 1
fi
echo ${pg_version}
pg_class_permissions=$(${PSQL} -c "select has_table_privilege('pg_catalog.pg_class', 'update');")
if [[ $pg_class_permissions != t ]]; then
  echo "Error, user ${PSQL_USER} cannot update pg_class in order to live update primary keys. Exiting..."
  exit 1
fi

#Create procedures that generate the sql to rebuild indexes
${PSQL} 1>&2 <<\EOSQL
-- Following two functions are used to re-index a primary key.
-- Use as follows:
--   select rebuild_pkey_index('public','test');
-- Then execute the SQL it generates.
CREATE OR REPLACE FUNCTION swap_for_pkey(text,text,text) returns integer
AS
$$
   DECLARE
     cmd text;
     oid1 integer;
     oid2 integer;
     filenode1 integer;
     filenode2 integer;
     relation text;
   BEGIN
      select oid::integer into oid1 from pg_class where relname=$2 and relnamespace = (select oid from pg_namespace where nspname=$1);
     RAISE NOTICE 'PKEY OID: %',oid1;
      select relfilenode::integer into filenode1 from pg_class where oid=oid1;
      select oid::integer into oid2 from pg_class where relname=$3 and relnamespace = (select oid from pg_namespace where nspname=$1);
     RAISE NOTICE 'PKEY OID: %',oid2;
      select relfilenode::integer into filenode2 from pg_class where oid=oid2;
      select (indrelid::regclass)::text into relation from pg_index where indexrelid=oid1;
    RAISE NOTICE 'RELATION NAME: %',relation;
      cmd:='LOCK '||relation||';';
      RAISE NOTICE 'Executing :- %',cmd;
      Execute cmd;
      cmd:='UPDATE pg_class SET relfilenode='||filenode2|| ' WHERE oid='||oid1||';';
      RAISE NOTICE 'Executing :- %',cmd;
      Execute cmd;
      cmd:='UPDATE pg_class SET relfilenode='||filenode1|| ' WHERE oid='||oid2||';';
      RAISE NOTICE 'Executing :- %',cmd;
      Execute cmd;
      cmd:='DROP INDEX '||$1||'.'||$3||';';
      RAISE NOTICE 'Executing :- %',cmd;
      Execute cmd;
      return 0;
   END;
$$language plpgsql;

CREATE OR REPLACE Function rebuild_pkey_index(text,text) returns setof text
AS
$$
  DECLARE
     reloid integer;
     cmd text;
     rec record;
     oid1 integer;
     oid2 integer;
     filenode1 integer;
     filenode2 integer;
     relname1 text;
  BEGIN
    select oid into reloid from pg_class where relname=$2 and relnamespace=(select oid from pg_namespace where nspname=$1);
    for rec in SELECT c2.relname as indexname, substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true),0, position( 'INDEX'in pg_catalog.pg_get_indexdef(i.indexrelid, 0, true))+6) ||'CONCURRENTLY '|| c2.relname||'_new '|| substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true),position( 'ON 'in pg_catalog.pg_get_indexdef(i.indexrelid, 0, true)))||';' as command, substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true) from ' ON (.*) USING') as table FROM pg_catalog.pg_class c, pg_catalog.pg_class c2, pg_catalog.pg_index i WHERE c2.oid=reloid  AND c.oid = i.indrelid AND i.indexrelid = c2.oid and i.indisprimary=true ORDER BY i.indisprimary DESC, i.indisunique DESC, c2.relname
    LOOP
       return next rec.command;
       cmd:= 'ANALYZE VERBOSE '||rec.table||';';
       return next cmd;
       cmd:='SELECT swap_for_pkey('||''''||$1||''''||','||''''||rec.indexname||''''||','||''''||rec.indexname||'_new'||''''||');';
       return next cmd;
   END LOOP;
   END;
$$language plpgsql;


-- Following function will re-index a non-primary key.
-- Use as follows:
--   select rebuild_nonpkey_index('public','test');
-- Then execute the SQL it generates.
CREATE OR REPLACE Function rebuild_nonpkey_index(text,text) returns setof text
AS
$$
  DECLARE
     reloid integer;
     cmd text;
     rec record;
  BEGIN

    select oid into reloid from pg_class where relname=$2 and relnamespace=(select oid from pg_namespace where nspname=$1);
    for rec in SELECT c2.relname as indexname, substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true),0, position( 'INDEX'in pg_catalog.pg_get_indexdef(i.indexrelid, 0, true))+6) ||'CONCURRENTLY '|| c2.relname||'_new '|| substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true),position( 'ON 'in pg_catalog.pg_get_indexdef(i.indexrelid, 0, true)))||';' as command, substring(pg_catalog.pg_get_indexdef(i.indexrelid, 0, true) from ' ON (.*) USING') as table FROM pg_catalog.pg_class c, pg_catalog.pg_class c2, pg_catalog.pg_index i WHERE c2.oid=reloid  AND c.oid = i.indrelid AND i.indexrelid = c2.oid and i.indisprimary=false ORDER BY i.indisprimary DESC, i.indisunique DESC, c2.relname
    LOOP
       return next rec.command;
       cmd:= 'ANALYZE VERBOSE '||rec.table||';';
       return next cmd;
       cmd:= 'DROP INDEX '||rec.indexname||';';
       return next cmd;
       cmd:='ALTER INDEX '||rec.indexname||'_new'||' RENAME TO '||rec.indexname||';';
       return next cmd;
     END LOOP;
   END;
$$language plpgsql;

--Create stored procs to iterate over all indexes and run procedures above.
CREATE OR REPLACE Function generate_rebuild_pkey_index(text) returns setof text
AS
$$
  DECLARE
     cmd text;
     rec record;
  BEGIN

    for rec in SELECT s.indexrelname as index_name from pg_stat_user_indexes s join pg_index i on i.indexrelid=s.indexrelid where i.indisprimary=true and s.schemaname='public' order by s.relname, s.indexrelname
    LOOP
       cmd:= 'SELECT rebuild_pkey_index(''public'', ''' || rec.index_name || ''');';
       return next cmd;
     END LOOP;
   END;
$$language plpgsql;

CREATE OR REPLACE Function generate_rebuild_unique_nonpkey_index(text) returns setof text
AS
$$
  DECLARE
     cmd text;
     rec record;
  BEGIN

    for rec in SELECT s.indexrelname as index_name from pg_stat_user_indexes s join pg_index i on i.indexrelid=s.indexrelid where i.indisunique=true and i.indisprimary=false and s.schemaname='public' order by s.relname, s.indexrelname
    LOOP
       cmd:= 'SELECT rebuild_nonpkey_index(''public'', ''' || rec.index_name || ''');';
       return next cmd;
     END LOOP;
   END;
$$language plpgsql;

CREATE OR REPLACE Function generate_rebuild_nonunique_nonpkey_index(text) returns setof text
AS
$$
  DECLARE
     cmd text;
     rec record;
  BEGIN

    for rec in SELECT s.indexrelname as index_name from pg_stat_user_indexes s join pg_index i on i.indexrelid=s.indexrelid where i.indisunique=false and i.indisprimary=false and s.schemaname='public' order by s.relname, s.indexrelname
    LOOP
       cmd:= 'SELECT rebuild_nonpkey_index(''public'', ''' || rec.index_name || ''');';
       return next cmd;
     END LOOP;
   END;
$$language plpgsql;
EOSQL

if (( $? != 0 )); then
  echo "Error creating stored procedures. Exiting..."
  exit 1
fi

#Execute the stored procs that generate the sql to run
generate_reindex_sql=$(${PSQL} <<\EOSQL
select generate_rebuild_pkey_index('public');
-- Do not run rebuild on unique non-primary key indexes (some may be constraints which
-- this code does not cope with yet)
--select generate_rebuild_unique_nonpkey_index('public');
select generate_rebuild_nonunique_nonpkey_index('public');
EOSQL
)
reindex_sql=$(echo $generate_reindex_sql | ${PSQL})

if (( $? != 0 )); then
  echo "Error generating reindex sql. Exiting..."
  exit 1
fi

#Execute the sql that will reindex everything
echo $reindex_sql | ${PSQL}

if (( $? != 0 )); then
  echo "Error executing stored procedures. Exiting..."
  exit 1
fi

#Cleanup all the functions
${PSQL} 1>&2 <<\EOSQL
drop function swap_for_pkey(text,text,text);
drop function rebuild_pkey_index(text,text);
drop function rebuild_nonpkey_index(text,text);

drop function generate_rebuild_pkey_index(text);
drop function generate_rebuild_unique_nonpkey_index(text);
drop function generate_rebuild_nonunique_nonpkey_index(text);
EOSQL

if (( $? != 0 )); then
  echo "Error dropping stored procedures. Exiting..."
  exit 1
fi