|
| 1 | +#!/bin/bash |
| 2 | +set -o errexit |
| 3 | +set -o pipefail |
| 4 | +set -o nounset |
| 5 | + |
| 6 | +# Allow overriding via environment for local testing |
| 7 | +readonly WB_EXE="${WB_EXE:-/usr/bin/wb}" |
| 8 | +readonly PGWEB_BASE="${PGWEB_BASE:-/root/.pgweb}" |
| 9 | +readonly BOOKMARK_DIR="${PGWEB_BASE}/bookmarks" |
| 10 | + |
| 11 | +# Create base directory if it doesn't exist |
| 12 | +mkdir -p "${PGWEB_BASE}" |
| 13 | + |
| 14 | +# Helper function to get credentials and generate IAM auth token |
| 15 | +generate_iam_token() { |
| 16 | + local resource_id="${1}" |
| 17 | + local scope="${2}" |
| 18 | + local endpoint="${3}" |
| 19 | + local port="${4}" |
| 20 | + local username="${5}" |
| 21 | + local region="${6}" |
| 22 | + |
| 23 | + # Get credentials from Workbench |
| 24 | + local wb_creds |
| 25 | + wb_creds=$(${WB_EXE} resource credentials --id "${resource_id}" --scope "${scope}" --format json 2>/dev/null) || return 1 |
| 26 | + readonly wb_creds |
| 27 | + |
| 28 | + # Extract AWS credentials |
| 29 | + local access_key secret_key session_token |
| 30 | + access_key=$(echo "${wb_creds}" | jq -r '.AccessKeyId') |
| 31 | + secret_key=$(echo "${wb_creds}" | jq -r '.SecretAccessKey') |
| 32 | + session_token=$(echo "${wb_creds}" | jq -r '.SessionToken') |
| 33 | + readonly access_key secret_key session_token |
| 34 | + |
| 35 | + # Generate IAM token |
| 36 | + AWS_ACCESS_KEY_ID="${access_key}" \ |
| 37 | + AWS_SECRET_ACCESS_KEY="${secret_key}" \ |
| 38 | + AWS_SESSION_TOKEN="${session_token}" \ |
| 39 | + aws rds generate-db-auth-token \ |
| 40 | + --hostname "${endpoint}" \ |
| 41 | + --port "${port}" \ |
| 42 | + --username "${username}" \ |
| 43 | + --region "${region}" |
| 44 | +} |
| 45 | + |
| 46 | +# Helper function to create bookmark TOML file |
| 47 | +create_bookmark() { |
| 48 | + local output_file="${1}" |
| 49 | + local endpoint="${2}" |
| 50 | + local port="${3}" |
| 51 | + local username="${4}" |
| 52 | + local password="${5}" |
| 53 | + local database="${6}" |
| 54 | + |
| 55 | + cat > "${output_file}" <<EOF |
| 56 | +host = "${endpoint}" |
| 57 | +port = ${port} |
| 58 | +user = "${username}" |
| 59 | +password = "${password}" |
| 60 | +database = "${database}" |
| 61 | +sslmode = "require" |
| 62 | +EOF |
| 63 | +} |
| 64 | + |
| 65 | +refresh_bookmarks() { |
| 66 | + echo "$(date): Refreshing pgweb bookmarks from Workbench resources..." |
| 67 | + |
| 68 | + # Create temporary directory for new bookmarks (using PID for uniqueness) |
| 69 | + local TEMP_DIR="${PGWEB_BASE}/bookmarks.tmp.$$" |
| 70 | + readonly TEMP_DIR |
| 71 | + rm -rf "${TEMP_DIR}" |
| 72 | + mkdir -p "${TEMP_DIR}" |
| 73 | + |
| 74 | + # Get list of Aurora databases from Workbench |
| 75 | + local RESOURCES |
| 76 | + RESOURCES=$(${WB_EXE} resource list --format json) |
| 77 | + readonly RESOURCES |
| 78 | + |
| 79 | + # Process each resource |
| 80 | + echo "${RESOURCES}" | jq -c '.[]' | while read -r resource; do |
| 81 | + local RESOURCE_TYPE |
| 82 | + RESOURCE_TYPE=$(echo "${resource}" | jq -r '.resourceType') |
| 83 | + |
| 84 | + # Skip non-Aurora resources |
| 85 | + if [[ ! "${RESOURCE_TYPE}" =~ AURORA_DATABASE ]]; then |
| 86 | + continue |
| 87 | + fi |
| 88 | + |
| 89 | + local RESOURCE_ID |
| 90 | + RESOURCE_ID=$(echo "${resource}" | jq -r '.id') |
| 91 | + echo " Processing: ${RESOURCE_ID} (type: ${RESOURCE_TYPE})" |
| 92 | + |
| 93 | + # Extract database details from top level (controlled) or referencedResource (reference) |
| 94 | + local DB_DATA |
| 95 | + if [[ "${RESOURCE_TYPE}" == "AWS_AURORA_DATABASE" ]]; then |
| 96 | + DB_DATA="${resource}" |
| 97 | + else |
| 98 | + DB_DATA=$(echo "${resource}" | jq -r '.referencedResource') |
| 99 | + fi |
| 100 | + |
| 101 | + # Extract database connection info |
| 102 | + local DB_NAME RO_ENDPOINT RO_USER RW_ENDPOINT RW_USER PORT REGION |
| 103 | + DB_NAME=$(echo "${DB_DATA}" | jq -r '.databaseName') |
| 104 | + RO_ENDPOINT=$(echo "${DB_DATA}" | jq -r '.roEndpoint') |
| 105 | + RO_USER=$(echo "${DB_DATA}" | jq -r '.roUser') |
| 106 | + RW_ENDPOINT=$(echo "${DB_DATA}" | jq -r '.rwEndpoint') |
| 107 | + RW_USER=$(echo "${DB_DATA}" | jq -r '.rwUser') |
| 108 | + PORT=$(echo "${DB_DATA}" | jq -r '.port') |
| 109 | + REGION=$(echo "${DB_DATA}" | jq -r '.region // "us-east-1"') |
| 110 | + |
| 111 | + # Validate all required fields are present |
| 112 | + if [[ -z "${DB_NAME}" || "${DB_NAME}" == "null" ]] || \ |
| 113 | + [[ -z "${RO_ENDPOINT}" || "${RO_ENDPOINT}" == "null" ]] || \ |
| 114 | + [[ -z "${RO_USER}" || "${RO_USER}" == "null" ]] || \ |
| 115 | + [[ -z "${RW_ENDPOINT}" || "${RW_ENDPOINT}" == "null" ]] || \ |
| 116 | + [[ -z "${RW_USER}" || "${RW_USER}" == "null" ]] || \ |
| 117 | + [[ -z "${PORT}" || "${PORT}" == "null" ]]; then |
| 118 | + echo " Missing required database fields, skipping" |
| 119 | + continue |
| 120 | + fi |
| 121 | + |
| 122 | + # Try to create READ_ONLY bookmark |
| 123 | + echo " Checking read access..." |
| 124 | + local RO_TOKEN |
| 125 | + if RO_TOKEN=$(generate_iam_token "${RESOURCE_ID}" "READ_ONLY" "${RO_ENDPOINT}" "${PORT}" "${RO_USER}" "${REGION}"); then |
| 126 | + echo " Read access confirmed" |
| 127 | + echo " Creating read-only bookmark..." |
| 128 | + create_bookmark "${TEMP_DIR}/${RESOURCE_ID} (Read-Only).toml" "${RO_ENDPOINT}" "${PORT}" "${RO_USER}" "${RO_TOKEN}" "${DB_NAME}" |
| 129 | + echo " Created bookmark: ${RESOURCE_ID} (Read-Only)" |
| 130 | + else |
| 131 | + echo " No read access to ${RESOURCE_ID}, skipping" |
| 132 | + continue |
| 133 | + fi |
| 134 | + |
| 135 | + # Try to create WRITE_READ bookmark |
| 136 | + echo " Checking write access..." |
| 137 | + local RW_TOKEN |
| 138 | + if RW_TOKEN=$(generate_iam_token "${RESOURCE_ID}" "WRITE_READ" "${RW_ENDPOINT}" "${PORT}" "${RW_USER}" "${REGION}"); then |
| 139 | + echo " Write access confirmed" |
| 140 | + echo " Creating write-read bookmark..." |
| 141 | + create_bookmark "${TEMP_DIR}/${RESOURCE_ID} (Write-Read).toml" "${RW_ENDPOINT}" "${PORT}" "${RW_USER}" "${RW_TOKEN}" "${DB_NAME}" |
| 142 | + echo " Created bookmark: ${RESOURCE_ID} (Write-Read)" |
| 143 | + else |
| 144 | + echo " No write access, skipping write-read bookmark" |
| 145 | + fi |
| 146 | + done |
| 147 | + |
| 148 | + # Count bookmarks - must use find since the while loop runs in a subshell (due to pipe), |
| 149 | + # so a counter variable incremented in the loop would not be visible here |
| 150 | + local BOOKMARK_COUNT |
| 151 | + BOOKMARK_COUNT=$(find "${TEMP_DIR}" -name "*.toml" -type f 2>/dev/null | wc -l) |
| 152 | + readonly BOOKMARK_COUNT |
| 153 | + echo "$(date): Refresh complete. Created ${BOOKMARK_COUNT} bookmarks." |
| 154 | + |
| 155 | + # Atomically update symlink to point to new bookmark directory |
| 156 | + ln -sfn "$(basename "${TEMP_DIR}")" "${BOOKMARK_DIR}" |
| 157 | + |
| 158 | + # Cleanup old bookmark directories (all except current) |
| 159 | + find "${PGWEB_BASE}" -maxdepth 1 -type d -name "bookmarks.tmp.*" ! -name "bookmarks.tmp.$$" -exec rm -rf {} \; |
| 160 | +} |
| 161 | + |
| 162 | +# Run single refresh |
| 163 | +if ! refresh_bookmarks; then |
| 164 | + echo "$(date): ERROR: Bookmark refresh failed" |
| 165 | + exit 1 |
| 166 | +fi |
0 commit comments