More pgconn handling fixes (#4420)

Signed-off-by: Yee Hing Tong <[email protected]>
flyteorg · Nov 14, 2023 · 63b577e · 63b577e
1 parent 05b4ea0
commit 63b577e
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 9 deletions.
diff --git a/datacatalog/pkg/repositories/errors/postgres.go b/datacatalog/pkg/repositories/errors/postgres.go
@@ -3,6 +3,7 @@ package errors
 import (
 	"errors"
 	"fmt"
+	"github.com/flyteorg/flyte/flytestdlib/database"
 	"reflect"
 
 	"github.com/jackc/pgconn"
@@ -39,6 +40,11 @@ func (p *postgresErrorTransformer) fromGormError(err error) error {
 }
 
 func (p *postgresErrorTransformer) ToDataCatalogError(err error) error {
+	// First try the stdlib error handling
+	if database.IsPgErrorWithCode(err, uniqueConstraintViolationCode) {
+		return catalogErrors.NewDataCatalogErrorf(codes.AlreadyExists, uniqueConstraintViolation, err.Error())
+	}
+
 	if unwrappedErr := errors.Unwrap(err); unwrappedErr != nil {
 		err = unwrappedErr
 	}

diff --git a/flyteadmin/pkg/async/cloudevent/implementations/cloudevent_publisher.go b/flyteadmin/pkg/async/cloudevent/implementations/cloudevent_publisher.go
@@ -288,10 +288,6 @@ func (c *CloudEventWrappedPublisher) TransformTaskExecutionEvent(ctx context.Con
 			return nil, err
 		}
 	}
-	if outputs == nil {
-		// todo: remove
-		fmt.Printf("No output data found for task execution %v\n", rawEvent)
-	}
 
 	return &event.CloudEventTaskExecution{
 		RawEvent:            rawEvent,

diff --git a/flytestdlib/database/postgres.go b/flytestdlib/database/postgres.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"github.com/flyteorg/flyte/flytestdlib/logger"
+	oldPgConn "github.com/jackc/pgconn"
 	"github.com/jackc/pgx/v5/pgconn"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
@@ -92,14 +93,21 @@ func CreatePostgresDbIfNotExists(ctx context.Context, gormConfig *gorm.Config, p
 }
 
 func IsPgErrorWithCode(err error, code string) bool {
-	// Make sure the pgconn you're using is "github.com/jackc/pgx/v5/pgconn"
+	// Newer versions of the gorm postgres driver seem to use
+	// "github.com/jackc/pgx/v5/pgconn"
 	// See https://github.com/go-gorm/gorm/issues/4135
+	// Let's just try both of them to make sure.
 	pgErr := &pgconn.PgError{}
-	if !errors.As(err, &pgErr) {
+	if errors.As(err, &pgErr) {
 		// err chain does not contain a pgconn.PgError
-		return false
+		return pgErr.Code == code
 	}
 
-	// pgconn.PgError found in chain and set to code specified
-	return pgErr.Code == code
+	oldPgErr := &oldPgConn.PgError{}
+	if errors.As(err, &oldPgErr) {
+		// err chain does not contain a pgconn.PgError
+		return oldPgErr.Code == code
+	}
+
+	return false
 }