apiVersion: v1
kind: ConfigMap
metadata:
  name: datadog-init-script
data:
  init.sh: |
    #!/bin/bash

    set -xe
    export POSTGRES_HOST=$1
    export POSTGRES_PASSWORD=$2
    export DATADOG_PASSWORD=$3
    export PGPASSWORD=$POSTGRES_PASSWORD

    # Check if datadog user already exists
    existsCheckResult=$(psql -h $POSTGRES_HOST -U postgres postgres -t -c "SELECT 1 FROM pg_catalog.pg_user WHERE usename = 'datadog';")
    result=$(echo $result | xargs)


    # If user doesn't exist, create user and schema
    if [[ -z "$result" ]]; then
      echo "Datadog user not found. Creating user and schema..."
      psql -h $POSTGRES_HOST -d postgres -U postgres <<EOF
      CREATE USER datadog WITH password '$DATADOG_PASSWORD';
      CREATE SCHEMA datadog;
      GRANT USAGE ON SCHEMA datadog TO datadog;
      GRANT USAGE ON SCHEMA public TO datadog;
      GRANT pg_monitor TO datadog;
      CREATE EXTENSION IF NOT EXISTS pg_stat_statements;

      CREATE OR REPLACE FUNCTION datadog.explain_statement(
          l_query TEXT,
          OUT explain JSON
      )
      RETURNS SETOF JSON AS
      \$\$
      DECLARE
      curs REFCURSOR;
      plan JSON;

      BEGIN
          OPEN curs FOR EXECUTE pg_catalog.concat('EXPLAIN (FORMAT JSON) ', l_query);
          FETCH curs INTO plan;
          CLOSE curs;
          RETURN QUERY SELECT plan;
      END;
      \$\$
      LANGUAGE 'plpgsql'
      RETURNS NULL ON NULL INPUT
      SECURITY DEFINER;
    EOF
    else
      echo "Datadog user already exists. Skipping user and schema creation."
    fi

    # Verify permissions on remote host
    echo "Verifying permissions on $POSTGRES_HOST..."
    export PGPASSWORD=$DATADOG_PASSWORD

    psql -h $POSTGRES_HOST -U datadog postgres -A \
      -c "select * from pg_stat_database limit 1;" \
      && echo -e "\e[0;32mPostgres connection - OK\e[0m" \
      || echo -e "\e[0;31mCannot connect to Postgres\e[0m"

    psql -h $POSTGRES_HOST -U datadog postgres -A \
      -c "select * from pg_stat_activity limit 1;" \
      && echo -e "\e[0;32mPostgres pg_stat_activity read OK\e[0m" \
      || echo -e "\e[0;31mCannot read from pg_stat_activity\e[0m"

    psql -h $POSTGRES_HOST -U datadog postgres -A \
      -c "select * from pg_stat_statements limit 1;" \
      && echo -e "\e[0;32mPostgres pg_stat_statements read OK\e[0m" \
      || echo -e "\e[0;31mCannot read from pg_stat_statements\e[0m"

    echo "Script completed!"
  init.py: |
    #!/usr/bin/env python3

    import os
    import json
    import boto3
    import subprocess

    def fetch_secret(secret_arn):
        client = boto3.client('secretsmanager')
        response = client.get_secret_value(SecretId=secret_arn)
        secret = response['SecretString']
        return secret

    def run_init_script(endpoint, secret_value):
        password = json.loads(secret_value)['password']
        subprocess.run(['bash', 'init.sh', endpoint, password, os.getenv('DATADOG_PASSWORD')])

    if __name__ == "__main__":
        databases_env = os.getenv('DATABASES')
        # Parse the JSON string into a list of dictionaries looking like this:
        # [
        #   {
        #     "endpoint": "rds-endpoint
        #     "secretArn": "arn:aws:secretsmanager:us-west-2:123456789012:secret:my-secret
        #   }
        # ]
        databases = json.loads(databases_env)
        for entry in databases:
            if 'endpoint' in entry and 'secretArn' in entry:
                secret_value = fetch_secret(entry['secretArn'])
                run_init_script(entry['endpoint'], secret_value)
---
apiVersion: batch/v1
kind: Job
metadata:
  name: datadog-init-job
  annotations:
    helm.sh/hook: pre-install,pre-upgrade
spec:
  template:
    spec:
      restartPolicy: OnFailure
      serviceAccountName: datadog-init-job
      containers:
      - name: datadog-init
        image: alpine:latest
        command: ["/bin/sh", "-c"]
        args:
        - |-
          apk add bash postgresql-client
          apk add --no-cache python3 py3-pip
          apk add py3-boto3
          python3 init.py
        envFrom:
        - secretRef:
            name: datadog-init-job
        volumeMounts:
        - name: datadog-init-script
          mountPath: /init.sh
          subPath: init.sh
        - name: datadog-init-script
          mountPath: /init.py
          subPath: init.py
      volumes:
      - name: datadog-init-script
        configMap:
          name: datadog-init-script