emit metric for source peer variant (#3587)

Co-authored-by: Copilot <[email protected]>

Author: heavycrystal

flow/activities/flowable.go

@@ -1626,6 +1626,21 @@ func (a *FlowableActivity) GetFlowMetadata(
 			Type: peerTypes[input.DestinationName],
 		}
 	}
+
+	// Detect source database variant
+	if input.SourceName != "" {
+		if srcConn, err := connectors.GetByNameAs[connectors.DatabaseVariantConnector](ctx, nil, a.CatalogPool, input.SourceName); err == nil {
+			if variant, variantErr := srcConn.GetDatabaseVariant(ctx); variantErr == nil {
+				sourcePeer.Variant = variant
+			} else {
+				logger.Warn("failed to get source database variant", slog.Any("error", variantErr))
+			}
+			connectors.CloseConnector(ctx, srcConn)
+		} else {
+			logger.Warn("failed to get source connector to detect database variant", slog.Any("error", err))
+		}
+	}
+
 	logger.Debug("loaded peer types for flow", slog.String("flowName", input.FlowName),
 		slog.String("sourceName", input.SourceName), slog.String("destinationName", input.DestinationName),
 		slog.Any("peerTypes", peerTypes))

flow/connectors/core.go

@@ -311,6 +311,12 @@ type GetLogRetentionConnector interface {
 	GetLogRetentionHours(ctx context.Context) (float64, error)
 }
 
