Adding performance schema collector for query assessments (GoogleCloudPlatform#1168)

* Adding performance collector for query assessments

* addressing comments

* addressing comments

Author: shreyakhajanchi

assessment/assessment_engine.go

@@ -32,9 +32,10 @@ import (
 )
 
 type assessmentCollectors struct {
-	sampleCollector        *assessment.SampleCollector
-	infoSchemaCollector    *assessment.InfoSchemaCollector
-	appAssessmentCollector *assessment.MigrationCodeSummarizer
+	sampleCollector            *assessment.SampleCollector
+	infoSchemaCollector        *assessment.InfoSchemaCollector
+	appAssessmentCollector     *assessment.MigrationCodeSummarizer
+	performanceSchemaCollector *assessment.PerformanceSchemaCollector
 }
 
 type assessmentTaskInput struct {
@@ -152,6 +153,17 @@ func initializeCollectors(conv *internal.Conv, sourceProfile profiles.SourceProf
 		logger.Log.Info("app code info unavailable")
 	}
 
+	// Initialize Performance Schema Collector
+	logger.Log.Info("initializing performance schema collector")
+	performanceSchemaCollector, err := assessment.GetDefaultPerformanceSchemaCollector(sourceProfile)
+	if err != nil {
+		logger.Log.Warn("failed to initialize performance schema collector", zap.Error(err))
+		logger.Log.Info("performance schema assessment will be skipped")
+	} else {
+		c.performanceSchemaCollector = &performanceSchemaCollector
+		logger.Log.Info("initialized performance schema collector")
+	}
+
 	return c, err
 }
 

assessment/collectors/common/db_connector.go

@@ -0,0 +1,65 @@
+/* Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.*/
+
+package common
+
+import (
+	"database/sql"
+	"fmt"
+
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/conversion"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/profiles"
+)
+
+// ConnectionConfigProvider interface for getting database connection configuration
+type ConnectionConfigProvider interface {
+	GetConnectionConfig(sourceProfile profiles.SourceProfile) (interface{}, error)
+}
+
+// DefaultConnectionConfigProvider provides default connection configuration
+type DefaultConnectionConfigProvider struct{}
+
+// GetConnectionConfig returns the connection configuration for the given source profile
+func (d DefaultConnectionConfigProvider) GetConnectionConfig(sourceProfile profiles.SourceProfile) (interface{}, error) {
+	return conversion.ConnectionConfig(sourceProfile)
+}
+
+// DBConnector interface for establishing database connections
+type DBConnector interface {
+	Connect(driver string, connectionConfig interface{}) (*sql.DB, error)
+}
+
+// SQLDBConnector provides default SQL database connection functionality
+type SQLDBConnector struct{}
+
+// Connect establishes a database connection using the provided configuration
+func (d SQLDBConnector) Connect(driver string, connectionConfig interface{}) (*sql.DB, error) {
+	connectionStr, ok := connectionConfig.(string)
+	if !ok {
+		return nil, fmt.Errorf("invalid connection configuration type")
+	}
+
+	db, err := sql.Open(driver, connectionStr)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open database connection: %w", err)
+	}
+
+	// Test the connection
+	if err := db.Ping(); err != nil {
+		db.Close()
+		return nil, fmt.Errorf("failed to ping database: %w", err)
+	}
+
+	return db, nil
+}

assessment/collectors/infoschema_collector.go

