From 8fe088b47983f01f514c49a53f55d956f633ffab Mon Sep 17 00:00:00 2001
From: Aneesh Makala <amakala@yugabyte.com>
Date: Tue, 10 Dec 2024 16:16:03 +0530
Subject: [PATCH] fix unlogged table issue

---
 .../src/query/queryissue/issues_ddl_test.go   | 46 ++++++++++++++++++-
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/yb-voyager/src/query/queryissue/issues_ddl_test.go b/yb-voyager/src/query/queryissue/issues_ddl_test.go
index 105f56587..59842b7d8 100644
--- a/yb-voyager/src/query/queryissue/issues_ddl_test.go
+++ b/yb-voyager/src/query/queryissue/issues_ddl_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
 	"github.com/stretchr/testify/assert"
 	"github.com/testcontainers/testcontainers-go/modules/yugabytedb"
 )
@@ -52,6 +53,32 @@ func getConn() (*pgx.Conn, error) {
 	return conn, nil
 }
 
+func getConnWithNoticeHandler(noticeHandler func(*pgconn.PgConn, *pgconn.Notice)) (*pgx.Conn, error) {
+	ctx := context.Background()
+	var connStr string
+	var err error
+	if yugabytedbConnStr != "" {
+		connStr = yugabytedbConnStr
+	} else {
+		connStr, err = yugabytedbContainer.YSQLConnectionString(ctx, "sslmode=disable")
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	conf, err := pgx.ParseConfig(connStr)
+	if err != nil {
+		return nil, err
+	}
+	conf.OnNotice = noticeHandler
+	conn, err := pgx.ConnectConfig(ctx, conf)
+	if err != nil {
+		return nil, err
+	}
+
+	return conn, nil
+}
+
 func testXMLFunctionIssue(t *testing.T) {
 	ctx := context.Background()
 	conn, err := getConn()
@@ -79,13 +106,28 @@ func testStoredGeneratedFunctionsIssue(t *testing.T) {
 }
 
 func testUnloggedTableIssue(t *testing.T) {
+	noticeFound := false
+	noticeHandler := func(conn *pgconn.PgConn, notice *pgconn.Notice) {
+		if notice != nil {
+			if notice.Message != "" {
+				assert.Equal(t, "unlogged option is currently ignored in YugabyteDB, all non-temp tables will be logged", notice.Message)
+				noticeFound = true
+			}
+		}
+	}
 	ctx := context.Background()
-	conn, err := getConn()
+	conn, err := getConnWithNoticeHandler(noticeHandler)
 	assert.NoError(t, err)
 
 	defer conn.Close(context.Background())
 	_, err = conn.Exec(ctx, "CREATE UNLOGGED TABLE unlogged_table (a int)")
-	assert.ErrorContains(t, err, "UNLOGGED database object not supported yet")
+	// in 2024.2, UNLOGGED no longer throws an error, just a notice
+	if noticeFound {
+		return
+	} else {
+		assert.ErrorContains(t, err, "UNLOGGED database object not supported yet")
+	}
+
 }
 
 func TestDDLIssuesInYBVersion(t *testing.T) {