From 6e9219b52ba19a7e1d76f5af6cbf8024995bdcd6 Mon Sep 17 00:00:00 2001
From: parthiv11 <parthivmakwana11@gmail.com>
Date: Thu, 29 Jun 2023 23:48:20 +0530
Subject: [PATCH] Updated with .pgpass file

---
 plugins/yugabytedb/database_credentials.go    | 72 ++++++++++++++++++-
 .../yugabytedb/database_credentials_test.go   | 36 +++++-----
 2 files changed, 87 insertions(+), 21 deletions(-)

diff --git a/plugins/yugabytedb/database_credentials.go b/plugins/yugabytedb/database_credentials.go
index a941f7fa5..9ec8c7c1e 100644
--- a/plugins/yugabytedb/database_credentials.go
+++ b/plugins/yugabytedb/database_credentials.go
@@ -41,14 +41,80 @@ func DatabaseCredentials() schema.CredentialType {
 				Optional:            true,
 			},
 		},
-		DefaultProvisioner: provision.EnvVars(defaultEnvVarMapping),
-		Importer:           importer.TryEnvVarPair(defaultEnvVarMapping),
+		DefaultProvisioner: provision.TryAll(
+			provision.EnvVars(defaultEnvVarMapping),
+			provision.TempFile(
+				mysqlConfig, 
+				provision.Filename(".pgpass"), 
+				provision.AtFixedPath("~/")
+			)
+		),
+		Importer:           importer.TryAll(
+			importer.TryEnvVarPair(defaultEnvVarMapping),
+			TryMySQLConfigFile("~/.pgpass"),
+		)
 	}
 }
 var defaultEnvVarMapping = map[string]sdk.FieldName{
 	"PGHOST":     fieldname.Host,
 	"PGPORT":     fieldname.Port,
 	"PGUSER":     fieldname.User,
-	"PGPASSWORD": fieldname.Password,
 	"PGDATABASE": fieldname.Database,
 }
+
+func TryPgPassFile() sdk.Importer {
+	return importer.TryFile("~/.pgpass", func(ctx context.Context, contents importer.FileContents, in sdk.ImportInput, out *sdk.ImportAttempt) {
+		fileLines := strings.Split(contents.String(), "\n")
+
+		for _, line := range fileLines {
+			line = strings.TrimSpace(line)
+			if line == "" || strings.HasPrefix(line, "#") {
+				continue
+			}
+
+			fields := strings.Split(line, ":")
+
+			if len(fields) < 5 {
+				out.AddError(errors.New("Invalid .pgpass entry: " + line))
+				continue
+			}
+
+			out.AddCandidate(sdk.ImportCandidate{
+				NameHint: importer.SanitizeNameHint(fields[2]),
+				Fields: map[sdk.FieldName]string{
+					fieldname.Host:     fields[0],
+					fieldname.Port:     fields[1],
+					fieldname.Database: fields[2],
+					fieldname.User:     fields[3],
+					fieldname.Password: fields[4],
+				},
+			})
+		}
+	})
+}
+
+func pgPass(in sdk.ProvisionInput) ([]byte, error) {
+	content := ""
+
+	if host, ok := in.ItemFields[fieldname.Host]; ok {
+		content += host+':'
+	}
+	
+	if port, ok := in.ItemFields[fieldname.Port]; ok {
+		content += port+':'
+	}
+
+	if database, ok := in.ItemFields[fieldname.Database]; ok {
+		content += database+':'
+	}
+
+	if user, ok := in.ItemFields[fieldname.User]; ok {
+		content += user+':'
+	}
+	
+	if password, ok := in.ItemFields[fieldname.Password]; ok {
+		content += password
+	}
+
+	return []byte(content), nil
+}
\ No newline at end of file
diff --git a/plugins/yugabytedb/database_credentials_test.go b/plugins/yugabytedb/database_credentials_test.go
index eea2a8669..4928c365d 100644
--- a/plugins/yugabytedb/database_credentials_test.go
+++ b/plugins/yugabytedb/database_credentials_test.go
@@ -9,28 +9,24 @@ import (
 )
 
 func TestDatabaseCredentialsImporter(t *testing.T) {
+	expectedFields := map[sdk.FieldName]string{
+		fieldname.User:     "root",
+		fieldname.Password: "123456",
+		fieldname.Database: "test",
+		fieldname.Port:     "3306",
+		fieldname.Host:     "localhost",
+	}
+
 	plugintest.TestImporter(t, DatabaseCredentials().Importer, map[string]plugintest.ImportCase{
-		"default": {
-			Environment: map[string]string{
-				"PGHOST":     "localhost",
-				"PGPORT":     "5432",
-				"PGUSER":     "root",
-				"PGPASSWORD": "123456",
-				"PGDATABASE": "test",
+		"pgpass": {
+			Files: map[string]string{
+				"~/.pgpass":       plugintest.LoadFixture(t, ".pgpass"),
 			},
 			ExpectedCandidates: []sdk.ImportCandidate{
-				{
-					Fields: map[sdk.FieldName]string{
-						fieldname.Host:     "localhost",
-						fieldname.Port:     "5432",
-						fieldname.User:     "root",
-						fieldname.Password: "123456",
-						fieldname.Database: "test",
-					},
+				{Fields: expectedFields},
 				},
 			},
-		},
-	})
+		})
 }
 
 func TestDatabaseCredentialsProvisioner(t *testing.T) {
@@ -48,9 +44,13 @@ func TestDatabaseCredentialsProvisioner(t *testing.T) {
 					"PGHOST":     "localhost",
 					"PGPORT":     "5432",
 					"PGUSER":     "root",
-					"PGPASSWORD": "123456",
 					"PGDATABASE": "test",
 				},
+				Files: map[string]OutputFile{
+					"~/.pgpass": {
+						Contents: []byte(plugintest.LoadFixture(t, ".pgpass")),
+					},
+				}
 			},
 		},
 	})