From 204660fc4bcb2f8bab2cfbef37c883230e887ccf Mon Sep 17 00:00:00 2001
From: saikumar-vs <112166980+saikumar-vs@users.noreply.github.com>
Date: Thu, 19 Dec 2024 12:35:22 +0530
Subject: [PATCH] Updates to fix PGSQL for PGSM (#103)

---
 pmm_qa/pdpgsql_pgsm_setup.yml   | 5 +++--
 pmm_qa/pg_stat_monitor_setup.sh | 2 +-
 pmm_qa/pmm-framework.py         | 9 ++++++---
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/pmm_qa/pdpgsql_pgsm_setup.yml b/pmm_qa/pdpgsql_pgsm_setup.yml
index 3046548..fd3ef2d 100644
--- a/pmm_qa/pdpgsql_pgsm_setup.yml
+++ b/pmm_qa/pdpgsql_pgsm_setup.yml
@@ -12,9 +12,10 @@
     client_version: "{{ lookup('vars', 'extra_client_version', default=lookup('env','CLIENT_VERSION') | default('3-dev-latest', true) ) }}"
     admin_password: "{{ lookup('vars', 'extra_admin_password', default=lookup('env','ADMIN_PASSWORD') | default('admin', true) ) }}"
     pmm_qa_branch: "{{ lookup('vars', 'extra_pmm_qa_branch', default=lookup('env','PMM_QA_GIT_BRANCH') | default('main', true) ) }}"
-    pgstat_monitor_branch: "{{ lookup('vars', 'extra_pgstat_monitor_branch', default=lookup('env','PGSTAT_MONITOR_BRANCH') | default('1.1.1', true) ) }}"
+    pgstat_monitor_branch: "{{ lookup('vars', 'extra_pgstat_monitor_branch', default=lookup('env','PGSTAT_MONITOR_BRANCH') | default('2.1.0', true) ) }}"
     pgstat_monitor_repo: "{{ lookup('vars', 'extra_pgstat_monitor_repo', default=lookup('env','PGSTAT_MONITOR_REPO') | default('percona/pg_stat_monitor', true) ) }}"
     use_socket: "{{ lookup('vars', 'extra_pdpgsql_version', default=lookup('env','USE_SOCKET') | default('', true) ) }}"
+    distribution: "{{ lookup('vars', 'extra_pdpgsql_distribution', default=lookup('env','DISTRIBUTION') | default('PPG', true) ) }}"
 
   tasks:
   - name: cleanup container for client and DB setup
@@ -45,7 +46,7 @@
   - name: Execute Setup script inside the pdpgsql pdpgsql_pgsm_container
     shell: "{{ item }}"
     with_items:
-      - docker exec {{ pdpgsql_pgsm_container }} bash -xe ./pg_stat_monitor_setup.sh --distribution PPG --pgsql_version {{ pdpgsql_version }} --pgstat_monitor_branch {{ pgstat_monitor_branch }} --pgstat_monitor_repo {{ pgstat_monitor_repo }} > setup_pdpgsql_pgsm_{{ pdpgsql_version }}.log
+      - docker exec {{ pdpgsql_pgsm_container }} bash -xe ./pg_stat_monitor_setup.sh --distribution {{ distribution }} --pgsql_version {{ pdpgsql_version }} --pgstat_monitor_branch {{ pgstat_monitor_branch }} --pgstat_monitor_repo {{ pgstat_monitor_repo }} > setup_pdpgsql_pgsm_{{ pdpgsql_version }}.log
 
   - name: Install pmm2-client on the pdpgsql_pgsm_container
     shell: "{{ item }}"
diff --git a/pmm_qa/pg_stat_monitor_setup.sh b/pmm_qa/pg_stat_monitor_setup.sh
index c3f5d6c..cc6fb54 100644
--- a/pmm_qa/pg_stat_monitor_setup.sh
+++ b/pmm_qa/pg_stat_monitor_setup.sh
@@ -19,7 +19,7 @@ fi
 # If branch/tag is not provided then it will default to main branch
 if [ -z "$pgstat_monitor_branch" ]
 then
-      export pgstat_monitor_branch="1.1.1"
+      export pgstat_monitor_branch="2.1.0"
 fi
 
 # If repo is not provided then it will default to percona PGSM repository
diff --git a/pmm_qa/pmm-framework.py b/pmm_qa/pmm-framework.py
index cf5ea16..129d560 100755
--- a/pmm_qa/pmm-framework.py
+++ b/pmm_qa/pmm-framework.py
@@ -284,6 +284,7 @@ def setup_pdpgsql(db_type, db_version=None, db_config=None, args=None):
         'CLIENT_VERSION': get_value('CLIENT_VERSION', db_type, args, db_config),
         'USE_SOCKET': get_value('USE_SOCKET', db_type, args, db_config),
         'ADMIN_PASSWORD': os.getenv('ADMIN_PASSWORD') or args.pmm_server_password or 'admin',
+        'DISTRIBUTION': '',
         'PMM_QA_GIT_BRANCH': os.getenv('PMM_QA_GIT_BRANCH') or 'v3'
     }
 
@@ -335,17 +336,19 @@ def setup_pgsql(db_type, db_version=None, db_config=None, args=None):
 
     # Define environment variables for playbook
     env_vars = {
-        'PGSQL_VERSION': pgsql_version,
+        'PDPGSQL_VERSION': pgsql_version,
+        'PGSTAT_MONITOR_BRANCH': 'main',
         'PMM_SERVER_IP': args.pmm_server_ip or container_name or '127.0.0.1',
-        'PGSQL_PGSS_CONTAINER': 'pgsql_pgss_pmm_' + str(pgsql_version),
+        'PDPGSQL_PGSM_CONTAINER': 'pgsql_pgsm_pmm_' + str(pgsql_version),
         'CLIENT_VERSION': get_value('CLIENT_VERSION', db_type, args, db_config),
         'USE_SOCKET': get_value('USE_SOCKET', db_type, args, db_config),
         'ADMIN_PASSWORD': os.getenv('ADMIN_PASSWORD') or args.pmm_server_password or 'admin',
+        'DISTRIBUTION': 'PGDG',
         'PMM_QA_GIT_BRANCH': os.getenv('PMM_QA_GIT_BRANCH') or 'v3'
     }
 
     # Ansible playbook filename
-    playbook_filename = 'pgsql_pgss_setup.yml'
+    playbook_filename = 'pdpgsql_pgsm_setup.yml'
 
     # Call the function to run the Ansible playbook
     run_ansible_playbook(playbook_filename, env_vars, args)