From ea3fe4e52cee1000c0a8eea38c4a7a85c557e1e8 Mon Sep 17 00:00:00 2001
From: Nityananda Gohain <nityanandagohain@gmail.com>
Date: Sat, 26 Oct 2024 13:14:16 +0530
Subject: [PATCH] feat: support for modify query mat col in migrations (#436)

---
 .../schema_migrator/table_operations.go       | 54 +++++++++++++++++++
 .../schema_migrator/table_operations_test.go  | 24 +++++++++
 2 files changed, 78 insertions(+)

diff --git a/cmd/signozschemamigrator/schema_migrator/table_operations.go b/cmd/signozschemamigrator/schema_migrator/table_operations.go
index c7d202d6..b96e6bc1 100644
--- a/cmd/signozschemamigrator/schema_migrator/table_operations.go
+++ b/cmd/signozschemamigrator/schema_migrator/table_operations.go
@@ -203,3 +203,57 @@ func (c CreateMaterializedViewOperation) ToSQL() string {
 	sql.WriteString(c.Query)
 	return sql.String()
 }
+
+// ModifyQueryMaterializedViewOperation is used to represent the ALTER TABLE ... MODIFY QUERY statement in the SQL.
+type ModifyQueryMaterializedViewOperation struct {
+	cluster  string
+	Database string
+	ViewName string
+	Query    string
+}
+
+// OnCluster is used to specify the cluster on which the operation should be performed.
+// This is useful when the operation is to be performed on a cluster setup.
+func (c ModifyQueryMaterializedViewOperation) OnCluster(cluster string) Operation {
+	c.cluster = cluster
+	return &c
+}
+
+func (c ModifyQueryMaterializedViewOperation) WithReplication() Operation {
+	// no-op
+	return &c
+}
+
+func (c ModifyQueryMaterializedViewOperation) ShouldWaitForDistributionQueue() (bool, string, string) {
+	return false, c.Database, c.ViewName
+}
+
+func (c ModifyQueryMaterializedViewOperation) IsMutation() bool {
+	// Modify materialized view is not a mutation.
+	return false
+}
+
+func (c ModifyQueryMaterializedViewOperation) IsIdempotent() bool {
+	// Modify materialized view is idempotent. It will not change the materialized view if the materialized view already exists.
+	return true
+}
+
+func (c ModifyQueryMaterializedViewOperation) IsLightweight() bool {
+	// Modify materialized view is lightweight.
+	return true
+}
+
+func (c ModifyQueryMaterializedViewOperation) ToSQL() string {
+	var sql strings.Builder
+	sql.WriteString("ALTER TABLE ")
+	sql.WriteString(c.Database)
+	sql.WriteString(".")
+	sql.WriteString(c.ViewName)
+	if c.cluster != "" {
+		sql.WriteString(" ON CLUSTER ")
+		sql.WriteString(c.cluster)
+	}
+	sql.WriteString(" MODIFY QUERY ")
+	sql.WriteString(c.Query)
+	return sql.String()
+}
diff --git a/cmd/signozschemamigrator/schema_migrator/table_operations_test.go b/cmd/signozschemamigrator/schema_migrator/table_operations_test.go
index c360efe4..9bed0ad7 100644
--- a/cmd/signozschemamigrator/schema_migrator/table_operations_test.go
+++ b/cmd/signozschemamigrator/schema_migrator/table_operations_test.go
@@ -203,6 +203,30 @@ func TestCreateMaterializedView(t *testing.T) {
 	}
 }
 
+func TestModifyQueryMaterializedView(t *testing.T) {
+	testCases := []struct {
+		name string
+		op   Operation
+		want string
+	}{
+		{
+			name: "modify-materialized-view",
+			op: ModifyQueryMaterializedViewOperation{
+				Database: "db",
+				ViewName: "view",
+				Query:    "SELECT id, ts, value FROM db.table",
+			},
+			want: "ALTER TABLE db.view MODIFY QUERY SELECT id, ts, value FROM db.table",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.want, tc.op.ToSQL())
+		})
+	}
+}
+
 func TestCreateWithCluster(t *testing.T) {
 	testCases := []struct {
 		name string