Skip to content

feat: add support for the YugabyteDB PGX smart driver #4345

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sumanpatrik wants to merge 6 commits into
base: main
Choose a base branch
from
+48 −17
Open

feat: add support for the YugabyteDB PGX smart driver #4345

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
4b4596e
feat: add support for yugabyte smart driver
Mar 18, 2026
46428da
refactor parseDriver so that we can reuse the PGXV5
Mar 19, 2026
ece0a7e
refactor
Mar 19, 2026
520e301
docs: clarify the function need
Mar 19, 2026
b24def7
go fmt
Mar 19, 2026
ebf9ad7
refactor: proper sql_package
Mar 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions internal/codegen/golang/driver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,21 @@ 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
}
}

// custom packages based on pgx/v5 (e.g., YugabyteDB smart drivers) should return
// only the original driver when determining columnType, refer to the postgresType func.
func parseDriverPGType(sqlPackage string) opts.SQLDriver {
switch sqlPackage {
case opts.SQLPackagePGXV4:
return opts.SQLDriverPGXV4
case opts.SQLPackagePGXV5, opts.SQLPackageYugaBytePGXV5:
return opts.SQLDriverPGXV5
default:
return opts.SQLDriverLibPQ
}
Expand Down
10 changes: 10 additions & 0 deletions internal/codegen/golang/imports.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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{}{}
}
Expand All @@ -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{}{}
}
Expand Down Expand Up @@ -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)
Expand Down
20 changes: 13 additions & 7 deletions internal/codegen/golang/opts/enum.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "yugabyte/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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -58,6 +62,8 @@ func (d SQLDriver) Package() string {
return SQLPackagePGXV4
case SQLDriverPGXV5:
return SQLPackagePGXV5
case SQLDriverYugaBytePGXV5:
return SQLPackageYugaBytePGXV5
default:
return SQLPackageStandard
}
Expand Down
2 changes: 1 addition & 1 deletion internal/codegen/golang/postgresql_type.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ func parseIdentifierString(name string) (*plugin.Identifier, error) {
func postgresType(req *plugin.GenerateRequest, options *opts.Options, col *plugin.Column) string {
columnType := sdk.DataType(col.Type)
notNull := col.NotNull || col.IsArray
driver := parseDriver(options.SqlPackage)
driver := parseDriverPGType(options.SqlPackage)
emitPointersForNull := driver.IsPGX() && options.EmitPointersForNullTypes

switch columnType {
Expand Down
4 changes: 2 additions & 2 deletions internal/config/config.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,8 +125,8 @@ type SQL struct {
// AnalyzerDatabase represents the database analyzer setting.
// It can be a boolean (true/false) or the string "only" for database-only mode.
type AnalyzerDatabase struct {
value *bool // nil means not set, true/false for boolean values
isOnly bool // true when set to "only"
value *bool // nil means not set, true/false for boolean values
isOnly bool // true when set to "only"
}

// IsEnabled returns true if the database analyzer should be used.
Expand Down
4 changes: 2 additions & 2 deletions internal/engine/clickhouse/convert.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -485,8 +485,8 @@ func (c *cc) convertFunctionCall(n *chast.FunctionCall) *ast.FuncCall {
Funcname: &ast.List{
Items: []ast.Node{&ast.String{Str: n.Name}},
},
Location: n.Pos().Offset,
AggDistinct: n.Distinct,
Location: n.Pos().Offset,
AggDistinct: n.Distinct,
}

// Convert arguments
Expand Down
6 changes: 3 additions & 3 deletions internal/engine/clickhouse/parse.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,8 @@ func (p *Parser) Parse(r io.Reader) ([]ast.Statement, error) {
// https://clickhouse.com/docs/en/sql-reference/syntax#comments
func (p *Parser) CommentSyntax() source.CommentSyntax {
return source.CommentSyntax{
Dash: true, // -- comment
SlashStar: true, // /* comment */
Hash: true, // # comment (ClickHouse supports this)
Dash: true, // -- comment
SlashStar: true, // /* comment */
Hash: true, // # comment (ClickHouse supports this)
}
}
4 changes: 2 additions & 2 deletions internal/engine/sqlite/convert.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -826,7 +826,7 @@ func (c *cc) convertUnaryExpr(n *parser.Expr_unaryContext) ast.Node {
if opCtx.MINUS() != nil {
// Negative number: -expr
return &ast.A_Expr{
Name: &ast.List{Items: []ast.Node{&ast.String{Str: "-"}}},
Name: &ast.List{Items: []ast.Node{&ast.String{Str: "-"}}},
Rexpr: expr,
}
}
Expand All @@ -837,7 +837,7 @@ func (c *cc) convertUnaryExpr(n *parser.Expr_unaryContext) ast.Node {
if opCtx.TILDE() != nil {
// Bitwise NOT: ~expr
return &ast.A_Expr{
Name: &ast.List{Items: []ast.Node{&ast.String{Str: "~"}}},
Name: &ast.List{Items: []ast.Node{&ast.String{Str: "~"}}},
Rexpr: expr,
}
}
Expand Down
Loading