Repository: ariga/atlas
Pull request #1885 (2023-07-21 05:18:56)

sql/mysql/internal/mysqlversion: fetch collation/collate defaults from information_schema if possible


Patch:
diff --git a/sql/mysql/diff.go b/sql/mysql/diff.go
index 868fcbf5..3368e5ca 100644
--- a/sql/mysql/diff.go
+++ b/sql/mysql/diff.go
@@ -644,7 +644,7 @@ func (d *diff) defaultCollate(attrs *[]schema.Attr) error {
 		return nil
 	}
 	d.ch2co.Do(func() {
-		d.ch2co.v, d.ch2co.err = d.CharsetToCollate()
+		d.ch2co.v, d.ch2co.err = d.CharsetToCollate(d.ExecQuerier)
 	})
 	if d.ch2co.err != nil {
 		return d.ch2co.err
@@ -664,7 +664,7 @@ func (d *diff) defaultCharset(attrs *[]schema.Attr) error {
 		return nil
 	}
 	d.co2ch.Do(func() {
-		d.co2ch.v, d.co2ch.err = d.CollateToCharset()
+		d.co2ch.v, d.co2ch.err = d.CollateToCharset(d.ExecQuerier)
 	})
 	if d.co2ch.err != nil {
 		return d.co2ch.err
diff --git a/sql/mysql/internal/mysqlversion/mysqlversion.go b/sql/mysql/internal/mysqlversion/mysqlversion.go
index 2eb4582a..70a742ae 100644
--- a/sql/mysql/internal/mysqlversion/mysqlversion.go
+++ b/sql/mysql/internal/mysqlversion/mysqlversion.go
@@ -5,11 +5,15 @@
 package mysqlversion
 
 import (
+	"context"
+	"database/sql"
 	"embed"
 	"encoding/json"
 	"fmt"
 	"strings"
 
+	"ariga.io/atlas/sql/internal/sqlx"
+	"ariga.io/atlas/sql/schema"
 	"golang.org/x/mod/semver"
 )
 
@@ -91,21 +95,53 @@ func (v V) SupportsViewUsage() bool {
 }
 
 // CharsetToCollate returns the mapping from charset to its default collation.
-func (v V) CharsetToCollate() (map[string]string, error) {
+func (v V) CharsetToCollate(conn schema.ExecQuerier) (map[string]string, error) {
 	name := "is/charset2collate"
 	if v.Maria() {
 		name += ".maria"
 	}
-	return decode(name)
+	c2c, err := decode(name)
+	if err != nil {
+		return nil, err
+	}
+	if conn != nil {
+		mayExtend(conn, "SELECT CHARACTER_SET_NAME, DEFAULT_COLLATE_NAME FROM INFORMATION_SCHEMA.CHARACTER_SETS", c2c)
+	}
+	return c2c, nil
 }
 
 // CollateToCharset returns the mapping from a collation to its charset.
-func (v V) CollateToCharset() (map[string]string, error) {
+func (v V) CollateToCharset(conn schema.ExecQuerier) (map[string]string, error) {
 	name := "is/collate2charset"
 	if v.Maria() {
 		name += ".maria"
 	}
-	return decode(name)
+	c2c, err := decode(name)
+	if err != nil {
+		return nil, err
+	}
+	if conn != nil {
+		mayExtend(conn, "SELECT COLLATION_NAME, CHARACTER_SET_NAME FROM INFORMATION_SCHEMA.COLLATIONS", c2c)
+	}
+	return c2c, nil
+}
+
+// mayExtend the given collation/charset map from information schema.
+func mayExtend(conn schema.ExecQuerier, query string, to map[string]string) {
+	rows, err := conn.QueryContext(context.Background(), query)
+	if err != nil {
+		return
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var c1, c2 sql.NullString
+		if err := rows.Scan(&c1, &c2); err != nil {
+			return
+		}
+		if sqlx.ValidString(c1) && sqlx.ValidString(c2) {
+			to[c1.String] = c2.String
+		}
+	}
 }
 
 // Maria reports if the MySQL version is MariaDB.


Test patch:
diff --git a/sql/mysql/internal/mysqlversion/mysqlversion_test.go b/sql/mysql/internal/mysqlversion/mysqlversion_test.go
index 2001bd69..1d3536d9 100644
--- a/sql/mysql/internal/mysqlversion/mysqlversion_test.go
+++ b/sql/mysql/internal/mysqlversion/mysqlversion_test.go
@@ -8,6 +8,8 @@ import (
 	"testing"
 
 	"ariga.io/atlas/sql/mysql/internal/mysqlversion"
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/stretchr/testify/require"
 )
 
 func TestV_SupportsGeneratedColumns(t *testing.T) {
@@ -32,3 +34,18 @@ func TestV_SupportsGeneratedColumns(t *testing.T) {
 		})
 	}
 }
+
+func TestV_CollateToCharset(t *testing.T) {
+	c2c, err := mysqlversion.V("8.0.0").CollateToCharset(nil)
+	require.NoError(t, err)
+	require.Equal(t, "utf8mb4", c2c["utf8mb4_general_ci"])
+	require.Empty(t, c2c["custom"])
+
+	db, mk, err := sqlmock.New()
+	require.NoError(t, err)
+	mk.ExpectQuery("SELECT COLLATION_NAME, CHARACTER_SET_NAME FROM INFORMATION_SCHEMA.COLLATIONS").
+		WillReturnRows(sqlmock.NewRows([]string{"COLLATION_NAME", "CHARACTER_SET_NAME"}).AddRow("custom", "unknown"))
+	c2c, err = mysqlversion.V("8.0.0").CollateToCharset(db)
+	require.NoError(t, err)
+	require.Equal(t, "unknown", c2c["custom"])
+}