feat: create a common function to be used in the operator

Signed-off-by: Manan Gupta <[email protected]>
planetscale · Dec 28, 2023 · 637d481 · 637d481
1 parent dc4976b
commit 637d481
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 47 deletions.
diff --git a/pkg/controller/vitesskeyspace/reconcile_handler.go b/pkg/controller/vitesskeyspace/reconcile_handler.go
@@ -24,14 +24,12 @@ import (
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/client-go/tools/record"
 	"sigs.k8s.io/controller-runtime/pkg/client"
-	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/vt/logutil"
-	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/vttablet/tmclient"
 	"vitess.io/vitess/go/vt/wrangler"
 
 	v2 "planetscale.dev/vitess-operator/pkg/apis/planetscale/v2"
+	"planetscale.dev/vitess-operator/pkg/operator/environment"
 	"planetscale.dev/vitess-operator/pkg/operator/reconciler"
 	"planetscale.dev/vitess-operator/pkg/operator/toposerver"
 )
@@ -77,12 +75,7 @@ func (r *reconcileHandler) tsInit(ctx context.Context) error {
 		r.tmc = tmclient.NewTabletManagerClient()
 	}
 
-	collationEnv := collations.NewEnvironment(servenv.MySQLServerVersion())
-	parser, err := sqlparser.New(sqlparser.Options{
-		MySQLServerVersion: servenv.MySQLServerVersion(),
-		TruncateUILen:      servenv.TruncateUILen,
-		TruncateErrLen:     servenv.TruncateErrLen,
-	})
+	collationEnv, parser, err := environment.GetCollationEnvAndParser()
 	if err != nil {
 		return err
 	}

