Support for cluster function as table expression (#1408)

This PR is about supporting `FROM cluster(<name>, <table>)`

---------

Signed-off-by: Przemyslaw Delewski <[email protected]>
Co-authored-by: Przemysław Hejman <[email protected]>

Author: pdelewski

platform/clickhouse/clickhouse.go

@@ -137,6 +137,7 @@ type discoveredTable struct {
 	createTableQuery   string
 	timestampFieldName string
 	virtualTable       bool
+	existsOnAllNodes   bool // => table can be queried using `FROM cluster(<clusterName>,<tableName>)` syntax
 }
 
 func (lm *LogManager) ReloadTables() {

platform/clickhouse/table.go

@@ -28,7 +28,8 @@ type Table struct {
 
 	DiscoveredTimestampFieldName *string
 
-	VirtualTable bool
+	VirtualTable     bool
+	ExistsOnAllNodes bool
 }
 
 func (t *Table) createTableOurFieldsString() []string {

platform/clickhouse/table_discovery.go

@@ -217,6 +217,17 @@ func (td *tableDiscovery) ReloadTableDefinitions() {
 			configuredTables = td.configureTables(tables, databaseName)
 		}
 	}
+	var tablePresence map[string][]TablePresence
+	if td.cfg.ClusterName != "" {
+		var err error
+		tablePresence, err = td.getTablePresenceAcrossClusters(databaseName)
+		if err != nil {
+			logger.Warn().Msgf("could not get table presence across clusters: %v", err)
+		}
+		logger.Info().Msgf("Table presence across clusters: %v", tablePresence)
+		configuredTables = td.populateClusterNodes(configuredTables, databaseName, tablePresence)
+	}
+
 	configuredTables = td.readVirtualTables(configuredTables)
 
 	td.ReloadTablesError = nil
@@ -328,7 +339,7 @@ func (td *tableDiscovery) configureTables(tables map[string]map[string]columnMet
 				comment := td.tableComment(databaseName, table)
 				createTableQuery := td.createTableQuery(databaseName, table)
 				// we assume here that @timestamp field is always present in the table, or it's explicitly configured
-				configuredTables[table] = discoveredTable{table, databaseName, columns, indexConfig, comment, createTableQuery, "", false}
+				configuredTables[table] = discoveredTable{table, databaseName, columns, indexConfig, comment, createTableQuery, "", false, false}
 			}
 		} else {
 			notConfiguredTables = append(notConfiguredTables, table)
@@ -358,7 +369,7 @@ func (td *tableDiscovery) autoConfigureTables(tables map[string]map[string]colum
 			maybeTimestampField = td.tableTimestampField(databaseName, table, ClickHouse)
 		}
 		const isVirtualTable = false
-		configuredTables[table] = discoveredTable{table, databaseName, columns, config.IndexConfiguration{}, comment, createTableQuery, maybeTimestampField, isVirtualTable}
+		configuredTables[table] = discoveredTable{table, databaseName, columns, config.IndexConfiguration{}, comment, createTableQuery, maybeTimestampField, isVirtualTable, false}
 
 	}
 	for tableName, table := range configuredTables {
@@ -415,6 +426,7 @@ func (td *tableDiscovery) populateTableDefinitions(configuredTables map[string]d
 				CreateTableQuery:             resTable.createTableQuery,
 				DiscoveredTimestampFieldName: timestampFieldName,
 				VirtualTable:                 resTable.virtualTable,
+				ExistsOnAllNodes:             resTable.existsOnAllNodes,
 			}
 			if containsAttributes(resTable.columnTypes) {
 				table.Config.Attributes = []Attribute{NewDefaultStringAttribute()}
@@ -699,6 +711,90 @@ func (td *tableDiscovery) enrichTableWithMapFields(inputTable map[string]map[str
 	return outputTable
 }
 
+type TablePresence struct {
+	Database         string
+	Table            string
+	FoundNodes       int
+	TotalNodes       int
+	ExistsOnAllNodes bool
+}
+
+func (td *tableDiscovery) getTablePresenceAcrossClusters(database string) (map[string][]TablePresence, error) {
+	// Step 1: Get all cluster names
+	clusterQuery := `SELECT DISTINCT cluster FROM system.clusters ORDER BY cluster`
+	rows, err := td.dbConnPool.Query(context.Background(), clusterQuery)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query cluster list: %w", err)
+	}
+	defer rows.Close()
+
+	var clusters []string
+	for rows.Next() {
+		var cluster string
+		if err := rows.Scan(&cluster); err != nil {
+			return nil, fmt.Errorf("failed to scan cluster name: %w", err)
+		}
+		clusters = append(clusters, cluster)
+	}
+
+	// Step 2: For each cluster, safely query for the given database
+	presenceData := make(map[string][]TablePresence)
+
+	for _, cluster := range clusters {
+		query := `
+            WITH (
+                SELECT count(DISTINCT host_name)
+                FROM system.clusters
+                WHERE cluster = ?
+            ) AS total_nodes
+
+            SELECT
+                database,
+                name AS table_name,
+                count(DISTINCT hostName()) AS found_nodes,
+                total_nodes,
+                count(DISTINCT hostName()) = total_nodes AS exists_on_all_nodes
+            FROM cluster(?, system.tables)
+            WHERE database = ?
+            GROUP BY database, name, total_nodes
+        `
+
+		rows, err := td.dbConnPool.Query(context.Background(), query, cluster, cluster, database)
+		if err != nil {
+			return nil, fmt.Errorf("failed to query tables for cluster %s: %w", cluster, err)
+		}
+		defer rows.Close()
+
+		var tables []TablePresence
+		for rows.Next() {
+			var tp TablePresence
+			if err := rows.Scan(&tp.Database, &tp.Table, &tp.FoundNodes, &tp.TotalNodes, &tp.ExistsOnAllNodes); err != nil {
+				return nil, fmt.Errorf("failed to scan table row: %w", err)
+			}
+			tables = append(tables, tp)
+		}
+
+		if len(tables) > 0 {
+			presenceData[cluster] = tables
+		}
+	}
+
+	return presenceData, nil
+}
+
+func (td *tableDiscovery) populateClusterNodes(configuredTables map[string]discoveredTable, databaseName string, tablePresence map[string][]TablePresence) map[string]discoveredTable {
+	for _, tables := range tablePresence {
+		for _, table := range tables {
+			if table.Database == databaseName {
+				if discoTable, ok := configuredTables[table.Table]; ok {
+					discoTable.existsOnAllNodes = table.ExistsOnAllNodes
+				}
+			}
+		}
+	}
+	return configuredTables
+}
+
 func (td *tableDiscovery) readTables(database string) (map[string]map[string]columnMetadata, error) {
 	logger.Debug().Msgf("describing tables: %s", database)
 

platform/frontend_connectors/schema_transformer.go

@@ -1008,6 +1008,7 @@ func (s *SchemaCheckPass) Transform(queries []*model.Query) ([]*model.Query, err
 		{TransformationName: "MapTransformation", Transformation: s.applyMapTransformations},
 		{TransformationName: "MatchOperatorTransformation", Transformation: s.applyMatchOperator},
 		{TransformationName: "AggOverUnsupportedType", Transformation: s.checkAggOverUnsupportedType},
+		{TransformationName: "ClusterFunction", Transformation: s.applyFromClusterExpression},
 
 		// Section 4: compensations and checks
 		{TransformationName: "BooleanLiteralTransformation", Transformation: s.applyBooleanLiteralLowering},
@@ -1192,3 +1193,27 @@ func (s *SchemaCheckPass) applyMatchOperator(indexSchema schema.Schema, query *m
 	return query, nil
 
 }
+
+// applyFromClusterExpression transforms query so that `FROM table` becomes `FROM cluster(clusterName,table)` if applicable
+func (s *SchemaCheckPass) applyFromClusterExpression(currentSchema schema.Schema, query *model.Query) (*model.Query, error) {
+	if s.cfg.ClusterName == "" {
+		return query, nil
+	}
+	visitor := model.NewBaseVisitor()
+	table, ok := s.tableDiscovery.TableDefinitions().Load(query.TableName)
+	if !ok {
+		return nil, fmt.Errorf("table %s not found", query.TableName)
+	}
+	if !table.ExistsOnAllNodes {
+		return query, nil
+	}
+	visitor.OverrideVisitTableRef = func(b *model.BaseExprVisitor, e model.TableRef) interface{} {
+		return model.NewFunction("cluster", model.NewLiteral(s.cfg.ClusterName), e)
+	}
+	logger.Debug().Msgf("applyClusterFunction: %s", s.cfg.ClusterName)
+	expr := query.SelectCommand.Accept(visitor)
+	if _, ok := expr.(*model.SelectCommand); ok {
+		query.SelectCommand = *expr.(*model.SelectCommand)
+	}
+	return query, nil
+}

platform/frontend_connectors/schema_transformer_test.go

@@ -1764,3 +1764,89 @@ func Test_mapKeys(t *testing.T) {
 	}
 
 }
+
+func Test_cluster(t *testing.T) {
+	indexConfig := map[string]config.IndexConfiguration{
+		"kibana_sample_data_ecommerce": {},
+	}
+	fields := map[schema.FieldName]schema.Field{
+		"@timestamp":         {PropertyName: "@timestamp", InternalPropertyName: "@timestamp", InternalPropertyType: "DateTime64", Type: schema.QuesmaTypeDate},
+		"order_date":         {PropertyName: "order_date", InternalPropertyName: "order_date", InternalPropertyType: "DateTime64", Type: schema.QuesmaTypeDate},
+		"taxful_total_price": {PropertyName: "taxful_total_price", InternalPropertyName: "taxful_total_price", InternalPropertyType: "Float64", Type: schema.QuesmaTypeFloat},
+	}
+
+	indexSchema := schema.Schema{
+		Fields: fields,
+	}
+
+	tableMap := clickhouse.NewTableMap()
+
+	tableDiscovery := clickhouse.NewEmptyTableDiscovery()
+	tableDiscovery.TableMap = tableMap
+	for indexName := range indexConfig {
+		table := clickhouse.NewEmptyTable(indexName)
+		table.ExistsOnAllNodes = true
+		tableMap.Store(indexName, table)
+	}
+
+	clickhouseUrl := &config.Url{
+		Scheme: "clickhouse",
+		Host:   "localhost:9000",
+	}
+
+	clusterName := "my_cluster"
+
+	clickhouseConnector := config.RelationalDbConfiguration{
+		ConnectorType: "clickhouse-os",
+		Url:           clickhouseUrl,
+		ClusterName:   clusterName,
+	}
+	transform := NewSchemaCheckPass(&config.QuesmaConfiguration{IndexConfig: indexConfig, ClickHouse: clickhouseConnector, ClusterName: clusterName}, tableDiscovery, defaultSearchAfterStrategy)
+
+	tests := []struct {
+		name     string
+		query    *model.Query
+		expected *model.Query
+	}{
+		{
+			name: "simple array",
+			query: &model.Query{
+				TableName: "kibana_sample_data_ecommerce",
+				SelectCommand: model.SelectCommand{
+					FromClause: model.NewTableRef("kibana_sample_data_ecommerce"),
+					Columns:    []model.Expr{model.NewWildcardExpr},
+				},
+			},
+			expected: &model.Query{
+				TableName: "kibana_sample_data_ecommerce",
+				SelectCommand: model.SelectCommand{
+					FromClause: model.NewFunction("cluster", model.NewLiteral(clusterName), model.NewLiteral("kibana_sample_data_ecommerce")),
+					Columns:    []model.Expr{model.NewColumnRef("@timestamp"), model.NewColumnRef("order_date"), model.NewColumnRef("taxful_total_price")},
+				},
+			},
+		},
+	}
+	asString := func(query *model.Query) string {
+		return query.SelectCommand.String()
+	}
+
+	for i, tt := range tests {
+		t.Run(util.PrettyTestName(tt.name, i), func(t *testing.T) {
+			tt.query.Schema = indexSchema
+			tt.query.Indexes = []string{tt.query.TableName}
+			actual, err := transform.Transform([]*model.Query{tt.query})
+			assert.NoError(t, err)
+
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			assert.True(t, len(actual) == 1, "len queries == 1")
+
+			expectedJson := asString(tt.expected)
+			actualJson := asString(actual[0])
+
+			assert.Equal(t, expectedJson, actualJson)
+		})
+	}
+}