#!/usr/bin/perl

use strict;
use warnings;
use PgCommon;
use Storable qw(dclone);
use YAML;

# from list of queries, remove all that do not satisfy the scope/version/cluster/database constraints
sub filter_queries($$$$;$)
{
  my ($queries, $scope, $version, $cluster, $database) = @_;
  my @result;
  foreach my $query (@$queries) {
    $query->{scope} //= 'database';
    die "scope '$query->{scope}' must be either 'cluster' or 'database'"
      unless ($query->{scope} =~ /^(cluster|database)$/);
    next if $query->{scope} ne $scope;
    next if $query->{min_version} and $version < $query->{min_version};
    next if $query->{max_version} and $version > $query->{max_version};
    next if $query->{version} and $version !~ /$query->{version}/;
    next if $query->{cluster} and $cluster !~ /$query->{cluster}/;
    next if $query->{database} and $database !~ /$query->{database}/;

    my $q = dclone($query);
    delete $q->{scope}; delete $q->{min_version}; delete $q->{max_version};
    delete $q->{version}; delete $q->{cluster}; delete $q->{database};
    push @result, $q;
  }
  return \@result;
}

# Generate jobs from the given queries. One job is created for all the queries
# with null interval, and one job is created per query with non-null interval.
sub push_queries_as_jobs($$$)
{
  my ($jobs, $queries, $job_settings) = @_;
  my @instant;
  foreach my $query (@$queries) {
    if (!$query->{interval}) {
      push @instant, $query;
    } else {
      # Generate separate job for query with non-null interval
      my $job_settings_query = dclone($job_settings);
      $job_settings_query->{interval} = $query->{interval};
      $job_settings_query->{name} .= '/' . $query->{name};
      my $q = dclone($query);
      delete $q->{interval};
      $job_settings_query->{queries} = [$q];

      push @$jobs, $job_settings_query;
    }
  }

  if (@instant) {
    my $instant_job = dclone($job_settings);
    # Debian switched to upstream sql-exporter which does not have synchronous
    # jobs yet, so interval cannot be 0 for instant, or else sql-exporter would
    # run the queries non-stop. Our default Prometheus scrape interval is 300s,
    # so set this to 150s for now.
    $instant_job->{interval} = '150s';
    $instant_job->{queries} = \@instant;
    push @$jobs, $instant_job;
  }
}

my $queries_directory = $ARGV[0] // die "No directory for *.yml query files specified";
my $output_yaml = $ARGV[1] // die "No output yaml file specified";

# load all *.yml files from input directory and collect the contained query lists
my $queries;
foreach my $yml (glob "$queries_directory/*.yml") {
  my $q;
  eval { $q = YAML::LoadFile($yml); };
  die "Error loading $yml: $@" if ($@);
  next if (ref($q) eq ''); # file is empty
  die "$yml is not a yaml list:" . ref($q) unless (ref($q) eq 'ARRAY');
  push @$queries, @$q;
}

# walk all clusters and databases and produce jobs for them
my $jobs = [];
foreach my $version (get_versions()) {
  foreach my $cluster (get_version_clusters($version)) {
    my %info = cluster_info($version, $cluster);
    next unless $info{running}; # cluster is down, skip it
    my $owner = (getpwuid $info{owneruid})[0] // die "Could not determine owner name of cluster $version $cluster";
    my $socket = get_cluster_socketdir($version, $cluster);

    # jobs for cluster-wide queries
    push_queries_as_jobs $jobs, filter_queries($queries, 'cluster', $version, $cluster), {
      connections => [ "postgres://$owner\@:$info{port}/postgres?sslmode=disable&host=$socket" ],
      name => "$version/$cluster",
    };

    # jobs for per-database queries
    my @cluster_databases = get_cluster_databases($version, $cluster);
    foreach my $database (grep { $_ and $_ !~ /^template[01]$/ } @cluster_databases) {
      push_queries_as_jobs $jobs, filter_queries($queries, 'database', $version, $cluster, $database), {
        connections => [ "postgres://$owner\@:$info{port}/$database?sslmode=disable&host=$socket" ],
        name => "$version/$cluster/$database",
      };
    }

    # create pg_stat_statements in "postgres" database if missing
    next if $info{recovery}; # cluster is in recovery mode, skip it
    open PSQL, "|-", "su -c 'psql -q -h $socket -p $info{port} postgres' $owner";
    print PSQL "DO \$\$DECLARE ext name; BEGIN SELECT INTO ext extname FROM pg_extension WHERE extname = 'pg_stat_statements'; IF NOT FOUND THEN RAISE NOTICE 'Creating pg_stat_statements extension in cluster $version/$cluster'; CREATE EXTENSION pg_stat_statements; END IF; END\$\$ LANGUAGE plpgsql;\n";
    close PSQL;
  }
}

# write output yml
YAML::DumpFile($output_yaml, { jobs => $jobs });
