transfer-report

#!/bin/bash

set -o errexit -o nounset -o pipefail

export PGHOST
export PGUSER

readonly ExecAbsPath=$(readlink --canonicalize "$0")
readonly ExecName=$(basename "$ExecAbsPath")


main()
{
  local opts
  if ! opts=$(getopt --name "$ExecName" --longoptions host:,user: --options H:U: -- "$@")
  then
    printf 'failed to parse command line\n' >&2
    return 1
  fi

  eval set -- "$opts"

  while true
  do
    case "$1" in
      -H|--host)
        PGHOST="$2"
        shift 2
        ;;
      -U|--user)
        PGUSER="$2"
        shift 2
        ;;
      --)
        shift
        break
        ;;
      *)
        printf 'Unknown option %s\n' "$1" >&2
        return 1
        ;;
    esac
  done

  make_report
}


make_report()
{
  local msgs=tr-msg-entries
  ensure "$msgs" filter_msgs

  local openMsgs=tr-opens.csv
  ensure "$openMsgs" extract_open_msgs_csv < "$msgs"

  local addModMsgs=tr-add-mod.json
  ensure "$addModMsgs" extract_addmod_msgs < "$msgs"

  local addModMsgDir=tr-add-mod
  if ! [[ -e "$addModMsgDir" ]]
  then
    track_call bin_addmod_msgs "$addModMsgDir".tmp < "$addModMsgs"
    mv "$addModMsgDir".tmp "$addModMsgDir"
  fi

  local sizeMap=tr-size-map.csv
  ensure "$sizeMap" csv_entity_size_map_binned "$(date '+%Y-%m-%d %H:%M:%S')" "$addModMsgDir"

  local uploads=tr-uploads.csv
  ensure "$uploads" mk_uploads < "$addModMsgs"

  local downloadsQuery=tr-downloads-query.sql
  ensure "$downloadsQuery" mk_downloads_query "$sizeMap" "$openMsgs"

  local downloads=tr-downloads.csv
  ensure "$downloads" mk_downloads < "$downloadsQuery"

  local report=tr-report.csv
  ensure "$report" combine_reports "$uploads" "$downloads"

  cat "$report"
}


ensure()
{
  local product="$1"
  shift

  if ! [[ -e "$product" ]]
  then
    track_call "$@" > "$product".tmp
    mv "$product".tmp "$product"
  fi
}


track_call()
{
  local func="$1"

  disp_begin_func "$func"
  "$@"
  disp_end_func "$func"
}


extract_addmod_msgs()
{
  filter_addmod_msgs | compact_addmod_msgs
}


extract_open_msgs_csv()
{
  filter_open_msgs | csv_open_msgs
}


filter_msgs()
{
  awk --file - <(cat) \
<<'EOF'
  $4 == "NOTICE:" && $5 == "execCmd:cmd/amqptopicsend.py" {
    msg = gensub($3 " " $4 " ", "", 1, $0);

    # Remove potential trailing garabage
    msg = gensub(/[^\r]\\n.*/, "", 1, msg);

    # switch to unicode escapes
    print gensub(/\\\\x/, "\\\\\\\\u00", "g", msg);
  }
EOF
}