+type DatabaseVariantConnector interface {
+	Connector
+
+	GetDatabaseVariant(ctx context.Context) (protos.DatabaseVariant, error)
+}
+
 func LoadPeerType(ctx context.Context, catalogPool shared.CatalogPool, peerName string) (protos.DBType, error) {
 	row := catalogPool.QueryRow(ctx, "SELECT type FROM peers WHERE name = $1", peerName)
 	var dbtype protos.DBType
@@ -639,4 +645,7 @@ var (
 
 	_ GetLogRetentionConnector = &connmysql.MySqlConnector{}
 	_ GetLogRetentionConnector = &connmongo.MongoConnector{}
+
+	_ DatabaseVariantConnector = &connpostgres.PostgresConnector{}
+	_ DatabaseVariantConnector = &connmysql.MySqlConnector{}
 )

flow/connectors/mysql/mysql.go

@@ -8,6 +8,7 @@ import (
 	"iter"
 	"log/slog"
 	"net"
+	"strings"
 	"sync/atomic"
 	"time"
 
@@ -476,3 +477,36 @@ func (c *MySqlConnector) StatActivity(
 		StatData: statInfoRows,
 	}, nil
 }
+
+func (c *MySqlConnector) GetDatabaseVariant(ctx context.Context) (protos.DatabaseVariant, error) {
+	query := `SHOW VARIABLES WHERE Variable_name IN ('aurora_version', 'cloudsql_iam_authentication', 'azure_server_name', 'basedir')`
+
+	rs, err := c.Execute(ctx, query)
+	if err != nil {
+		c.logger.Error("failed to execute SHOW VARIABLES for getting database variant", slog.Any("error", err))
+		return protos.DatabaseVariant_VARIANT_UNKNOWN, err
+	}
+
+	var basedirValue string
+	for _, row := range rs.Values {
+		varName := string(row[0].AsString())
+		varValue := string(row[1].AsString())
+		switch varName {
+		case "aurora_version":
+			return protos.DatabaseVariant_AWS_AURORA, nil
+		case "cloudsql_iam_authentication":
+			return protos.DatabaseVariant_GOOGLE_CLOUD_SQL, nil
+		case "azure_server_name":
+			return protos.DatabaseVariant_AZURE_DATABASE, nil
+		case "basedir":
+			basedirValue = varValue
+		}
+	}
+
+	// true for RDS and Aurora, but we check for Aurora via aurora_version above
+	if strings.Contains(basedirValue, "/rdsdbbin/") {
+		return protos.DatabaseVariant_AWS_RDS, nil
+	}
+
+	return protos.DatabaseVariant_VARIANT_UNKNOWN, nil
+}

flow/connectors/postgres/postgres.go

@@ -1725,3 +1725,62 @@ func (c *PostgresConnector) GetVersion(ctx context.Context) (string, error) {
 	c.logger.Info("[postgres] version", slog.String("version", version))
 	return version, nil
 }
+
+func (c *PostgresConnector) GetDatabaseVariant(ctx context.Context) (protos.DatabaseVariant, error) {
+	// First check for Aurora by trying to call aurora_version()
+	var auroraVersion string
+	err := c.conn.QueryRow(ctx, "SELECT aurora_version()").Scan(&auroraVersion)
+	if err == nil && auroraVersion != "" {
+		return protos.DatabaseVariant_AWS_AURORA, nil
+	}
+
+	// If aurora_version() fails, it's not Aurora - continue checking other variants
+	settingsQuery := `
+		SELECT name, setting
+		FROM pg_settings
+		WHERE name IN (
+			'rds.extensions',
+			'cloudsql.logical_decoding',
+			'azure.extensions',
+			'neon.endpoint_id',
+			'extwlist.pscale_allowed_extensions',
+			'supautils.privileged_extensions'
+		) AND setting IS NOT NULL AND setting != ''`
+
+	rows, err := c.conn.Query(ctx, settingsQuery)
+	if err != nil {
+		c.logger.Error("failed to query pg_settings for determining variant", slog.Any("error", err))
+		return protos.DatabaseVariant_VARIANT_UNKNOWN, err
+	}
+	defer rows.Close()
+
+	for rows.Next() {
+		var name, setting string
+		if err := rows.Scan(&name, &setting); err != nil {
+			c.logger.Warn("failed to scan from pg_settings", slog.Any("error", err))
+			continue
+		}
+
+		switch name {
+		case "rds.extensions":
+			return protos.DatabaseVariant_AWS_RDS, nil
+		case "cloudsql.logical_decoding":
+			return protos.DatabaseVariant_GOOGLE_CLOUD_SQL, nil
+		case "azure.extensions":
+			return protos.DatabaseVariant_AZURE_DATABASE, nil
+		case "neon.endpoint_id":
+			return protos.DatabaseVariant_NEON, nil
+		case "extwlist.pscale_allowed_extensions":
+			return protos.DatabaseVariant_PLANETSCALE, nil
+		case "supautils.privileged_extensions":
+			return protos.DatabaseVariant_SUPABASE, nil
+		}
+	}
+
+	if err := rows.Err(); err != nil {
+		c.logger.Error("error iterating pg_settings rows", slog.Any("error", err))
+		return protos.DatabaseVariant_VARIANT_UNKNOWN, err
+	}
+
+	return protos.DatabaseVariant_VARIANT_UNKNOWN, nil
+}

flow/otel_metrics/attributes.go

@@ -15,6 +15,7 @@ const (
 	WorkflowTypeKey            = "workflowType"
 	BatchIdKey                 = "batchId"
 	SourcePeerType             = "sourcePeerType"
+	SourcePeerVariant          = "sourcePeerVariant"
 	DestinationPeerType        = "destinationPeerType"
 	SourcePeerName             = "sourcePeerName"
 	DestinationPeerName        = "destinationPeerName"

flow/otel_metrics/observables.go

@@ -110,6 +110,7 @@ func buildContextualAttributes(ctx context.Context) metric.MeasurementOption {
 		if flowMetadata.Source != nil {
 			attributes = append(attributes,
 				attribute.Stringer(SourcePeerType, flowMetadata.Source.Type),
+				attribute.Stringer(SourcePeerVariant, flowMetadata.Source.Variant),
 				attribute.String(SourcePeerName, flowMetadata.Source.Name))
 		}
 		if flowMetadata.Destination != nil {

protos/flow.proto

@@ -541,6 +541,7 @@ message MaintenanceMirrors {
 message PeerContextMetadata {
   string name = 1;
   peerdb_peers.DBType type = 2;
+  peerdb_peers.DatabaseVariant variant = 3;
 }
 
 message FlowContextMetadataInput {

protos/peers.proto

@@ -274,6 +274,21 @@ enum DBType {
   DBTYPE_UNKNOWN = -1;
 }
 
+enum DatabaseVariant {
+  // either vanilla or something not listed here
+  VARIANT_UNKNOWN = 0;
+  AWS_RDS = 2;
+  AWS_AURORA = 3;
+  GOOGLE_CLOUD_SQL = 4;
+  AZURE_DATABASE = 5;
+  // only Postgres
+  NEON = 6;
+  // has Vitess, but only Postgres for now
+  PLANETSCALE = 7;
+  // only Postgres
+  SUPABASE = 8;
+}
+
 message Peer {
   string name = 1;
   DBType type = 2;