db2: add db2i:// scheme for DB2 for i (IBM i / AS-400)

pkg/database/database.go

@@ -456,7 +456,9 @@ func Connect(ctx context.Context, opts ConnectOptions) (*sql.DB, DbEngine, error
 		}
 		return db, Vertica, nil
 
-	case "db2":
+	case "db2", "db2i":
+		// "db2" = DB2 LUW; "db2i" = DB2 for i (IBM i / AS-400). Same engine/driver;
+		// they differ only in the default DRDA port (see db2.convertToDB2DSN).
 		db, err := db2.Connect(ctx, parsedDsn.String())
 		if err != nil {
 			return nil, Unknown, err

pkg/database/db2/dsn.go

@@ -9,11 +9,12 @@ import (
 
 // Keywords derived from the URL itself; query parameters may not override them.
 // Anyone needing full control over these can pass a native DB2 DSN instead.
+// PROTOCOL is intentionally NOT reserved: it defaults to TCPIP but may be
+// overridden via a ?PROTOCOL= query parameter (needed by some DB2-for-i paths).
 var reservedDSNKeywords = map[string]bool{
 	"HOSTNAME": true,
 	"DATABASE": true,
 	"PORT":     true,
-	"PROTOCOL": true,
 	"UID":      true,
 	"PWD":      true,
 }
@@ -35,7 +36,7 @@ func quoteDB2Value(v string) (string, error) {
 func convertToDB2DSN(dsn string) (string, error) {
 	// If it's already in DB2 format (contains HOSTNAME= or DATABASE=), return as-is.
 	// URL-format DSNs are exempt from this check so those markers may appear in credentials.
-	if !strings.HasPrefix(dsn, "db2://") && (strings.Contains(dsn, "HOSTNAME=") || strings.Contains(dsn, "DATABASE=")) {
+	if !strings.HasPrefix(dsn, "db2://") && !strings.HasPrefix(dsn, "db2i://") && (strings.Contains(dsn, "HOSTNAME=") || strings.Contains(dsn, "DATABASE=")) {
 		return dsn, nil
 	}
 
@@ -44,17 +45,25 @@ func convertToDB2DSN(dsn string) (string, error) {
 		return "", fmt.Errorf("invalid DSN format: %w", err)
 	}
 
-	if parsedURL.Scheme != "db2" {
-		return "", fmt.Errorf("expected db2:// scheme, got %s", parsedURL.Scheme)
+	// "db2" targets DB2 LUW; "db2i" targets DB2 for i (IBM i / AS-400), which
+	// differs only in its default DRDA port. All other handling is identical.
+	if parsedURL.Scheme != "db2" && parsedURL.Scheme != "db2i" {
+		return "", fmt.Errorf("expected db2:// or db2i:// scheme, got %s", parsedURL.Scheme)
 	}
+	isIBMi := parsedURL.Scheme == "db2i"
 
 	hostname, err := quoteDB2Value(parsedURL.Hostname())
 	if err != nil {
 		return "", fmt.Errorf("invalid hostname: %w", err)
 	}
 	port := parsedURL.Port()
 	if port == "" {
-		port = "50000" // Default DB2 port
+		// DB2 LUW defaults to 50000; DB2 for i (DRDA) defaults to 446.
+		if isIBMi {
+			port = "446"
+		} else {
+			port = "50000"
+		}
 	}
 
 	database := strings.TrimPrefix(parsedURL.Path, "/")
@@ -78,7 +87,19 @@ func convertToDB2DSN(dsn string) (string, error) {
 	dsnParts = append(dsnParts, fmt.Sprintf("HOSTNAME=%s", hostname))
 	dsnParts = append(dsnParts, fmt.Sprintf("DATABASE=%s", quotedDatabase))
 	dsnParts = append(dsnParts, fmt.Sprintf("PORT=%s", port))
-	dsnParts = append(dsnParts, "PROTOCOL=TCPIP")
+	// PROTOCOL defaults to TCPIP but is overridable via a ?PROTOCOL= query param
+	// (appended, uppercased, by the query-parameter loop below). Only set the
+	// default here when the caller did not supply one.
+	hasProtocol := false
+	for k := range parsedURL.Query() {
+		if strings.EqualFold(k, "PROTOCOL") {
+			hasProtocol = true
+			break
+		}
+	}
+	if !hasProtocol {
+		dsnParts = append(dsnParts, "PROTOCOL=TCPIP")
+	}
 
 	if username != "" {
 		quoted, err := quoteDB2Value(username)

pkg/database/db2/dsn_test.go

@@ -28,10 +28,25 @@ func TestConvertToDB2DSN(t *testing.T) {
 			dsn:  "HOSTNAME=dbhost;PORT=50000;DATABASE=testdb;UID=user;PWD=pass",
 			want: "HOSTNAME=dbhost;PORT=50000;DATABASE=testdb;UID=user;PWD=pass",
 		},
+		{
+			name: "db2i scheme without port uses DRDA default 446",
+			dsn:  "db2i://user:pass@ibmi/MYRDB",
+			want: "HOSTNAME=ibmi;DATABASE=MYRDB;PORT=446;PROTOCOL=TCPIP;UID=user;PWD=pass",
+		},
+		{
+			name: "db2i scheme with explicit port",
+			dsn:  "db2i://user:pass@ibmi:8471/MYRDB",
+			want: "HOSTNAME=ibmi;DATABASE=MYRDB;PORT=8471;PROTOCOL=TCPIP;UID=user;PWD=pass",
+		},
+		{
+			name: "protocol overridable via query param (no duplicate default)",
+			dsn:  "db2://user:pass@dbhost:50000/testdb?PROTOCOL=IPC",
+			want: "HOSTNAME=dbhost;DATABASE=testdb;PORT=50000;UID=user;PWD=pass;PROTOCOL=IPC",
+		},
 		{
 			name:    "wrong scheme",
 			dsn:     "postgres://dbhost/testdb",
-			wantErr: "expected db2:// scheme",
+			wantErr: "expected db2:// or db2i:// scheme",
 		},
 		{
 			name:    "missing database name",