#!/usr/bin/bash

warn=50
crit=100
cache_max_age=3600
DEBUG=""

declare -A SKIP
ONLY=""
EXCLUDED_SCHEMAS="'pg_catalog','information_schema','bucardo'"

while [ "$1" ]; do
  case "$1" in
    -w)
        warn=$2
    	shift
    	;;
    -c)
        crit=$2
        shift
        ;;
    -d|-D)
        DBNAME=$2
        shift
        ;;
    -u|-U)
        DBUSER="-U $2"
        shift
        ;;
    -h|-H)
    	HOST="$2"
    	shift
    	;;
    -i)
	if [ "$ONLY" ]; then
		ONLY="$ONLY,'$2'"
	else
		ONLY="'$2'"
	fi
	shift
	;;
    -s)
    	SKIP[$2]=1
    	shift
    	;;
    -S)
    	EXCLUDED_SCHEMAS="$EXCLUDED_SCHEMAS,'$2'"
    	shift
    	;;
    --cache-max-age)
        cache_max_age=$2
        shift
        ;;
    --use-cache)
        USE_CACHE=$2
    	shift
    	;;
    --make-cache)
    	MAKE_CACHE=$2
    	shift
    	;;
    --debug)
    	DEBUG=1
    	;;
  esac
  shift
done

if [ -z "$HOST" -o -z "$DBNAME" ]; then
  echo "Not ehough parameters!"
  echo "Usage: $0 -H primary_hostname -d db_name \\"
  echo "          [-u username] [-w warn] [-c critical] [--debug] \\"
  echo "          [-i only_table] [-s skipped_table] [-S excluded_schema]"
  echo "          [--make-cache=filename]"
  echo "          [--cache-max-age=sec] [--use-cache=filename]"
  exit 1
fi

SERVER1="$DBUSER $DBNAME"
SERVER2="-h $HOST $DBUSER $DBNAME"

debug() {
  if [ ! -z "$DEBUG" ]; then
    echo "$@" 1>&2
  fi
}
select1() {
  debug "1: $@"
  psql -At $SERVER1 -c "$@" 
}
select2() {
  debug "2: $@"
  psql -At $SERVER2 -c "$@"
}

sum=0
state=""
script_start=`date +%s.%N`
if [ "$MAKE_CACHE" ]; then
  echo "#script_start $script_start" > "$MAKE_CACHE.tmp"
fi
if [ "$USE_CACHE" ]; then
  current_ts=`date +%s`
  cache_ts=`stat -L --format %Y "$USE_CACHE"`
  if [ 0$cache_ts -lt $(( $current_ts - $cache_max_age )) ]; then
    echo "CRITICAL Cache file too old."
    exit 2
  fi
  TABLES=`grep -v '^#' "$USE_CACHE" | sort -k 2 -n -r | cut -d" " -f1`
else
  if [ "$ONLY" ]; then
    ONLY="AND (table_schema || '.' || table_name) IN ($ONLY)"
  fi
  TABLES=`select1 \
    "SELECT table_schema || '.' || table_name FROM information_schema.tables
     WHERE table_catalog='$DBNAME' AND table_type='BASE TABLE' AND
       table_schema NOT IN ($EXCLUDED_SCHEMAS) $ONLY;" \
    | sed 's/^public\.//'`
fi
for table in $TABLES; do
  if [ "${SKIP[$table]}" ]; then
    continue
  fi
  if [ "$USE_CACHE" ]; then
    # use cache
    diff=`grep "^$table " "$USE_CACHE" | cut -d" " -f2`
    if [ "$diff" = "" ]; then
      diff="10000"
      state="$state $table=MISSING_DATA"
    elif [ ! "$diff" = 0 ]; then
      state="$state $table=$diff"
    fi
  else
    start=`date +%s.%N`
    pkey=`select1 \
     "SELECT column_name FROM information_schema.table_constraints c
        LEFT JOIN information_schema.key_column_usage k
          ON c.constraint_name=k.constraint_name
             AND c.table_catalog=k.table_catalog
             AND c.table_name=k.table_name
      WHERE constraint_type='PRIMARY KEY' AND c.table_catalog='$DBNAME'
            AND c.table_name='$table';"`
    # multicolumn key?
    pkey="`echo "$pkey" | xargs echo | sed 's/^.* .*$//'`"
    if [ "$pkey" ]; then
      row1=`select1 "SELECT max($pkey) FROM $table;"`
      if [ -z "$row1" ]; then
        row1=0
      fi
      row2=`select2 "SELECT max($pkey) FROM $table;"`
      if [ -z "$row2" ]; then
        row2=0
      fi
    else
      row1=`select1 "SELECT count(*) FROM $table;"`
      row2=`select2 "SELECT count(*) FROM $table;"`
    fi
    end=`date +%s.%N`
    len=`echo $end $start - p | dc`
    diff=`echo $row2 $row1 - p | dc 2>/dev/null`
    debug "$len $table $row1 $row2 $diff"
    if [ "$diff" = "" ]; then
      if [ "$row1" = "$row2" ]; then
        diff="0"
      else
        state="$state $table=X"
        diff="100"
      fi
    elif [ ! "$diff" = 0 ]; then
      state="$state $table=$diff"
    fi
    if [ "$MAKE_CACHE" ]; then
      echo "$table $diff $row1 $row2 $len" >> "$MAKE_CACHE.tmp"
    fi
  fi
  if [[ "$diff" =~ ^-?[0-9.]+$ ]]; then
    absdiff="$diff"
    if [ "$diff" -lt 0 ]; then
      absdiff=$((-diff))
    fi
    sum=$((sum+absdiff))
  fi
done
script_end=`date +%s.%N`
script_time=`echo $script_end $script_start - p | dc`

if [ "$MAKE_CACHE" ]; then
  echo "#script_end $script_end $script_time" >> "$MAKE_CACHE.tmp"
  mv -f "$MAKE_CACHE.tmp" "$MAKE_CACHE"
fi

if [ "$MAKE_CACHE" ]; then
  # No critical or warn defined. Cache is updated, exit now.
  exit 0
fi

perfdata="diffsum=$sum time=$script_time"
if [ $sum -lt $warn ]; then
  echo "OK $perfdata$state|$perfdata"
  exit 0
elif [ $sum -lt $crit ]; then
  echo "WARNING $perfdata$state|$perfdata"
  exit 1
else
  echo "CRITICAL $perfdata$state|$perfdata"
  exit 2
fi
