Repository: ariga/atlas
Pull request #3325 (2025-01-19 19:13:27)

sql/postgres: load default access method from settings


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 52d8f365..ff381ce8 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -83,11 +83,11 @@ func ScanStrings(rows *sql.Rows) ([]string, error) {
 	defer rows.Close()
 	var vs []string
 	for rows.Next() {
-		var v string
+		var v sql.NullString
 		if err := rows.Scan(&v); err != nil {
 			return nil, err
 		}
-		vs = append(vs, v)
+		vs = append(vs, v.String)
 	}
 	return vs, nil
 }
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
index c03e6beb..56511200 100644
--- a/sql/postgres/driver.go
+++ b/sql/postgres/driver.go
@@ -36,6 +36,8 @@ type (
 		schema.ExecQuerier
 		// The schema in the `search_path` parameter (if given).
 		schema string
+		// Maps to the connection default_table_access_method parameter.
+		accessMethod string
 		// System variables that are set on `Open`.
 		version int
 		crdb    bool
@@ -97,21 +99,18 @@ func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 	if err != nil {
 		return nil, fmt.Errorf("postgres: scanning system variables: %w", err)
 	}
-	params, err := sqlx.ScanStrings(rows)
-	if err != nil {
-		return nil, fmt.Errorf("postgres: failed scanning rows: %w", err)
-	}
-	if len(params) != 1 && len(params) != 2 {
-		return nil, fmt.Errorf("postgres: unexpected number of rows: %d", len(params))
+	var ver, am, crdb sql.NullString
+	if err := sqlx.ScanOne(rows, &ver, &am, &crdb); err != nil {
+		return nil, fmt.Errorf("postgres: scanning system variables: %w", err)
 	}
-	if c.version, err = strconv.Atoi(params[0]); err != nil {
-		return nil, fmt.Errorf("postgres: malformed version: %s: %w", params[0], err)
+	if c.version, err = strconv.Atoi(ver.String); err != nil {
+		return nil, fmt.Errorf("postgres: malformed version: %s: %w", ver.String, err)
 	}
 	if c.version < 10_00_00 {
 		return nil, fmt.Errorf("postgres: unsupported postgres version: %d", c.version)
 	}
-	// Means we are connected to CockroachDB because we have a result for name='crdb_version'. see `paramsQuery`.
-	if c.crdb = len(params) == 2; c.crdb {
+	c.accessMethod = am.String
+	if c.crdb = sqlx.ValidString(crdb); c.crdb {
 		return noLockDriver{
 			&Driver{
 				conn:        c,
diff --git a/sql/postgres/inspect_oss.go b/sql/postgres/inspect_oss.go
index 25ba5c95..b26e4fa1 100644
--- a/sql/postgres/inspect_oss.go
+++ b/sql/postgres/inspect_oss.go
@@ -1500,7 +1500,7 @@ func parseFmtType(t string) (s, n string) {
 
 const (
 	// Query to list runtime parameters.
-	paramsQuery = `SELECT setting FROM pg_settings WHERE name IN ('server_version_num', 'crdb_version') ORDER BY name DESC`
+	paramsQuery = `SELECT current_setting('server_version_num'), current_setting('default_table_access_method', true), current_setting('crdb_version', true)`
 
 	// Query to list database schemas.
 	schemasQuery = `


Test patch:
diff --git a/sql/postgres/inspect_oss_test.go b/sql/postgres/inspect_oss_test.go
index b0c4fd4c..29e39bc1 100644
--- a/sql/postgres/inspect_oss_test.go
+++ b/sql/postgres/inspect_oss_test.go
@@ -440,11 +440,10 @@ func TestDriver_InspectCRDBSchema(t *testing.T) {
 	mk := mock{m}
 	mk.ExpectQuery(sqltest.Escape(paramsQuery)).
 		WillReturnRows(sqltest.Rows(`
-	setting
-------------
-130000
-cockroach
-				`))
+  version       |  am  | crdb
+----------------|------|-----
+ 130000         | heap | cockroach
+`))
 	drv, err := Open(db)
 	require.NoError(t, err)
 	mk.ExpectQuery(sqltest.Escape(fmt.Sprintf(schemasQueryArgs, "= $1"))).
@@ -727,14 +726,14 @@ type mock struct {
 func (m mock) version(version string) {
 	m.ExpectQuery(sqltest.Escape(paramsQuery)).
 		WillReturnRows(sqltest.Rows(`
-  setting
-------------
- ` + version + `
+  setting       |  am  | crdb
+----------------|------|-----
+ ` + version + `| heap | NULL
 `))
 }
 
 func (m mock) tableExists(schema, table string, exists bool) {
-	rows := sqlmock.NewRows([]string{"oid", "table_schema", "table_name", "table_comment", "partition_attrs", "partition_strategy", "partition_exprs", "attrs"})
+	rows := sqlmock.NewRows([]string{"oid", "table_schema", "table_name", "table_comment", "partition_attrs", "partition_strategy", "partition_exprs", "row_security"})
 	if exists {
 		rows.AddRow(nil, schema, table, nil, nil, nil, nil, nil)
 	}