Add VCP attribute support and add msg about extension

nancy_run.sh

@@ -111,6 +111,7 @@ function dbg_cli_parameters() {
 --delta-sql_do: ${DELTA_SQL_DO}
 --delta-sql_undo: ${DELTA_SQL_UNDO}
 --delta-config: ${DELTA_CONFIG}
+--custom-script: ${CUSTOM_SCRIPT}
 
 --aws-ec2-type: ${AWS_EC2_TYPE}
 --aws-keypair-name: $AWS_KEYPAIR_NAME
@@ -121,6 +122,7 @@ function dbg_cli_parameters() {
 --aws-block-duration: ${AWS_BLOCK_DURATION}
 --aws-zfs: ${AWS_ZFS}
 --s3-cfg-path: ${S3_CFG_PATH}
+--aws-vpc-id: ${AWS_VPC_ID}
 
 --no-perf: ${NO_PERF}
 
@@ -229,6 +231,7 @@ function check_cli_parameters() {
   ([[ ! -z ${SQL_AFTER_DB_RESTORE+x} ]] && [[ -z $SQL_AFTER_DB_RESTORE ]]) && unset -v SQL_AFTER_DB_RESTORE
   ([[ ! -z ${AWS_ZONE+x} ]] && [[ -z $AWS_ZONE ]]) && unset -v AWS_ZONE
   ([[ ! -z ${CONFIG+x} ]] && [[ -z $CONFIG ]]) && unset -v CONFIG
+  ([[ ! -z ${CUSTOM_SCRIPT+x} ]] && [[ -z $CUSTOM_SCRIPT ]]) && unset -v CUSTOM_SCRIPT
   ### CLI parameters checks ###
   if [[ "${RUN_ON}" == "aws" ]]; then
     if [ ! -z ${CONTAINER_ID+x} ]; then
@@ -260,6 +263,9 @@ function check_cli_parameters() {
     if [[ -z ${AWS_ZONE+x} ]]; then
       err "NOTICE: AWS EC2 zone is not specified. Will be determined during the price optimization process."
     fi
+    if [[ -z ${AWS_VPC_ID+x} ]]; then
+      err "NOTICE: AWS VPC is not specified. Default will be used."
+    fi
     if [[ -z ${AWS_ZFS+x} ]]; then
       err "NOTICE: Ext4 will be used for PGDATA."
     else
@@ -577,6 +583,15 @@ function check_cli_parameters() {
     fi
   fi
 
+  if [[ ! -z ${CUSTOM_SCRIPT+x} ]]; then
+    check_path CUSTOM_SCRIPT
+    if [[ "$?" -ne "0" ]]; then
+      dbg "WARNING: Value given as after_db_init_code: '$CUSTOM_SCRIPT' not found as file will use as content"
+      echo "$CUSTOM_SCRIPT" > $CUSTOM_SCRIPT
+      msg "Custom script $CUSTOM_SCRIPT executed"
+    fi
+  fi
+
   if [[ ! -z ${SQL_AFTER_DB_RESTORE+x} ]]; then
     check_path SQL_AFTER_DB_RESTORE
     if [[ "$?" -ne "0" ]]; then
@@ -615,11 +630,12 @@ function check_cli_parameters() {
 #   (text) [7] The AWS region to launch the instance
 #              (for example us-east-1, eu-central-1)
 #   (text) [8] The AWS zone to launch the instance in (one of a,b,c,d,e)
+#   (text) [9] The AWS VPC to launch the instance in (one of a,b,c,d,e)
 # Returns:
 #   None
 #######################################
 function create_ec2_docker_machine() {
-  msg "Attempting to provision a Docker machine in region $7 with price $3..."
+  msg "Attempting to provision a Docker machine in region $7 with price $3 and vpc $9..."
   docker-machine create --driver=amazonec2 \
     --amazonec2-request-spot-instance \
     --amazonec2-instance-type=$2 \
@@ -629,6 +645,8 @@ function create_ec2_docker_machine() {
     --amazonec2-ssh-keypath="$6" \
     --amazonec2-region="$7" \
     --amazonec2-zone="$8" \
+    --amazonec2-vpc-id="$9" \
+    --amazonec2-use-private-address="true" \
     $1 2> >(grep -v "failed waiting for successful resource state" >&2) &
 }
 
@@ -1146,6 +1164,8 @@ while [ $# -gt 0 ]; do
       ARTIFACTS_DESTINATION="$2"; shift 2 ;;
     --artifacts-dirname )
       ARTIFACTS_DIRNAME="$2"; shift 2 ;;
+    --custom-script )
+      CUSTOM_SCRIPT="$2"; shift 2 ;;
 
     --aws-ec2-type )
       AWS_EC2_TYPE="$2"; shift 2 ;;
@@ -1163,6 +1183,8 @@ while [ $# -gt 0 ]; do
         AWS_BLOCK_DURATION=$2; shift 2 ;;
     --aws-zfs )
         AWS_ZFS=1; shift ;;
+    --aws-vpc-id )
+        AWS_VPC_ID="$2"; shift 2 ;;
     --db-ebs-volume-id )
       DB_EBS_VOLUME_ID=$2; shift 2;;
     --db-local-pgdata )
@@ -1212,7 +1234,7 @@ if [ -n "$CIRCLE_JOB" ]; then
   IS_CIRCLE_CI=true
 else
   IS_CIRCLE_CI=false
-fi	
+fi
 
 if $DEBUG ; then
   set -xueo pipefail
@@ -1248,7 +1270,7 @@ if [[ "$RUN_ON" == "localhost" ]]; then
       )
     else
       CONTAINER_HASH=$(docker run --cap-add SYS_ADMIN --name="pg_nancy_${CURRENT_TS}" \
-        ${DB_EXPOSE_PORT} \
+        ${DB_EXPOSE_PORT} \$AWS_BLOCK_DURATION $AWS_KEYPAIR_NAME $AWS_SSH_KEY_PATH $AWS_REGION $AWS_ZONE $AWS_VPC_ID
         -v $TMP_PATH:/machine_home \
         -dit "postgresmen/postgres-nancy:${PG_VERSION}" \
       )
@@ -1264,7 +1286,7 @@ if [[ "$RUN_ON" == "localhost" ]]; then
 elif [[ "$RUN_ON" == "aws" ]]; then
   determine_history_ec2_spot_price
   create_ec2_docker_machine $DOCKER_MACHINE $AWS_EC2_TYPE $EC2_PRICE \
-    $AWS_BLOCK_DURATION $AWS_KEYPAIR_NAME $AWS_SSH_KEY_PATH $AWS_REGION $AWS_ZONE
+    $AWS_BLOCK_DURATION $AWS_KEYPAIR_NAME $AWS_SSH_KEY_PATH $AWS_REGION $AWS_ZONE $AWS_VPC_ID
   status=$(wait_ec2_docker_machine_ready "$DOCKER_MACHINE" true)
   if [[ "$status" == "price-too-low" ]]; then
     msg "Price $price is too low for $AWS_EC2_TYPE instance. Getting the up-to-date value from the error message..."
@@ -1278,7 +1300,7 @@ elif [[ "$RUN_ON" == "aws" ]]; then
     DOCKER_MACHINE="${DOCKER_MACHINE//_/-}"
     #try start docker machine name with new price
     create_ec2_docker_machine $DOCKER_MACHINE $AWS_EC2_TYPE $EC2_PRICE \
-      $AWS_BLOCK_DURATION $AWS_KEYPAIR_NAME $AWS_SSH_KEY_PATH $AWS_REGION $AWS_ZONE
+      $AWS_BLOCK_DURATION $AWS_KEYPAIR_NAME $AWS_SSH_KEY_PATH $AWS_REGION $AWS_ZONE $AWS_VPC_ID
     wait_ec2_docker_machine_ready "$DOCKER_MACHINE" false
   fi
 
@@ -1647,6 +1669,30 @@ function apply_commands_after_container_init() {
   fi
 }
 
+#######################################
+# Execute custom shell commands in container after it started
+# Globals:
+#   CUSTOM_SCRIPT
+# Arguments:
+#   None
+# Returns:
+#   None
+#######################################
+function apply_custom_script() {
+  OP_START_TIME=$(date +%s)
+  if ([ ! -z ${CUSTOM_SCRIPT+x} ] && [ "$CUSTOM_SCRIPT" != "" ])
+  then
+    msg "Apply custom script after docker init"
+    CUSTOM_SCRIPT_FILENAME=$(basename $CUSTOM_SCRIPT)
+    copy_file $CUSTOM_SCRIPT
+    docker_exec bash -c "chmod +x ${MACHINE_HOME}/${CUSTOM_SCRIPT_FILENAME}"
+    output=$(docker_exec sh $MACHINE_HOME/$CUSTOM_SCRIPT_FILENAME)
+    END_TIME=$(date +%s)
+    DURATION=$(echo $((END_TIME-OP_START_TIME)) | awk '{printf "%d:%02d:%02d", $1/3600, ($1/60)%60, $1%60}')
+    msg "Time taken to apply \"custom script code\": $DURATION."
+  fi
+}
+
 #######################################
 # Execute SQL code before database restore
 # Globals:
@@ -1697,6 +1743,9 @@ function restore_dump() {
   ;;
       pgdmp)
   docker_exec bash -c "pg_restore -j $CPU_CNT --no-owner --no-privileges -U postgres -d $DB_NAME $MACHINE_HOME/$DB_DUMP_FILENAME" || true
+  ;;
+      *)
+  msg "The file should be with extension *.sql, *bz2, *gz or *.pgdmp: $DB_DUMP_FILENAME."
   ;;
     esac
   fi
@@ -2394,6 +2443,7 @@ docker_exec bash -c "psql -U postgres $DB_NAME -b -c 'create extension if not ex
 docker_exec bash -c "psql -U postgres $DB_NAME -b -c 'create extension if not exists pg_stat_kcache;' $VERBOSE_OUTPUT_REDIRECT"
 
 apply_commands_after_container_init
+apply_custom_script
 pg_config_init
 apply_sql_before_db_restore
 if [[ ! -z ${DB_DUMP+x} ]] || [[ ! -z ${DB_PGBENCH+x} ]]; then