filter_addmod_msgs()
{
  awk --file - <(cat) \
<<'EOF'
  /"data-object\.(add|mod)"/ {
    timestamp = $1 "." $2;

    # Remove through message type, two blanks, and the leading quote
    match($0, /"data-object\.(add|mod)"/);
    offset = RSTART + RLENGTH + 3;
    msg = gensub(/\\("|')/, "\\1", "g", substr($0, offset));

    msg = substr(msg, 1, length(msg) - 1);  # remove trailing quote
    msg = gensub(/\r/, "\\r", "g", msg);    # escape carriage returns

    # Add timestamp to start of message and record separator before
    print "\x1e{\"timestamp\":\"" timestamp "\"," substr(msg, 2);
  }
EOF
}


filter_open_msgs()
{
  awk --file - <(cat) \
<<'EOF'
  /"data-object\.open"/ {
    # Remove through "data-object.open"m two blanks, and the leading quote
    match($0, /"data-object\.open"/);
    offset = RSTART + RLENGTH + 3;
    msg = gensub(/\\("|')/, "\\1", "g", substr($0, offset));

    # remove trailing quote
    msg = substr(msg, 1, length(msg) - 1);

    # escape carriage returns and prefix with record sep
    printf "\x1e%s\n", gensub(/\r/, "\\r", "g", msg);
  }
EOF
}


bin_addmod_msgs()
{
  local binDir="$1"

  rm --force "$binDir"/*
  mkdir --parents "$binDir"

  local entry
  while IFS= read -r entry
  do
    printf '%s\n' "$entry" >> "$addModMsgDir".tmp/"${entry:12:1}".json
  done
}


csv_entity_size_map_binned()
{
  local endDate="$1"
  local binDir="$2"

  local part
  for part in $(ls "$binDir"/*)
  do
    csv_entity_size_map "$endDate" < "$part"
  done
}


csv_open_msgs()
{
  jq \
    --raw-output \
    --from-file <(cat \
<<'JQ'
  if (.timestamp | not) or (.entity | not) then
    empty
  else
    [ .entity, (.timestamp | split(".") | join(" ")) ]
  end |
  @csv
JQ
      )
}


csv_entity_size_map()
{
  local lastStopTs="$1"

  jq \
    --raw-output --slurp \
    --arg LAST_STOP_TS "$lastStopTs" \
    --from-file <(cat \
<<'JQ'
  # removes all of the leading objects that have a size of s from an array
  def ltrim_size(s):
    def trim_rec: if .[0].size != s then . else (.[1:] | trim_rec) end;
    if (length == 0) then [] else trim_rec end;

  # removes all objects that have the same size as the previous object in
  # an array
  def del_seq_sizes:
    def emit_keepers:
      if (length == 0) then
        empty
      else
        (first as $f |
         ($f,
          (.[1:] | ltrim_size($f.size) | emit_keepers)))
      end;
    [ emit_keepers ];

  # Transforms size entry to to intervaled size entry
  def to_interval(stop_ts):
    { obj: .entity, start_ts: .timestamp, stop_ts: stop_ts, size: .size };

  # Converts a size map to an intervaled size map
  def intervals:
    def emit_intervals(first_stop_ts):
      if (length == 0) then
        empty
      else
        ((first | to_interval(first_stop_ts)) as $f |
         ($f,
          (.[1:] | emit_intervals($f.start_ts))))
      end;
    reverse | [ emit_intervals($LAST_STOP_TS) ] | reverse;

  group_by(.entity) |
  map(sort_by(.timestamp) | del_seq_sizes | intervals) |
  flatten |
  .[] |
  [ .obj, .start_ts, .stop_ts, .size ] |
  @csv
JQ
      )
}


compact_addmod_msgs()
{
  jq \
    --compact-output --seq \
    --from-file <(cat \
<<'JQ'
  if (.timestamp | not) or (.entity | not) or (.size | not) then
    empty
  else
    { entity: .entity, timestamp: (.timestamp | split(".") | join(" ")),  size: .size }
  end
JQ
      )
}


mk_downloads_query()
{
  local sizeMapFile="$1"
  local downloadsFile="$2"

  cat \
<<'SQL'
  \timing on
  BEGIN;

  \echo 'Importing size mappings'
  CREATE TEMPORARY TABLE size_map(
      data_uuid CHAR(37),
      start_time TIMESTAMP,
      stop_time TIMESTAMP,
      size BIGINT)
    ON COMMIT DROP;

  COPY size_map FROM STDIN WITH (FORMAT CSV);
SQL

  cat "$sizeMapFile"
  echo '\.'

  cat \
<<'SQL'
  CREATE INDEX idx_size_map_data_uuid_time ON size_map(data_uuid, start_time, stop_time);

  \echo 'Importing download events'
  CREATE TEMPORARY TABLE downloads(data_uuid CHAR(37), time TIMESTAMP) ON COMMIT DROP;

  COPY downloads FROM STDIN WITH (FORMAT CSV);
SQL

  cat "$downloadsFile"
  echo '\.'

  cat \
<<'SQL'
  CREATE INDEX idx_downloads_data_all ON downloads(data_uuid, time);

  \echo 'Resolving download sizes using size mappings'
  CREATE TEMPORARY TABLE sized_downloads(data_uuid, count, size) ON COMMIT DROP AS
  SELECT d.data_uuid, COUNT(*), s.size
  FROM downloads AS d
    LEFT JOIN size_map AS s
      ON s.data_uuid = d.data_uuid AND s.start_time <= d.time AND d.time < s.stop_time
  GROUP BY d.data_uuid, s.size;

  CREATE INDEX idx_sized_downloads_uuid ON sized_downloads(data_uuid);
  CREATE INDEX idx_sized_downloads_all ON sized_downloads(data_uuid, count, size);

  \echo 'Resolving remaining download sizes using ICAT'
  CREATE TEMPORARY TABLE resolved_downloads(data_uuid, count, size) ON COMMIT DROP AS
  SELECT d.data_uuid, d.count, COALESCE(d.size, AVG(dm.data_size), 0)
  FROM sized_downloads AS d
    LEFT JOIN r_meta_main AS mm
      ON mm.meta_attr_name = 'ipc_UUID' AND mm.meta_attr_value = d.data_uuid
    LEFT JOIN r_objt_metamap AS om ON om.meta_id = mm.meta_id
    LEFT JOIN r_data_main AS dm ON dm.data_id = om.object_id
  GROUP BY d.data_uuid, d.count, d.size;

  \echo 'Exporting summary'
  COPY (SELECT SUM(count), CAST(SUM(count * size) AS BIGINT) FROM resolved_downloads)
  TO STDOUT
  WITH (FORMAT CSV);

  ROLLBACK;
SQL
}


mk_uploads()
{
  printf 'Upload Count,Upload Volume(B)\n'
  jq --raw-output '.size' | awk '{ c++; v+=$1 } END { printf "%d,%d\n", c, v }'
}


mk_downloads()
{
  printf 'Download Count,Download Volume(B)\n'
  psql ICAT | tee >(tail --lines 4 | head --lines 1) >(disp_dbms_prog) > /dev/null
}


combine_reports()
{
  local uploads="$1"
  local downloads="$2"

  paste --delimiters=, "$uploads" "$downloads"
}


disp_begin_func()
{
  local func="$1"

  printf '\e[32m%s: \e[1mbegin\e[0m' "$func" >&3
}


disp_end_func()
{
  local func="$1"

  printf '\033[1K\r\e[32m%s: \e[1mdone\e[0m\n' "$func" >&3
}


disp_dbms_prog()
{
  printf '\n' >&3

  while IFS= read -r
  do
    printf '  \e[32m%s\e[0m\n' "$REPLY" >&3
  done
}


disp_error()
{
  while IFS= read -r
  do
    if [[ "$(get_cursor_col)" -gt 0 ]]
    then
      printf '\n'
    fi

    printf '\e[31m%s\e[0m\n' "$REPLY"
  done >&2
}


get_cursor_col()
{
  # based on a script from http://invisible-island.net/xterm/xterm.faq.html
  exec < /dev/tty

  local oldStty
  oldStty=$(stty -g)
  stty raw -echo min 0

  local col
  echo -en "\033[6n" > /dev/tty
  IFS=';' read -r -d R _ col

  stty "$oldStty"

  printf '%s' "$((col - 1))"
}


(main "$@" 3>&4 2>&1 1>&4 | disp_error) 4>&1