Add flag to enable/disable public access for RDS instances created during Test

pkg/app/rds_postgres.go

@@ -16,9 +16,9 @@ package app
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"os"
+	"strconv"
 	"time"
 
 	awssdk "github.com/aws/aws-sdk-go/aws"
@@ -27,6 +27,7 @@ import (
 	"github.com/ghodss/yaml"
 	"github.com/pkg/errors"
 	v1 "k8s.io/api/core/v1"
+	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 
@@ -37,33 +38,34 @@ import (
 	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/kanisterio/kanister/pkg/log"
-
-	// Initialize pq driver
-	_ "github.com/lib/pq"
 )
 
 type RDSPostgresDB struct {
-	name              string
-	cli               kubernetes.Interface
-	namespace         string
-	id                string
-	host              string
-	databases         []string
-	username          string
-	password          string
-	accessID          string
-	secretKey         string
-	region            string
-	sessionToken      string
-	securityGroupID   string
-	securityGroupName string
-	sqlDB             *sql.DB
-	configMapName     string
-	secretName        string
+	name                     string
+	cli                      kubernetes.Interface
+	namespace                string
+	id                       string
+	host                     string
+	databases                []string
+	dbSubnetGroup            string
+	username                 string
+	password                 string
+	accessID                 string
+	secretKey                string
+	region                   string
+	sessionToken             string
+	securityGroupID          string
+	securityGroupName        string
+	configMapName            string
+	secretName               string
+	bastionDebugWorkloadName string
+	publicAccess             bool
+	vpcID                    string
 }
 
 const (
-	dbInstanceType = "db.t3.micro"
+	dbInstanceType           = "db.t3.micro"
+	postgresConnectionString = "PGPASSWORD=%s psql -h %s -p 5432 -U %s -d %s -t -c"
 )
 
 func NewRDSPostgresDB(name string, customRegion string) App {
@@ -131,30 +133,74 @@ func (pdb *RDSPostgresDB) Install(ctx context.Context, ns string) error {
 		return err
 	}
 
+	// Create rds client
+	rdsCli, err := rds.NewClient(ctx, awsConfig, region)
+	if err != nil {
+		return err
+	}
+
+	pdb.bastionDebugWorkloadName = fmt.Sprintf("%s-workload", pdb.name)
+
+	deploymentSpec := bastionDebugWorkloadSpec(ctx, pdb.bastionDebugWorkloadName, "postgres", pdb.namespace)
+	_, err = pdb.cli.AppsV1().Deployments(pdb.namespace).Create(ctx, deploymentSpec, metav1.CreateOptions{})
+	if err != nil {
+		return errors.Wrapf(err, "Failed to create test deployment %s , app: %s", pdb.bastionDebugWorkloadName, pdb.name)
+	}
+
+	if err := kube.WaitOnDeploymentReady(ctx, pdb.cli, pdb.namespace, pdb.bastionDebugWorkloadName); err != nil {
+		return errors.Wrapf(err, "Failed while waiting for deployment %s to be ready, app: %s", pdb.bastionDebugWorkloadName, pdb.name)
+	}
+
+	pdb.vpcID = os.Getenv("VPC_ID")
+
+	// VPCId is not provided, use Default VPC
+	if pdb.vpcID == "" {
+		defaultVpc, err := ec2Cli.DescribeDefaultVpc(ctx)
+		if err != nil {
+			return err
+		}
+		if len(defaultVpc.Vpcs) == 0 {
+			return fmt.Errorf("No default VPC found")
+		}
+		pdb.vpcID = *defaultVpc.Vpcs[0].VpcId
+	}
+
+	// describe subnets in the VPC
+	resp, err := ec2Cli.DescribeSubnets(ctx, pdb.vpcID)
+	if err != nil {
+		return errors.Wrapf(err, "Failed to describe subnets")
+	}
+
+	// Extract subnet IDs from the response
+	var subnetIDs []string
+	for _, subnet := range resp.Subnets {
+		subnetIDs = append(subnetIDs, *subnet.SubnetId)
+	}
+	// create a subnetgroup with subnets in the VPC
+	subnetGroup, err := rdsCli.CreateDBSubnetGroup(ctx, fmt.Sprintf("%s-subnetgroup", pdb.name), "kanister-test-subnet-group", subnetIDs)
+	if err != nil {
+		return errors.Wrapf(err, "Failed to create subnet group")
+	}
+	pdb.dbSubnetGroup = *subnetGroup.DBSubnetGroup.DBSubnetGroupName
+
 	// Create security group
 	log.Info().Print("Creating security group.", field.M{"app": pdb.name, "name": pdb.securityGroupName})
-	sg, err := ec2Cli.CreateSecurityGroup(ctx, pdb.securityGroupName, "kanister-test-security-group")
+	sg, err := ec2Cli.CreateSecurityGroup(ctx, pdb.securityGroupName, "kanister-test-security-group", pdb.vpcID)
 	if err != nil {
 		return err
 	}
 	pdb.securityGroupID = *sg.GroupId
 
 	// Add ingress rule
 	log.Info().Print("Adding ingress rule to security group.", field.M{"app": pdb.name})
-	_, err = ec2Cli.AuthorizeSecurityGroupIngress(ctx, pdb.securityGroupName, "0.0.0.0/0", "tcp", 5432)
-	if err != nil {
-		return err
-	}
-
-	// Create rds client
-	rdsCli, err := rds.NewClient(ctx, awsConfig, region)
+	_, err = ec2Cli.AuthorizeSecurityGroupIngress(ctx, pdb.securityGroupID, "0.0.0.0/0", "tcp", 5432)
 	if err != nil {
 		return err
 	}
 
 	// Create RDS instance
 	log.Info().Print("Creating RDS instance.", field.M{"app": pdb.name, "id": pdb.id})
-	_, err = rdsCli.CreateDBInstance(ctx, 20, dbInstanceType, pdb.id, "postgres", pdb.username, pdb.password, []string{pdb.securityGroupID})
+	_, err = rdsCli.CreateDBInstance(ctx, 20, dbInstanceType, pdb.id, "postgres", pdb.username, pdb.password, []string{pdb.securityGroupID}, pdb.publicAccess)
 	if err != nil {
 		return err
 	}
@@ -234,6 +280,7 @@ func (pdb *RDSPostgresDB) Object() crv1alpha1.ObjectReference {
 
 // Ping makes and tests DB connection
 func (pdb *RDSPostgresDB) Ping(ctx context.Context) error {
+	log.Print("Pinging rds postgres database", field.M{"app": pdb.name})
 	// Get connection info from configmap
 	dbconfig, err := pdb.cli.CoreV1().ConfigMaps(pdb.namespace).Get(ctx, pdb.configMapName, metav1.GetOptions{})
 	if err != nil {
@@ -255,63 +302,74 @@ func (pdb *RDSPostgresDB) Ping(ctx context.Context) error {
 		return errors.New("Databases are missing from configmap")
 	}
 
-	var connectionString string = fmt.Sprintf("host=%s user=%s password=%s dbname=%s sslmode=disable", dbconfig.Data["postgres.host"], dbconfig.Data["postgres.user"], dbsecret.Data["password"], databases[0])
+	isReadyQuery := fmt.Sprintf(postgresConnectionString+"'SELECT version();'", dbsecret.Data["password"], dbconfig.Data["postgres.host"], dbconfig.Data["postgres.user"], databases[0])
 
-	// Initialize connection object.
-	db, err := sql.Open("postgres", connectionString)
-	if err != nil {
-		return err
-	}
+	pingCommand := []string{"sh", "-c", isReadyQuery}
 
-	err = db.Ping()
+	_, stderr, err := pdb.execCommand(ctx, pingCommand)
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "Error while Pinging the database: %s, app: %s", stderr, pdb.name)
 	}
-
-	pdb.sqlDB = db
-	log.Info().Print("Connected to database.", field.M{"app": pdb.name})
+	log.Print("Ping to the application was successful.", field.M{"app": pdb.name})
 	return nil
 }
 
 func (pdb RDSPostgresDB) Insert(ctx context.Context) error {
+	log.Print("Adding entry to database", field.M{"app": pdb.name})
 	now := time.Now().Format(time.RFC3339Nano)
-	stmt := "INSERT INTO inventory (name) VALUES ($1);"
-	_, err := pdb.sqlDB.Exec(stmt, now)
+	insertQuery := fmt.Sprintf(postgresConnectionString+
+		"\"INSERT INTO inventory (name) VALUES ('%s');\"", pdb.password, pdb.host, pdb.username, pdb.databases[0], now)
+
+	insertCommand := []string{"sh", "-c", insertQuery}
+	_, stderr, err := pdb.execCommand(ctx, insertCommand)
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "Error while inserting data into table: %s, app: %s", stderr, pdb.name)
 	}
 	log.Info().Print("Inserted a row in test db.", field.M{"app": pdb.name})
 	return nil
 }
 
 func (pdb RDSPostgresDB) Count(ctx context.Context) (int, error) {
-	stmt := "SELECT COUNT(*) FROM inventory;"
-	row := pdb.sqlDB.QueryRow(stmt)
-	var count int
-	err := row.Scan(&count)
+	log.Print("Counting entries from database", field.M{"app": pdb.name})
+	countQuery := fmt.Sprintf(postgresConnectionString+
+		"\"SELECT COUNT(*) FROM inventory;\"", pdb.password, pdb.host, pdb.username, pdb.databases[0])
+
+	countCommand := []string{"sh", "-c", countQuery}
+	stdout, stderr, err := pdb.execCommand(ctx, countCommand)
 	if err != nil {
-		return 0, err
+		return 0, errors.Wrapf(err, "Error while counting data into table: %s, app: %s", stderr, pdb.name)
 	}
-	log.Info().Print("Counting rows in test db.", field.M{"app": pdb.name, "count": count})
-	return count, nil
+
+	rowsReturned, err := strconv.Atoi(stdout)
+	if err != nil {
+		return 0, errors.Wrapf(err, "Error while converting response of count query: %s, app: %s", stderr, pdb.name)
+	}
+
+	log.Info().Print("Counting rows in test db.", field.M{"app": pdb.name, "count": rowsReturned})
+	return rowsReturned, nil
 }
 
 func (pdb RDSPostgresDB) Reset(ctx context.Context) error {
-	_, err := pdb.sqlDB.Exec("DROP TABLE IF EXISTS inventory;")
+	log.Print("Resetting database", field.M{"app": pdb.name})
+	deleteQuery := fmt.Sprintf(postgresConnectionString+"\"DROP TABLE IF EXISTS inventory;\"", pdb.password, pdb.host, pdb.username, pdb.databases[0])
+	deleteCommand := []string{"sh", "-c", deleteQuery}
+	_, stderr, err := pdb.execCommand(ctx, deleteCommand)
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "Error while deleting data from table: %s, app: %s", stderr, pdb.name)
 	}
