feat: add support for yugabyte smart driver

Author: Suman De

internal/codegen/golang/driver.go

@@ -8,6 +8,8 @@ func parseDriver(sqlPackage string) opts.SQLDriver {
 		return opts.SQLDriverPGXV4
 	case opts.SQLPackagePGXV5:
 		return opts.SQLDriverPGXV5
+	case opts.SQLPackageYugaBytePGXV5:
+		return opts.SQLDriverYugaBytePGXV5
 	default:
 		return opts.SQLDriverLibPQ
 	}

internal/codegen/golang/imports.go

@@ -132,6 +132,10 @@ func (i *importer) dbImports() fileImports {
 	case opts.SQLDriverPGXV5:
 		pkg = append(pkg, ImportSpec{Path: "github.com/jackc/pgx/v5/pgconn"})
 		pkg = append(pkg, ImportSpec{Path: "github.com/jackc/pgx/v5"})
+
+	case opts.SQLDriverYugaBytePGXV5:
+		pkg = append(pkg, ImportSpec{Path: "github.com/yugabyte/pgx/v5/pgconn"})
+		pkg = append(pkg, ImportSpec{Path: "github.com/yugabyte/pgx/v5"})
 	default:
 		std = append(std, ImportSpec{Path: "database/sql"})
 		if i.Options.EmitPreparedQueries {
@@ -176,6 +180,8 @@ func buildImports(options *opts.Options, queries []Query, uses func(string) bool
 				pkg[ImportSpec{Path: "github.com/jackc/pgconn"}] = struct{}{}
 			case opts.SQLDriverPGXV5:
 				pkg[ImportSpec{Path: "github.com/jackc/pgx/v5/pgconn"}] = struct{}{}
+			case opts.SQLDriverYugaBytePGXV5:
+				pkg[ImportSpec{Path: "github.com/yugabyte/pgx/v5/pgconn"}] = struct{}{}
 			default:
 				std["database/sql"] = struct{}{}
 			}
@@ -191,6 +197,8 @@ func buildImports(options *opts.Options, queries []Query, uses func(string) bool
 	if uses("pgtype.") {
 		if sqlpkg == opts.SQLDriverPGXV5 {
 			pkg[ImportSpec{Path: "github.com/jackc/pgx/v5/pgtype"}] = struct{}{}
+		} else if sqlpkg == opts.SQLDriverYugaBytePGXV5 {
+			pkg[ImportSpec{Path: "github.com/yugabyte/pgx/v5/pgtype"}] = struct{}{}
 		} else {
 			pkg[ImportSpec{Path: "github.com/jackc/pgtype"}] = struct{}{}
 		}
@@ -489,6 +497,8 @@ func (i *importer) batchImports() fileImports {
 		pkg[ImportSpec{Path: "github.com/jackc/pgx/v4"}] = struct{}{}
 	case opts.SQLDriverPGXV5:
 		pkg[ImportSpec{Path: "github.com/jackc/pgx/v5"}] = struct{}{}
+	case opts.SQLDriverYugaBytePGXV5:
+		pkg[ImportSpec{Path: "github.com/yugabyte/pgx/v5"}] = struct{}{}
 	}
 
 	return sortedImports(std, pkg)

internal/codegen/golang/opts/enum.go

@@ -5,15 +5,17 @@ import "fmt"
 type SQLDriver string
 
 const (
-	SQLPackagePGXV4    string = "pgx/v4"
-	SQLPackagePGXV5    string = "pgx/v5"
-	SQLPackageStandard string = "database/sql"
+	SQLPackagePGXV4         string = "pgx/v4"
+	SQLPackagePGXV5         string = "pgx/v5"
+	SQLPackageStandard      string = "database/sql"
+	SQLPackageYugaBytePGXV5 string = "yb/pgx/v5"
 )
 
 var validPackages = map[string]struct{}{
-	string(SQLPackagePGXV4):    {},
-	string(SQLPackagePGXV5):    {},
-	string(SQLPackageStandard): {},
+	string(SQLPackagePGXV4):         {},
+	string(SQLPackagePGXV5):         {},
+	string(SQLPackageYugaBytePGXV5): {},
+	string(SQLPackageStandard):      {},
 }
 
 func validatePackage(sqlPackage string) error {
@@ -28,13 +30,15 @@ const (
 	SQLDriverPGXV5                      = "github.com/jackc/pgx/v5"
 	SQLDriverLibPQ                      = "github.com/lib/pq"
 	SQLDriverGoSQLDriverMySQL           = "github.com/go-sql-driver/mysql"
+	SQLDriverYugaBytePGXV5              = "github.com/yugabyte/pgx/v5"
 )
 
 var validDrivers = map[string]struct{}{
 	string(SQLDriverPGXV4):            {},
 	string(SQLDriverPGXV5):            {},
 	string(SQLDriverLibPQ):            {},
 	string(SQLDriverGoSQLDriverMySQL): {},
+	string(SQLDriverYugaBytePGXV5):    {},
 }
 
 func validateDriver(sqlDriver string) error {
@@ -45,7 +49,7 @@ func validateDriver(sqlDriver string) error {
 }
 
 func (d SQLDriver) IsPGX() bool {
-	return d == SQLDriverPGXV4 || d == SQLDriverPGXV5
+	return d == SQLDriverPGXV4 || d == SQLDriverPGXV5 || d == SQLDriverYugaBytePGXV5
 }
 
 func (d SQLDriver) IsGoSQLDriverMySQL() bool {
@@ -58,6 +62,8 @@ func (d SQLDriver) Package() string {
 		return SQLPackagePGXV4
 	case SQLDriverPGXV5:
 		return SQLPackagePGXV5
+	case SQLDriverYugaBytePGXV5:
+		return SQLPackageYugaBytePGXV5
 	default:
 		return SQLPackageStandard
 	}