add user facing logging when table is too large

Author: jgao54

flow/activities/flowable.go

@@ -514,6 +514,19 @@ func (a *FlowableActivity) GetQRepPartitions(ctx context.Context,
 	}
 	defer srcClose(ctx)
 
+	if tableSizeEstimatorConn, ok := srcConn.(connectors.TableSizeEstimatorConnector); ok {
+		// expect estimate query execution to be fast, so set a short timeout
+		timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+		defer cancel()
+		if bytes, connErr := tableSizeEstimatorConn.GetTableSizeEstimatedBytes(timeoutCtx, config.WatermarkTable); connErr == nil {
+			if bytes > 1024*1024*1024*2024 {
+				msg := fmt.Sprintf("Large table detected: %s (%s). Partitioning query may take several hours to execute. "+
+					"This is normal for tables over 1 TB.", config.WatermarkTable, utils.FormatTableSize(bytes))
+				a.Alerter.LogFlowInfo(ctx, config.FlowJobName, msg)
+			}
+		}
+	}
+
 	partitions, err := srcConn.GetQRepPartitions(ctx, config, last)
 	if err != nil {
 		return nil, a.Alerter.LogFlowWrappedError(ctx, config.FlowJobName, "failed to get partitions from source", err)

flow/connectors/core.go

@@ -331,6 +331,12 @@ type DatabaseVariantConnector interface {
 	GetDatabaseVariant(ctx context.Context) (protos.DatabaseVariant, error)
 }
 
+type TableSizeEstimatorConnector interface {
+	Connector
+
+	GetTableSizeEstimatedBytes(ctx context.Context, tableName string) (int64, 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
@@ -673,4 +679,7 @@ var (
 
 	_ DatabaseVariantConnector = &connpostgres.PostgresConnector{}
 	_ DatabaseVariantConnector = &connmysql.MySqlConnector{}
+
+	_ TableSizeEstimatorConnector = &connpostgres.PostgresConnector{}
+	_ TableSizeEstimatorConnector = &connmysql.MySqlConnector{}
 )

flow/connectors/mysql/mysql.go

@@ -515,3 +515,23 @@ func (c *MySqlConnector) GetDatabaseVariant(ctx context.Context) (protos.Databas
 
 	return protos.DatabaseVariant_VARIANT_UNKNOWN, nil
 }
+
+func (c *MySqlConnector) GetTableSizeEstimatedBytes(ctx context.Context, fullyQualifiedTable string) (int64, error) {
+	schemaTable, err := utils.ParseSchemaTable(fullyQualifiedTable)
+	if err != nil {
+		return 0, err
+	}
+	query := fmt.Sprintf(
+		"SELECT data_length FROM information_schema.tables WHERE table_schema = '%s' AND table_name = '%s'",
+		mysql.Escape(schemaTable.Schema),
+		mysql.Escape(schemaTable.Table),
+	)
+
+	rs, err := c.Execute(ctx, query)
+	if err != nil {
+		c.logger.Warn("failed to get estimated table size", slog.Any("error", err))
+		return 0, err
+	}
+	defer rs.Close()
+	return rs.GetInt(0, 0)
+}

flow/connectors/postgres/postgres.go

@@ -1803,3 +1803,16 @@ func (c *PostgresConnector) GetDatabaseVariant(ctx context.Context) (protos.Data
 
 	return protos.DatabaseVariant_VARIANT_UNKNOWN, nil
 }
+
+func (c *PostgresConnector) GetTableSizeEstimatedBytes(ctx context.Context, fullyQualifiedTableName string) (int64, error) {
+	tableSizeQuery := "SELECT pg_relation_size(to_regclass($1))"
+	var tableSizeBytes pgtype.Int8
+	if err := c.conn.QueryRow(ctx, tableSizeQuery, fullyQualifiedTableName).Scan(&tableSizeBytes); err != nil {
+		c.logger.Warn("failed to get estimated table size", slog.Any("error", err))
+		return 0, err
+	}
+	if !tableSizeBytes.Valid {
+		return 0, errors.New("table size is not valid")
+	}
+	return tableSizeBytes.Int64, nil
+}