sql/postgres: disable locking capability in crdb mode (#1421)

ariga · Jan 24, 2023 · 93eb984 · 93eb984
1 parent 747609c
commit 93eb984
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 25 deletions.
diff --git a/sql/postgres/crdb.go b/sql/postgres/crdb.go
@@ -12,24 +12,22 @@ import (
 	"strings"
 
 	"ariga.io/atlas/sql/internal/sqlx"
+	"ariga.io/atlas/sql/migrate"
 	"ariga.io/atlas/sql/schema"
-	"ariga.io/atlas/sql/sqlclient"
 )
 
-func init() {
-	sqlclient.Register(
-		"cockroach",
-		sqlclient.DriverOpener(Open),
-		sqlclient.RegisterCodec(MarshalHCL, EvalHCL),
-		sqlclient.RegisterFlavours("crdb"),
-		sqlclient.RegisterURLParser(parser{}),
-	)
-}
-
-// crdbDiff implements the sqlx.DiffDriver for CockroachDB.
 type (
 	crdbDiff    struct{ diff }
 	crdbInspect struct{ inspect }
+	noLocker    interface {
+		migrate.Driver
+		migrate.Snapshoter
+		migrate.CleanChecker
+		schema.Normalizer
+	}
+	noLockDriver struct {
+		noLocker
+	}
 )
 
 var _ sqlx.DiffDriver = (*crdbDiff)(nil)
@@ -266,10 +264,11 @@ const (
 	TypeGeometry = "geometry"
 )
 
-// CockroachDB query for getting schema indexes.
-// Scanning constraints is disabled due to internal CockroachDB error.
-// (internal error: unexpected type *tree.DOidWrapper for key value)
-const crdbIndexesQuery = `
+const (
+	// CockroachDB query for getting schema indexes.
+	// Scanning constraints is disabled due to internal CockroachDB error.
+	// (internal error: unexpected type *tree.DOidWrapper for key value)
+	crdbIndexesQuery = `
 SELECT
 	t.relname AS table_name,
 	i.relname AS index_name,
@@ -301,7 +300,7 @@ ORDER BY
 	table_name, index_name, idx.ord
 `
 
-const crdbColumnsQuery = `
+	crdbColumnsQuery = `
 SELECT
 	t1.table_name,
 	t1.column_name,
@@ -341,3 +340,4 @@ WHERE
 ORDER BY
 	t1.table_name, t1.ordinal_position
 `
+)
diff --git a/sql/postgres/driver.go b/sql/postgres/driver.go
@@ -28,12 +28,13 @@ type (
 		schema.Differ
 		schema.Inspector
 		migrate.PlanApplier
-		schema string // the schema given in the `search_path` parameter (if given)
 	}
 
 	// database connection and its information.
 	conn struct {
 		schema.ExecQuerier
+		// The schema in the `search_path` parameter (if given).
+		schema string
 		// System variables that are set on `Open`.
 		collate string
 		ctype   string
@@ -69,7 +70,12 @@ func opener(_ context.Context, u *url.URL) (*sqlclient.Client, error) {
 		}
 		return nil, err
 	}
-	drv.(*Driver).schema = ur.Schema
+	switch drv := drv.(type) {
+	case *Driver:
+		drv.schema = ur.Schema
+	case noLockDriver:
+		drv.noLocker.(*Driver).schema = ur.Schema
+	}
 	return &sqlclient.Client{
 		Name:   DriverName,
 		DB:     db,
@@ -101,11 +107,13 @@ func Open(db schema.ExecQuerier) (migrate.Driver, error) {
 	}
 	// Means we are connected to CockroachDB because we have a result for name='crdb_version'. see `paramsQuery`.
 	if c.crdb = len(params) == 4; c.crdb {
-		return &Driver{
-			conn:        c,
-			Differ:      &sqlx.Diff{DiffDriver: &crdbDiff{diff{c}}},
-			Inspector:   &crdbInspect{inspect{c}},
-			PlanApplier: &planApply{c},
+		return noLockDriver{
+			&Driver{
+				conn:        c,
+				Differ:      &sqlx.Diff{DiffDriver: &crdbDiff{diff{c}}},
+				Inspector:   &crdbInspect{inspect{c}},
+				PlanApplier: &planApply{c},
+			},
 		}, nil
 	}
 	return &Driver{

diff --git a/sql/postgres/driver_test.go b/sql/postgres/driver_test.go
@@ -105,7 +105,7 @@ func TestDriver_UnlockError(t *testing.T) {
 
 func TestDriver_CheckClean(t *testing.T) {
 	s := schema.New("test")
-	drv := &Driver{Inspector: &mockInspector{schema: s}, schema: "test"}
+	drv := &Driver{Inspector: &mockInspector{schema: s}, conn: conn{schema: "test"}}
 	// Empty schema.
 	err := drv.CheckClean(context.Background(), nil)
 	require.NoError(t, err)