Refactor GetConnection to be able to get the pq error code of the error (#72)

Author: pcallewaert

.gitignore

@@ -76,6 +76,9 @@ tags
 .history
 # End of https://www.gitignore.io/api/go,vim,emacs,visualstudiocode
 
+### MacOS
+.DS_Store
+
 .idea
 deploy/secret.yaml
 # build artifact

pkg/postgres/database.go

@@ -35,10 +35,13 @@ func (c *pg) CreateDB(dbname, role string) error {
 }
 
 func (c *pg) CreateSchema(db, role, schema string, logger logr.Logger) error {
-	tmpDb := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	if err != nil {
+		return err
+	}
 	defer tmpDb.Close()
 
-	_, err := tmpDb.Exec(fmt.Sprintf(CREATE_SCHEMA, schema, role))
+	_, err = tmpDb.Exec(fmt.Sprintf(CREATE_SCHEMA, schema, role))
 	if err != nil {
 		return err
 	}
@@ -58,22 +61,28 @@ func (c *pg) DropDatabase(database string, logger logr.Logger) error {
 }
 
 func (c *pg) CreateExtension(db, extension string, logger logr.Logger) error {
-	tmpDb := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	if err != nil {
+		return err
+	}
 	defer tmpDb.Close()
 
-	_, err := tmpDb.Exec(fmt.Sprintf(CREATE_EXTENSION, extension))
+	_, err = tmpDb.Exec(fmt.Sprintf(CREATE_EXTENSION, extension))
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
 func (c *pg) SetSchemaPrivileges(db, creator, role, schema, privs string, logger logr.Logger) error {
-	tmpDb := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+	if err != nil {
+		return err
+	}
 	defer tmpDb.Close()
 
 	// Grant role usage on schema
-	_, err := tmpDb.Exec(fmt.Sprintf(GRANT_USAGE_SCHEMA, schema, role))
+	_, err = tmpDb.Exec(fmt.Sprintf(GRANT_USAGE_SCHEMA, schema, role))
 	if err != nil {
 		return err
 	}

pkg/postgres/postgres.go

@@ -36,8 +36,13 @@ type pg struct {
 }
 
 func NewPG(host, user, password, uri_args, default_database, cloud_type string, logger logr.Logger) (PG, error) {
+	db, err := GetConnection(user, password, host, default_database, uri_args, logger)
+	if err != nil {
+		log.Fatalf("failed to connect to PostgreSQL server: %s", err.Error())
+	}
+	logger.Info("connected to postgres server")
 	postgres := &pg{
-		db:               GetConnection(user, password, host, default_database, uri_args, logger),
+		db:               db,
 		log:              logger,
 		host:             host,
 		user:             user,
@@ -64,15 +69,11 @@ func (c *pg) GetDefaultDatabase() string {
 	return c.default_database
 }
 
-func GetConnection(user, password, host, database, uri_args string, logger logr.Logger) *sql.DB {
+func GetConnection(user, password, host, database, uri_args string, logger logr.Logger) (*sql.DB, error) {
 	db, err := sql.Open("postgres", fmt.Sprintf("postgresql://%s:%s@%s/%s?%s", user, password, host, database, uri_args))
 	if err != nil {
 		log.Fatal(err)
 	}
 	err = db.Ping()
-	if err != nil {
-		log.Fatalf("failed to connect to PostgreSQL server: %s", err.Error())
-	}
-	logger.Info("connected to postgres server")
-	return db
+	return db, err
 }

pkg/postgres/role.go

@@ -62,8 +62,15 @@ func (c *pg) RevokeRole(role, revoked string) error {
 
 func (c *pg) DropRole(role, newOwner, database string, logger logr.Logger) error {
 	// REASSIGN OWNED BY only works if the correct database is selected
-	tmpDb := GetConnection(c.user, c.pass, c.host, database, c.args, logger)
-	_, err := tmpDb.Exec(fmt.Sprintf(REASIGN_OBJECTS, role, newOwner))
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, database, c.args, logger)
+	if err != nil {
+		if err.(*pq.Error).Code == "3D000" {
+			return nil // Database is does not exist (anymore)
+		} else {
+			return err
+		}
+	}
+	_, err = tmpDb.Exec(fmt.Sprintf(REASIGN_OBJECTS, role, newOwner))
 	defer tmpDb.Close()
 	// Check if error exists and if different from "ROLE NOT FOUND" => 42704
 	if err != nil && err.(*pq.Error).Code != "42704" {