From 8721b3e27bff4df5972e053f33cc1b447d2dce06 Mon Sep 17 00:00:00 2001
From: Chao-Han Tsai <milton0825@gmail.com>
Date: Tue, 8 Jun 2021 11:49:48 -0700
Subject: [PATCH] Fix connection error handling (#45)

---
 datacatalog/cmd/entrypoints/migrate.go        | 16 +++++++++++---
 datacatalog/go.mod                            |  2 +-
 .../pkg/repositories/errors/postgres.go       | 15 +++++++++----
 .../pkg/repositories/gormimpl/tag_test.go     | 21 +++++++++++++++++--
 4 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/datacatalog/cmd/entrypoints/migrate.go b/datacatalog/cmd/entrypoints/migrate.go
index 9663538123..91067ce627 100644
--- a/datacatalog/cmd/entrypoints/migrate.go
+++ b/datacatalog/cmd/entrypoints/migrate.go
@@ -1,11 +1,15 @@
 package entrypoints
 
 import (
+	"reflect"
+
 	"github.com/flyteorg/datacatalog/pkg/repositories"
+	errors2 "github.com/flyteorg/datacatalog/pkg/repositories/errors"
 	"github.com/flyteorg/datacatalog/pkg/runtime"
 	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/flyteorg/flytestdlib/promutils"
-	"github.com/lib/pq"
+
+	"github.com/jackc/pgconn"
 
 	"context"
 
@@ -38,8 +42,14 @@ var migrateCmd = &cobra.Command{
 
 		if err != nil {
 			// if db does not exist, try creating it
-			pqError, ok := err.(*pq.Error)
-			if ok && pqError.Code == pqInvalidDBCode {
+			cErr, ok := err.(errors2.ConnectError)
+			if !ok {
+				logger.Errorf(ctx, "Failed to cast error of type: %v, err: %v", reflect.TypeOf(err),
+					err)
+				panic(err)
+			}
+			pqError := cErr.Unwrap().(*pgconn.PgError)
+			if pqError.Code == pqInvalidDBCode {
 				logger.Warningf(ctx, "Database [%v] does not exist, trying to create it now", dbName)
 
 				dbConfigValues.DbName = defaultDB
diff --git a/datacatalog/go.mod b/datacatalog/go.mod
index 346c2e5c70..04267ceafa 100644
--- a/datacatalog/go.mod
+++ b/datacatalog/go.mod
@@ -8,7 +8,7 @@ require (
 	github.com/flyteorg/flytestdlib v0.3.13
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
 	github.com/golang/protobuf v1.4.3
-	github.com/lib/pq v1.3.0
+	github.com/jackc/pgconn v1.8.1
 	github.com/mitchellh/mapstructure v1.4.1
 	github.com/spf13/cobra v1.1.1
 	github.com/spf13/pflag v1.0.5
diff --git a/datacatalog/pkg/repositories/errors/postgres.go b/datacatalog/pkg/repositories/errors/postgres.go
index d68a93212f..bbdd7d595d 100644
--- a/datacatalog/pkg/repositories/errors/postgres.go
+++ b/datacatalog/pkg/repositories/errors/postgres.go
@@ -3,8 +3,9 @@ package errors
 import (
 	"fmt"
 
+	"github.com/jackc/pgconn"
+
 	"github.com/flyteorg/datacatalog/pkg/errors"
-	"github.com/lib/pq"
 	"google.golang.org/grpc/codes"
 	"gorm.io/gorm"
 )
@@ -20,7 +21,7 @@ type postgresErrorTransformer struct {
 
 const (
 	unexpectedType            = "unexpected error type for: %v"
-	uniqueConstraintViolation = "value with matching %s already exists (%s)"
+	uniqueConstraintViolation = "value with matching already exists (%s)"
 	defaultPgError            = "failed database operation with %s"
 	unsupportedTableOperation = "cannot query with specified table attributes: %s"
 )
@@ -35,13 +36,14 @@ func (p *postgresErrorTransformer) fromGormError(err error) error {
 }
 
 func (p *postgresErrorTransformer) ToDataCatalogError(err error) error {
-	pqError, ok := err.(*pq.Error)
+	cErr, ok := err.(ConnectError)
 	if !ok {
 		return p.fromGormError(err)
 	}
+	pqError := cErr.Unwrap().(*pgconn.PgError)
 	switch pqError.Code {
 	case uniqueConstraintViolationCode:
-		return errors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, pqError.Constraint, pqError.Message)
+		return errors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, pqError.Message)
 	case undefinedTable:
 		return errors.NewDataCatalogErrorf(codes.InvalidArgument, unsupportedTableOperation, pqError.Message)
 	default:
@@ -52,3 +54,8 @@ func (p *postgresErrorTransformer) ToDataCatalogError(err error) error {
 func NewPostgresErrorTransformer() ErrorTransformer {
 	return &postgresErrorTransformer{}
 }
+
+type ConnectError interface {
+	Unwrap() error
+	Error() string
+}
diff --git a/datacatalog/pkg/repositories/gormimpl/tag_test.go b/datacatalog/pkg/repositories/gormimpl/tag_test.go
index c7c9f3b154..ee33417a47 100644
--- a/datacatalog/pkg/repositories/gormimpl/tag_test.go
+++ b/datacatalog/pkg/repositories/gormimpl/tag_test.go
@@ -3,6 +3,8 @@ package gormimpl
 import (
 	"testing"
 
+	"github.com/jackc/pgconn"
+
 	"gorm.io/gorm"
 
 	"context"
@@ -20,7 +22,6 @@ import (
 	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/promutils"
 	"github.com/flyteorg/flytestdlib/promutils/labeled"
-	"github.com/lib/pq"
 	"google.golang.org/grpc/codes"
 )
 
@@ -28,8 +29,24 @@ func init() {
 	labeled.SetMetricKeys(contextutils.AppNameKey)
 }
 
+type pgError struct {
+	e   error
+	msg string
+}
+
+func (p *pgError) Error() string {
+	return p.msg
+}
+
+func (p *pgError) Unwrap() error {
+	return p.e
+}
+
 func getAlreadyExistsErr() error {
-	return &pq.Error{Code: "23505"}
+	return &pgError{
+		e:   &pgconn.PgError{Code: "23505"},
+		msg: "some error",
+	}
 }
 
 func getTestTag() models.Tag {