feat(otelgorm): add query formatter option

uptrace · Apr 12, 2022 · e2d9787 · e2d9787
1 parent 01faa43
commit e2d9787
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/otelgorm/option.go b/otelgorm/option.go
@@ -36,6 +36,13 @@ func WithoutQueryVariables() Option {
 	}
 }
 
+// WithQueryFormatter configures a query formatter
+func WithQueryFormatter(queryFormatter func(query string) string) Option {
+	return func(p *otelPlugin) {
+		p.queryFormatter = queryFormatter
+	}
+}
+
 // WithoutMetrics prevents DBStats metrics from being reported.
 func WithoutMetrics() Option {
 	return func(p *otelPlugin) {

diff --git a/otelgorm/otelgorm.go b/otelgorm/otelgorm.go
@@ -22,6 +22,7 @@ type otelPlugin struct {
 	attrs            []attribute.KeyValue
 	excludeQueryVars bool
 	excludeMetrics   bool
+	queryFormatter   func(query string) string
 }
 
 func NewPlugin(opts ...Option) gorm.Plugin {
@@ -124,8 +125,7 @@ func (p *otelPlugin) after() gormHookFunc {
 
 		query := tx.Dialector.Explain(tx.Statement.SQL.String(), vars...)
 
-		attrs = append(attrs, semconv.DBStatementKey.String(query))
-
+		attrs = append(attrs, semconv.DBStatementKey.String(p.formatQuery(query)))
 		if tx.Statement.Table != "" {
 			attrs = append(attrs, semconv.DBSQLTableKey.String(tx.Statement.Table))
 		}
@@ -141,10 +141,19 @@ func (p *otelPlugin) after() gormHookFunc {
 	}
 }
 
+func (p *otelPlugin) formatQuery(query string) string {
+	if p.queryFormatter != nil {
+		return p.queryFormatter(query)
+	}
+	return query
+}
+
 func dbSystem(tx *gorm.DB) attribute.KeyValue {
 	switch tx.Dialector.Name() {
 	case "mysql":
 		return semconv.DBSystemMySQL
+	case "mssql":
+		return semconv.DBSystemMSSQL
 	case "postgres", "postgresql":
 		return semconv.DBSystemPostgreSQL
 	case "sqlite":