-
 	log.Info().Print("Database reset successful!", field.M{"app": pdb.name})
 	return nil
 }
 
 // Initialize is used to initialize the database or create schema
 func (pdb RDSPostgresDB) Initialize(ctx context.Context) error {
 	// Create table.
-	_, err := pdb.sqlDB.Exec("CREATE TABLE inventory (id serial PRIMARY KEY, name VARCHAR(50));")
+	log.Print("Initializing database", field.M{"app": pdb.name})
+	createQuery := fmt.Sprintf(postgresConnectionString+"\"CREATE TABLE inventory (id serial PRIMARY KEY, name VARCHAR(50));\"", pdb.password, pdb.host, pdb.username, pdb.databases[0])
+	createCommand := []string{"sh", "-c", createQuery}
+	_, stderr, err := pdb.execCommand(ctx, createCommand)
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "Error while initializing the database: %s, app: %s", stderr, pdb.name)
 	}
 	return nil
 }
@@ -377,6 +435,21 @@ func (pdb RDSPostgresDB) Uninstall(ctx context.Context) error {
 		return errors.Wrap(err, "Failed to ec2 client. You may need to delete EC2 resources manually. app=rds-postgresql")
 	}
 
+	// Delete dbSubnetGroup
+	log.Info().Print("Deleting db subnet group.", field.M{"app": pdb.name})
+	_, err = rdsCli.DeleteDBSubnetGroup(ctx, pdb.dbSubnetGroup)
+	if err != nil {
+		// If the subnet group does not exist, ignore the error and return
+		if err, ok := err.(awserr.Error); ok {
+			switch err.Code() {
+			case awsrds.ErrCodeDBSubnetGroupNotFoundFault:
+				log.Info().Print("Subnet Group Does not exist: ErrCodeDBSubnetGroupNotFoundFault.", field.M{"app": pdb.name, "id": pdb.id})
+			default:
+				return errors.Wrapf(err, "Failed to delete db subnet group. You may need to delete it manually. app=rds-postgresql name=%s", pdb.dbSubnetGroup)
+			}
+		}
+	}
+
 	// Delete security group
 	log.Info().Print("Deleting security group.", field.M{"app": pdb.name})
 	_, err = ec2Cli.DeleteSecurityGroup(ctx, pdb.securityGroupID)
