Connect to Amazon Managed Service for Prometheus

# Example provisioning configuration
apiVersion: 1
datasources:
  - name: 'Amazon Managed Prometheus'
    type: 'grafana-amazonprometheus-datasource'
    url: 'https://aps-workspaces.us-west-2.amazonaws.com/workspaces/ws-12345678-1234-1234-1234-123456789012/'
    jsonData:
      httpMethod: 'POST'
      sigV4Auth: true
      sigV4AuthType: 'keys'
      sigV4Region: 'us-east-2'
    secureJsonData:
      sigV4AccessKey: '<access key>'
      sigV4SecretKey: '<secret key>'

#!/bin/bash

# Configuration
GRAFANA_URL=""
BEARER_TOKEN=""
LOG_FILE="grafana_migration_$(date +%Y%m%d_%H%M%S).log"

# Function to log messages to both console and file
log_message() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    echo "[$timestamp] $message" | tee -a "$LOG_FILE"
}

# Function to update a data source
update_data_source() {
    local uid="$1"
    local data="$2"

    response=$(curl -s -w "\n%{http_code}" -X PUT \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $BEARER_TOKEN" \
        -d "$data" \
        "$GRAFANA_URL/api/datasources/uid/$uid")

    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | sed '$d')

    if [[ "$http_code" -ge 200 && "$http_code" -lt 300 ]]; then
        log_message "$uid successful"
    else
        log_message "$uid error: HTTP $http_code - $response_body"
    fi
}

# Function to process and update data source types
update_data_source_type() {
    local result="$1"
    local processed_count=0
    local updated_count=0
    local readonly_count=0
    local skipped_count=0

    # Use jq to parse and process JSON
    echo "$result" | jq -c '.[]' | while read -r data; do
        uid=$(echo "$data" | jq -r '.uid')
        prometheus_type_migration=$(echo "$data" | jq -r '.jsonData["prometheus-type-migration"] // false')
        data_type=$(echo "$data" | jq -r '.type')
        read_only=$(echo "$data" | jq -r '.readOnly // false')

        processed_count=$((processed_count + 1))

        # Check conditions
        if [[ "$prometheus_type_migration" != "true" ]] || [[ "$data_type" != "grafana-amazonprometheus-datasource" ]]; then
            skipped_count=$((skipped_count + 1))
            continue
        fi

        if [[ "$read_only" == "true" ]]; then
            readonly_count=$((readonly_count + 1))
            log_message "$uid is readOnly. If this data source is provisioned, edit the data source type to be \`prometheus\` in the provisioning file."
            continue
        fi

        # Update the data
        updated_data=$(echo "$data" | jq '.type = "prometheus" | .jsonData["prometheus-type-migration"] = false')
        update_data_source "$uid" "$updated_data"
        updated_count=$((updated_count + 1))

        # Log the raw data for debugging (optional - uncomment if needed)
        # log_message "DEBUG - Updated data for $uid: $updated_data"
    done

    # Note: These counts won't work in the while loop due to subshell
    # Moving summary to the main function instead
}

# Function to get summary statistics
get_summary_stats() {
    local result="$1"
    local total_datasources=$(echo "$result" | jq '. | length')
    local migration_candidates=$(echo "$result" | jq '[.[] | select(.jsonData["prometheus-type-migration"] == true and .type == "grafana-amazonprometheus-datasource")] | length')
    local readonly_candidates=$(echo "$result" | jq '[.[] | select(.jsonData["prometheus-type-migration"] == true and .type == "grafana-amazonprometheus-datasource" and .readOnly == true)] | length')
    local updateable_candidates=$(echo "$result" | jq '[.[] | select(.jsonData["prometheus-type-migration"] == true and .type == "grafana-amazonprometheus-datasource" and (.readOnly == false or .readOnly == null))] | length')

    log_message "=== MIGRATION SUMMARY ==="
    log_message "Total data sources found: $total_datasources"
    log_message "Migration candidates found: $migration_candidates"
    log_message "Read-only candidates (will be skipped): $readonly_candidates"
    log_message "Updateable candidates: $updateable_candidates"
    log_message "=========================="
}

# Main function to remove Prometheus type migration
remove_prometheus_type_migration() {
    log_message "Starting remove Azure Prometheus migration"
    log_message "Log file: $LOG_FILE"
    log_message "Grafana URL: $GRAFANA_URL"

    response=$(curl -s -w "\n%{http_code}" -X GET \
        -H "Content-Type: application/json" \
        -H "Authorization: Bearer $BEARER_TOKEN" \
        "$GRAFANA_URL/api/datasources/")

    http_code=$(echo "$response" | tail -n1)
    response_body=$(echo "$response" | sed '$d')

    if [[ "$http_code" -ge 200 && "$http_code" -lt 300 ]]; then
        log_message "Successfully fetched data sources"
        get_summary_stats "$response_body"
        update_data_source_type "$response_body"
        log_message "Migration process completed"
    else
        log_message "error fetching data sources: HTTP $http_code - $response_body"
    fi
}

# Function to initialize log file
initialize_log() {
    echo "=== Grafana Azure Prometheus Migration Log ===" > "$LOG_FILE"
    echo "Started at: $(date)" >> "$LOG_FILE"
    echo "=============================================" >> "$LOG_FILE"
    echo "" >> "$LOG_FILE"
}

# Check if jq is installed
if ! command -v jq &> /dev/null; then
    echo "Error: jq is required but not installed. Please install jq to run this script."
    exit 1
fi

# Check if required variables are set
if [[ -z "$GRAFANA_URL" || -z "$BEARER_TOKEN" ]]; then
    echo "Error: Please set GRAFANA_URL and BEARER_TOKEN variables at the top of the script."
    exit 1
fi

# Initialize log file
initialize_log

# Execute main function
log_message "Script started"
remove_prometheus_type_migration
log_message "Script completed"

# Final log message
echo ""
echo "Migration completed. Full log available at: $LOG_FILE"