From 92ae18b1725a716ce6bcd7f30bd7102bd9e44032 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=BCrkan=20=C4=B0ndibay?= <gindibay@microsoft.com>
Date: Mon, 23 Aug 2021 13:19:27 +0300
Subject: [PATCH] nightly_version_assignment reference error when  nightlypg
 param is missing in pkgvars (#157)

---
 packaging_automation/citus_package.py | 36 +++++++++------------------
 packaging_automation/requirements.in  |  2 +-
 packaging_automation/requirements.txt |  4 ++-
 3 files changed, 16 insertions(+), 26 deletions(-)

diff --git a/packaging_automation/citus_package.py b/packaging_automation/citus_package.py
index 3c3bb082..62172929 100644
--- a/packaging_automation/citus_package.py
+++ b/packaging_automation/citus_package.py
@@ -5,9 +5,10 @@
 from enum import Enum
 from typing import List
 from typing import Tuple
-import gnupg
-import docker
 
+import docker
+import gnupg
+from dotenv import dotenv_values
 from parameters_validation import non_blank, non_empty, validate_parameters
 
 from .common_tool_methods import (run_with_output, PackageType, transform_key_into_base64_str,
@@ -160,32 +161,19 @@ def sign_packages(base_output_path: str, sub_folder: str, secret_key: str, passp
 
 
 def get_postgres_versions(os_name: str, input_files_dir: str) -> Tuple[List[str], List[str]]:
-    release_versions = []
-    nightly_versions = []
     if platform_postgres_version_source[os_name] == PostgresVersionDockerImageType.single:
         release_versions = ["all"]
         nightly_versions = ["all"]
     else:
-        with open(f"{input_files_dir}/pkgvars", "r") as reader:
-            content = reader.read()
-            lines = content.splitlines()
-            for line in lines:
-                if line.startswith("releasepg"):
-                    release_version_assignment = line
-                if line.startswith("nightlypg"):
-                    nightly_version_assignment = line
-            if release_version_assignment is None or "=" not in release_version_assignment or len(
-                    release_version_assignment.split("=")) != 2:
-                raise ValueError(
-                    f"Release version in pkglatest is not well formatted. Expected format: releasepg=12,13 "
-                    f"Actual Format:{release_version_assignment}")
-            if nightly_version_assignment is None or "=" not in nightly_version_assignment or len(
-                    nightly_version_assignment.split("=")) != 2:
-                raise ValueError(
-                    f"Nightly version in pkglatest is not well formatted. Expected format: nightlypg=12,13 "
-                    f"Actual Format:{nightly_version_assignment}")
-            release_versions = release_version_assignment.split("=")[1].split(",")
-            nightly_versions = nightly_version_assignment.split("=")[1].split(",")
+        pkgvars_config = dotenv_values(f"{input_files_dir}/{PKGVARS_FILE}")
+        release_versions_str = pkgvars_config['releasepg']
+        if "nightlypg" in pkgvars_config:
+            nightly_versions_str = pkgvars_config['nightlypg']
+        else:
+            nightly_versions_str = release_versions_str
+
+        release_versions = release_versions_str.split(",")
+        nightly_versions = nightly_versions_str.split(",")
     return release_versions, nightly_versions
 
 
diff --git a/packaging_automation/requirements.in b/packaging_automation/requirements.in
index 67db959e..1d9a29a8 100644
--- a/packaging_automation/requirements.in
+++ b/packaging_automation/requirements.in
@@ -14,4 +14,4 @@ requests
 SQLAlchemy
 urllib3
 wheel
-
+python-dotenv
diff --git a/packaging_automation/requirements.txt b/packaging_automation/requirements.txt
index 400713c4..7677d605 100644
--- a/packaging_automation/requirements.txt
+++ b/packaging_automation/requirements.txt
@@ -1,5 +1,5 @@
 #
-# This file is autogenerated by pip-compile with python 3.8
+# This file is autogenerated by pip-compile
 # To update, run:
 #
 #    pip-compile
@@ -56,6 +56,8 @@ pyparsing==2.4.7
     # via packaging
 pytest==6.2.4
     # via -r requirements.in
+python-dotenv==0.19.0
+    # via -r requirements.in
 python-gnupg==0.4.7
     # via -r requirements.in
 python-string-utils==1.0.0