@@ -20,11 +20,11 @@ import (
 	"sort"
 	"strings"
 
+	collectorCommon "github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/collectors/common"
 	common "github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/sources"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/sources/mysql"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/utils"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants"
-	"github.com/GoogleCloudPlatform/spanner-migration-tool/conversion"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/internal"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/logger"
 	"github.com/GoogleCloudPlatform/spanner-migration-tool/profiles"
@@ -50,10 +50,10 @@ func (c InfoSchemaCollector) IsEmpty() bool {
 }
 
 func GetDefaultInfoSchemaCollector(conv *internal.Conv, sourceProfile profiles.SourceProfile) (InfoSchemaCollector, error) {
-	return GetInfoSchemaCollector(conv, sourceProfile, SQLDBConnector{}, DefaultConnectionConfigProvider{}, getInfoSchema)
+	return GetInfoSchemaCollector(conv, sourceProfile, collectorCommon.SQLDBConnector{}, collectorCommon.DefaultConnectionConfigProvider{}, getInfoSchema)
 }
 
-func GetInfoSchemaCollector(conv *internal.Conv, sourceProfile profiles.SourceProfile, dbConnector DBConnector, configProvider ConnectionConfigProvider, infoSchemaProvider func(*sql.DB, profiles.SourceProfile) (common.InfoSchema, error)) (InfoSchemaCollector, error) {
+func GetInfoSchemaCollector(conv *internal.Conv, sourceProfile profiles.SourceProfile, dbConnector collectorCommon.DBConnector, configProvider collectorCommon.ConnectionConfigProvider, infoSchemaProvider func(*sql.DB, profiles.SourceProfile) (common.InfoSchema, error)) (InfoSchemaCollector, error) {
 	logger.Log.Info("initializing infoschema collector")
 	var errString string
 	connectionConfig, err := configProvider.GetConnectionConfig(sourceProfile)
@@ -428,27 +428,3 @@ func (c InfoSchemaCollector) ListStoredProcedures() map[string]utils.StoredProce
 func (c InfoSchemaCollector) ListSpannerSequences() map[string]ddl.Sequence {
 	return c.conv.SpSequences
 }
-
-type ConnectionConfigProvider interface {
-	GetConnectionConfig(sourceProfile profiles.SourceProfile) (interface{}, error)
-}
-
-type DefaultConnectionConfigProvider struct{}
-
-func (d DefaultConnectionConfigProvider) GetConnectionConfig(sourceProfile profiles.SourceProfile) (interface{}, error) {
-	return conversion.ConnectionConfig(sourceProfile)
-}
-
-type DBConnector interface {
-	Connect(driver string, connectionConfig interface{}) (*sql.DB, error)
-}
-
-type SQLDBConnector struct{}
-
-func (d SQLDBConnector) Connect(driver string, connectionConfig interface{}) (*sql.DB, error) {
-	db, err := sql.Open(driver, connectionConfig.(string))
-	if err != nil {
-		return nil, err
-	}
-	return db, nil
-}

assessment/collectors/performance_schema_collector.go

@@ -0,0 +1,99 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package assessment
+
+import (
+	"database/sql"
+	_ "embed"
+	"fmt"
+
+	collectorCommon "github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/collectors/common"
+	sourcesCommon "github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/sources"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/sources/mysql"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/assessment/utils"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/common/constants"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/logger"
+	"github.com/GoogleCloudPlatform/spanner-migration-tool/profiles"
+	"go.uber.org/zap"
+)
+
+// PerformanceSchemaCollector collects performance schema data from source databases
+type PerformanceSchemaCollector struct {
+	queries []utils.QueryAssessmentInfo
+}
+
+// IsEmpty checks if the collector has any data
+func (c PerformanceSchemaCollector) IsEmpty() bool {
+	return len(c.queries) == 0
+}
+
+// GetDefaultPerformanceSchemaCollector creates a new PerformanceSchemaCollector with default settings
+func GetDefaultPerformanceSchemaCollector(sourceProfile profiles.SourceProfile) (PerformanceSchemaCollector, error) {
+	return GetPerformanceSchemaCollector(sourceProfile, collectorCommon.SQLDBConnector{}, collectorCommon.DefaultConnectionConfigProvider{}, DefaultPerformanceSchemaProvider{})
+}
+
+// GetPerformanceSchemaCollector creates a new PerformanceSchemaCollector with custom dependencies
+func GetPerformanceSchemaCollector(sourceProfile profiles.SourceProfile, dbConnector collectorCommon.DBConnector, configProvider collectorCommon.ConnectionConfigProvider, performanceSchemaProvider PerformanceSchemaProvider) (PerformanceSchemaCollector, error) {
+	logger.Log.Info("initializing performance schema collector")
+
+	connectionConfig, err := configProvider.GetConnectionConfig(sourceProfile)
+	if err != nil {
+		return PerformanceSchemaCollector{}, fmt.Errorf("failed to get connection config: %w", err)
+	}
+
+	db, err := dbConnector.Connect(sourceProfile.Driver, connectionConfig)
+	if err != nil {
+		return PerformanceSchemaCollector{}, fmt.Errorf("failed to connect to database: %w", err)
+	}
+
+	performanceSchema, err := performanceSchemaProvider.getPerformanceSchema(db, sourceProfile)
+	if err != nil {
+		return PerformanceSchemaCollector{}, fmt.Errorf("failed to get performance schema: %w", err)
+	}
+
+	queries, err := performanceSchema.GetAllQueryAssessments()
+	if err != nil {
+		return PerformanceSchemaCollector{}, fmt.Errorf("failed to get all queries: %w", err)
+	}
+
+	logger.Log.Info("performance schema collector initialized successfully",
+		zap.Int("query_count", len(queries)))
+
+	return PerformanceSchemaCollector{
+		queries: queries,
+	}, nil
+}
+
+// DBConnector interface for establishing database connections
+type PerformanceSchemaProvider interface {
+	getPerformanceSchema(db *sql.DB, sourceProfile profiles.SourceProfile) (sourcesCommon.PerformanceSchema, error)
+}
+
+// SQLDBConnector provides default SQL database connection functionality
+type DefaultPerformanceSchemaProvider struct{}
+
+// getPerformanceSchema creates a performance schema implementation based on the database driver
+func (d DefaultPerformanceSchemaProvider) getPerformanceSchema(db *sql.DB, sourceProfile profiles.SourceProfile) (sourcesCommon.PerformanceSchema, error) {
+	driver := sourceProfile.Driver
+	switch driver {
+	case constants.MYSQL:
+		return mysql.PerformanceSchemaImpl{
+			Db:     db,
+			DbName: sourceProfile.Conn.Mysql.Db,
+		}, nil
+	default:
+		return nil, fmt.Errorf("driver %s not supported for performance schema", driver)
+	}
+}