Fix connection error handling (#45)

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

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

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
+}

datacatalog/pkg/repositories/gormimpl/tag_test.go

@@ -3,6 +3,8 @@ package gormimpl
 import (
 	"testing"
 
+	"github.com/jackc/pgconn"
+
 	"gorm.io/gorm"
 
 	"context"
@@ -20,16 +22,31 @@ 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"
 )
 
 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 {