#!/bin/bash
# -*- sh-indentation: 2; sh-basic-offset: 2-*-

# This is based on codes from Michael Paquier and Tomas Vondra:
# * https://www.postgresql.org/message-id/flat/Zbr6piWuVHDtFFOl%40paquier.xyz#dbbec4d5c54ef2317be01a54abaf495c
# * https://www.postgresql.org/message-id/flat/26541788-8853-4d93-86cd-5f701b13ae51%40enterprisedb.com

DIR=${1:-$(pwd)}

psql postgres > /dev/null 2>&1 <<EOF
DROP DATABASE IF EXISTS test;
CREATE DATABASE test;
EOF
psql test > /dev/null 2>&1 <<EOF
CREATE EXTENSION blackhole_am;
CREATE OR REPLACE FUNCTION create_table_cols(tabname text, num_cols int)
RETURNS VOID AS
\$func\$
DECLARE
  query text;
BEGIN
  query := 'CREATE UNLOGGED TABLE ' || tabname || ' (';
  FOR i IN 1..num_cols LOOP
    query := query || 'a_' || i::text || ' int2 default 1';
    IF i != num_cols THEN
      query := query || ', ';
    END IF;
  END LOOP;
  query := query || ')';
  EXECUTE format(query);
END
\$func\$ LANGUAGE plpgsql;
EOF

for format in text csv binary; do
  for n_columns in 1 10 100; do
    for n_rows in 10 100 1000 10000 10000 100000 1000000; do
      psql test > /dev/null 2>&1 <<EOF
DROP TABLE IF EXISTS to_table;
DROP TABLE IF EXISTS from_table;

SELECT create_table_cols ('to_table', ${n_columns});
SELECT create_table_cols ('from_table', ${n_columns});

INSERT INTO to_table SELECT 1 FROM generate_series(1, ${n_rows});

COPY to_table TO '$DIR/test.data' WITH (FORMAT ${format});

ALTER TABLE from_table SET ACCESS METHOD blackhole_am;
EOF

      # run COPY TO 5x
      for r in $(seq 1 5); do
        s=$(psql test -t -A -c "SELECT EXTRACT(EPOCH FROM now())")
        psql test -c "COPY to_table TO '/dev/null' WITH (FORMAT ${format})" > /dev/null 2>&1
        d=$(psql test -t -A -c "SELECT 1000 * (EXTRACT(EPOCH FROM now()) - $s)")
        echo "TO,${format},${n_columns},${n_rows},${r},${d}"
      done

      # run COPY FROM 5x
      for r in $(seq 1 5); do
        s=$(psql test -t -A -c "SELECT EXTRACT(EPOCH FROM now())")
        psql test -c "COPY from_table FROM '$DIR/test.data' WITH (FORMAT ${format})" > /dev/null 2>&1
        d=$(psql test -t -A -c "SELECT 1000 * (EXTRACT(EPOCH FROM now()) - $s)")
        echo "FROM,${format},${n_columns},${n_rows},${r},${d}"
      done
    done
  done
done