@@ -390,6 +463,12 @@ func (pdb RDSPostgresDB) Uninstall(ctx context.Context) error {
 			}
 		}
 	}
+	// Remove workload object created for executing commands
+	err = pdb.cli.AppsV1().Deployments(pdb.namespace).Delete(ctx, pdb.bastionDebugWorkloadName, metav1.DeleteOptions{})
+	if err != nil && !apierrors.IsNotFound(err) {
+		return errors.Wrapf(err, "Error deleting Workload name=%s app=%s", pdb.bastionDebugWorkloadName, pdb.name)
+	}
+
 	return nil
 }
 
@@ -413,3 +492,11 @@ func makeYamlList(dbs []string) string {
 	}
 	return dbsYaml
 }
+
+func (pdb RDSPostgresDB) execCommand(ctx context.Context, command []string) (string, string, error) {
+	podName, containerName, err := kube.GetPodContainerFromDeployment(ctx, pdb.cli, pdb.namespace, pdb.bastionDebugWorkloadName)
+	if err != nil || podName == "" {
+		return "", "", err
+	}
+	return kube.Exec(pdb.cli, pdb.namespace, podName, containerName, command, nil)
+}

pkg/app/utils.go

@@ -15,8 +15,12 @@
 package app
 
 import (
+	"context"
 	"fmt"
 
+	appsv1 "k8s.io/api/apps/v1"
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/rand"
 )
 
@@ -61,3 +65,32 @@ func getOpenShiftDBTemplate(appName string, templateVersion DBTemplate, storageT
 func getLabelOfApp(appName string, storageType storage) string {
 	return fmt.Sprintf("app=%s-%s", appName, storageType)
 }
+
+// bastionDebugWorkloadSpec creates Deployment Resource Manifest from which RDS database queries can be executed
+func bastionDebugWorkloadSpec(ctx context.Context, name string, image string, namespace string) *appsv1.Deployment {
+	return &appsv1.Deployment{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+		},
+		Spec: appsv1.DeploymentSpec{
+			Selector: &metav1.LabelSelector{MatchLabels: map[string]string{"app": name}},
+			Template: v1.PodTemplateSpec{
+				ObjectMeta: metav1.ObjectMeta{
+					Labels: map[string]string{
+						"app": name,
+					},
+				},
+				Spec: v1.PodSpec{
+					Containers: []v1.Container{
+						{
+							Name:    name,
+							Image:   image,
+							Command: []string{"sleep", "infinity"},
+						},
+					},
+				},
+			},
+		},
+	}
+}