diff --git a/pkg/controller/vitessshard/reconcile_topo.go b/pkg/controller/vitessshard/reconcile_topo.go
@@ -21,9 +21,6 @@ import (
 	"strings"
 	"time"
 
-	"vitess.io/vitess/go/mysql/collations"
-	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 
 	corev1 "k8s.io/api/core/v1"
@@ -34,6 +31,7 @@ import (
 	"vitess.io/vitess/go/vt/wrangler"
 
 	planetscalev2 "planetscale.dev/vitess-operator/pkg/apis/planetscale/v2"
+	"planetscale.dev/vitess-operator/pkg/operator/environment"
 	"planetscale.dev/vitess-operator/pkg/operator/k8s"
 	"planetscale.dev/vitess-operator/pkg/operator/results"
 	"planetscale.dev/vitess-operator/pkg/operator/toposerver"
@@ -66,12 +64,7 @@ func (r *ReconcileVitessShard) reconcileTopology(ctx context.Context, vts *plane
 		return resultBuilder.RequeueAfter(topoRequeueDelay)
 	}
 	defer ts.Close()
-	collationEnv := collations.NewEnvironment(servenv.MySQLServerVersion())
-	parser, err := sqlparser.New(sqlparser.Options{
-		MySQLServerVersion: servenv.MySQLServerVersion(),
-		TruncateUILen:      servenv.TruncateUILen,
-		TruncateErrLen:     servenv.TruncateErrLen,
-	})
+	collationEnv, parser, err := environment.GetCollationEnvAndParser()
 	if err != nil {
 		return resultBuilder.Error(err)
 	}

diff --git a/pkg/controller/vitessshardreplication/vitessshardreplication_controller.go b/pkg/controller/vitessshardreplication/vitessshardreplication_controller.go
@@ -22,10 +22,6 @@ import (
 	"time"
 
 	"github.com/sirupsen/logrus"
-	"vitess.io/vitess/go/mysql/collations"
-	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/sqlparser"
-
 	"vitess.io/vitess/go/vt/logutil"
 	"vitess.io/vitess/go/vt/vttablet/tmclient"
 	"vitess.io/vitess/go/vt/wrangler"
@@ -198,12 +194,7 @@ func (r *ReconcileVitessShard) Reconcile(cctx context.Context, request reconcile
 	tmc := tmclient.NewTabletManagerClient()
 	defer tmc.Close()
 
-	collationEnv := collations.NewEnvironment(servenv.MySQLServerVersion())
-	parser, err := sqlparser.New(sqlparser.Options{
-		MySQLServerVersion: servenv.MySQLServerVersion(),
-		TruncateUILen:      servenv.TruncateUILen,
-		TruncateErrLen:     servenv.TruncateErrLen,
-	})
+	collationEnv, parser, err := environment.GetCollationEnvAndParser()
 	if err != nil {
 		return resultBuilder.Error(err)
 	}

diff --git a/pkg/operator/environment/environment.go b/pkg/operator/environment/environment.go
@@ -25,12 +25,19 @@ import (
 	"time"
 
 	"github.com/spf13/pflag"
+	"vitess.io/vitess/go/mysql/collations"
+	"vitess.io/vitess/go/vt/sqlparser"
 
 	planetscalev2 "planetscale.dev/vitess-operator/pkg/apis/planetscale/v2"
 )
 
 var (
-	reconcileTimeout time.Duration
+	reconcileTimeout   time.Duration
+	mySQLServerVersion = "8.0.30-Vitess"
+	// truncateUILen truncate queries in debug UIs to the given length. 0 means unlimited.
+	truncateUILen = 512
+	// truncateErrLen truncate queries in error logs to the given length. 0 means unlimited.
+	truncateErrLen = 0
 )
 
 // FlagSet returns the FlagSet for the operator.
@@ -58,3 +65,14 @@ func FlagSet() *pflag.FlagSet {
 func ReconcileTimeout() time.Duration {
 	return reconcileTimeout
 }
+
+// GetCollationEnvAndParser gets the collation environment and parser to be used in the operator.
+func GetCollationEnvAndParser() (*collations.Environment, *sqlparser.Parser, error) {
+	collationEnv := collations.NewEnvironment(mySQLServerVersion)
+	parser, err := sqlparser.New(sqlparser.Options{
+		MySQLServerVersion: mySQLServerVersion,
+		TruncateUILen:      truncateUILen,
+		TruncateErrLen:     truncateErrLen,
+	})
+	return collationEnv, parser, err
+}
diff --git a/pkg/operator/vitesstopo/prune_keyspaces.go b/pkg/operator/vitesstopo/prune_keyspaces.go
@@ -24,14 +24,12 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/tools/record"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/vt/logutil"
-	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/wrangler"
 
 	planetscalev2 "planetscale.dev/vitess-operator/pkg/apis/planetscale/v2"
+	"planetscale.dev/vitess-operator/pkg/operator/environment"
 	"planetscale.dev/vitess-operator/pkg/operator/results"
 
 	vtctldatapb "vitess.io/vitess/go/vt/proto/vtctldata"
@@ -94,12 +92,7 @@ func KeyspacesToPrune(keyspaceNames []string, desiredKeyspaces sets.String, orph
 func DeleteKeyspaces(ctx context.Context, ts *topo.Server, recorder record.EventRecorder, eventObj runtime.Object, keyspaceNames []string) (reconcile.Result, error) {
 	resultBuilder := &results.Builder{}
 
-	collationEnv := collations.NewEnvironment(servenv.MySQLServerVersion())
-	parser, err := sqlparser.New(sqlparser.Options{
-		MySQLServerVersion: servenv.MySQLServerVersion(),
-		TruncateUILen:      servenv.TruncateUILen,
-		TruncateErrLen:     servenv.TruncateErrLen,
-	})
+	collationEnv, parser, err := environment.GetCollationEnvAndParser()
 	if err != nil {
 		return resultBuilder.Error(err)
 	}

diff --git a/pkg/operator/vitesstopo/prune_shards.go b/pkg/operator/vitesstopo/prune_shards.go
@@ -24,14 +24,12 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/tools/record"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
-	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/vt/logutil"
-	"vitess.io/vitess/go/vt/servenv"
-	"vitess.io/vitess/go/vt/sqlparser"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/wrangler"
 
 	planetscalev2 "planetscale.dev/vitess-operator/pkg/apis/planetscale/v2"
+	"planetscale.dev/vitess-operator/pkg/operator/environment"
 	"planetscale.dev/vitess-operator/pkg/operator/results"
 )
 
@@ -86,12 +84,7 @@ func ShardsToPrune(currentShards []string, desiredShards sets.String, orphanedSh
 func DeleteShards(ctx context.Context, ts *topo.Server, recorder record.EventRecorder, eventObj runtime.Object, keyspaceName string, shardNames []string) (reconcile.Result, error) {
 	resultBuilder := &results.Builder{}
 
-	collationEnv := collations.NewEnvironment(servenv.MySQLServerVersion())
-	parser, err := sqlparser.New(sqlparser.Options{
-		MySQLServerVersion: servenv.MySQLServerVersion(),
-		TruncateUILen:      servenv.TruncateUILen,
-		TruncateErrLen:     servenv.TruncateErrLen,
-	})
+	collationEnv, parser, err := environment.GetCollationEnvAndParser()
 	if err != nil {
 		return resultBuilder.Error(err)
 	}