github.com/pinpoint-apm/pinpoint-go-agent@v1.4.1-0.20240110120318-a50c2eb18c8c/sql_util.go (about)

     1  package pinpoint
     2  
     3  import (
     4  	"bufio"
     5  	"strconv"
     6  	"strings"
     7  )
     8  
     9  type sqlNormalizer struct {
    10  	r          *bufio.Reader
    11  	output     *strings.Builder
    12  	param      *strings.Builder
    13  	paramIndex int
    14  	sql        string
    15  	isChanged  bool
    16  }
    17  
    18  func newSqlNormalizer(sql string) *sqlNormalizer {
    19  	normalizer := sqlNormalizer{}
    20  
    21  	normalizer.r = bufio.NewReader(strings.NewReader(sql))
    22  	normalizer.output = &strings.Builder{}
    23  	normalizer.param = &strings.Builder{}
    24  	normalizer.paramIndex = 0
    25  	normalizer.sql = sql
    26  	normalizer.isChanged = false
    27  
    28  	return &normalizer
    29  }
    30  
    31  func (s *sqlNormalizer) run() (string, string) {
    32  	numberTokenStartEnable := false
    33  
    34  	for {
    35  		if ch := s.read(); ch == eof {
    36  			break
    37  		} else if ch == '/' {
    38  			s.output.WriteRune(ch)
    39  			if s.lookahead('/') {
    40  				s.consumeSingleLineComment()
    41  			} else if s.lookahead('*') {
    42  				s.consumeMultiLineComment()
    43  			} else {
    44  				numberTokenStartEnable = true
    45  			}
    46  		} else if ch == '-' {
    47  			s.output.WriteRune(ch)
    48  			if s.lookahead('-') {
    49  				s.consumeSingleLineComment()
    50  			} else {
    51  				numberTokenStartEnable = true
    52  			}
    53  		} else if ch == '\'' {
    54  			s.output.WriteRune(ch)
    55  			if s.lookahead('\'') {
    56  				s.output.WriteRune(s.read())
    57  			} else {
    58  				s.consumeCharLiteral()
    59  			}
    60  		} else if isDigit(ch) {
    61  			if numberTokenStartEnable {
    62  				s.unread()
    63  				s.consumeNumberLiteral()
    64  			} else {
    65  				s.output.WriteRune(ch)
    66  			}
    67  		} else if isLetter(ch) || ch == '.' || ch == '_' || ch == '@' || ch == ':' || ch == '$' {
    68  			numberTokenStartEnable = false
    69  			s.output.WriteRune(ch)
    70  		} else {
    71  			numberTokenStartEnable = true
    72  			s.output.WriteRune(ch)
    73  		}
    74  	}
    75  
    76  	if s.isChanged {
    77  		if s.param.Len() > 0 {
    78  			return s.output.String(), s.param.String()
    79  		} else {
    80  			return s.output.String(), ""
    81  		}
    82  	} else {
    83  		return s.sql, ""
    84  	}
    85  
    86  }
    87  
    88  func (s *sqlNormalizer) consumeSingleLineComment() {
    89  	var ch rune
    90  
    91  	for {
    92  		if ch = s.read(); ch == eof {
    93  			break
    94  		}
    95  		s.output.WriteRune(ch)
    96  		if ch == '\n' {
    97  			break
    98  		}
    99  	}
   100  }
   101  
   102  func (s *sqlNormalizer) consumeMultiLineComment() {
   103  	var ch rune
   104  	prev := eof
   105  	s.output.WriteRune(s.read()) /* cousume '*' */
   106  
   107  	for {
   108  		if ch = s.read(); ch == eof {
   109  			break
   110  		}
   111  		s.output.WriteRune(ch)
   112  		if prev == '*' && ch == '/' {
   113  			break
   114  		}
   115  		prev = ch
   116  	}
   117  }
   118  
   119  func (s *sqlNormalizer) consumeCharLiteral() {
   120  	var ch rune
   121  
   122  	s.isChanged = true
   123  	if s.param.Len() > 0 {
   124  		s.param.WriteRune(',')
   125  	}
   126  
   127  	for {
   128  		if ch = s.read(); ch == eof {
   129  			break
   130  		}
   131  
   132  		if ch == ',' {
   133  			s.param.WriteRune(ch)
   134  		} else if ch == '\'' {
   135  			if s.lookahead('\'') {
   136  				s.param.WriteRune(s.read())
   137  			} else {
   138  				s.output.WriteString(strconv.Itoa(s.paramIndex))
   139  				s.paramIndex++
   140  				s.output.WriteRune('$')
   141  				s.output.WriteRune('\'')
   142  				break
   143  			}
   144  		}
   145  
   146  		s.param.WriteRune(ch)
   147  	}
   148  }
   149  
   150  func (s *sqlNormalizer) consumeNumberLiteral() {
   151  	var ch rune
   152  
   153  	s.isChanged = true
   154  	if s.param.Len() > 0 {
   155  		s.param.WriteRune(',')
   156  	}
   157  	s.output.WriteString(strconv.Itoa(s.paramIndex))
   158  	s.paramIndex++
   159  	s.output.WriteRune('#')
   160  
   161  	for {
   162  		if ch = s.read(); ch == eof {
   163  			break
   164  		}
   165  
   166  		if isDigit(ch) || ch == '.' || ch == 'E' || ch == 'e' {
   167  			s.param.WriteRune(ch)
   168  		} else {
   169  			s.unread()
   170  			break
   171  		}
   172  	}
   173  }
   174  
   175  func (s *sqlNormalizer) read() rune {
   176  	ch, _, err := s.r.ReadRune()
   177  	if err != nil {
   178  		return eof
   179  	}
   180  	return ch
   181  }
   182  
   183  func (s *sqlNormalizer) unread() {
   184  	_ = s.r.UnreadRune()
   185  }
   186  
   187  func (s *sqlNormalizer) lookahead(expected rune) bool {
   188  	ch, _, err := s.r.ReadRune()
   189  	_ = s.r.UnreadRune()
   190  	if err != nil {
   191  		return false
   192  	}
   193  	return ch == expected
   194  }
   195  
   196  func isLetter(ch rune) bool {
   197  	return (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z')
   198  }
   199  
   200  func isDigit(ch rune) bool {
   201  	return ch >= '0' && ch <= '9'
   202  }
   203  
   204  var eof = rune(0)