refactor(player): normalize queries line breaks and spacing (#44299)

This commit is contained in:
Gabriel Corado 2024-07-16 23:28:52 -03:00 committed by GitHub
parent a4570813ee
commit ba68436058
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 73 additions and 1 deletions

View file

@ -98,7 +98,18 @@ func (p *PostgresTranslator) TranslateEvent(evt events.AuditEvent) *events.Sessi
const lineBreak = "\r\n"
func (p *PostgresTranslator) generateCommandPrint(metadata events.DatabaseMetadata, command string) string {
return lineBreak + fmt.Sprintf("%s=> %s", metadata.DatabaseName, command) + lineBreak
lead := metadata.DatabaseName + "=> "
leadSpacing := strings.Repeat(" ", len(lead))
var sb strings.Builder
commandLines := strings.Split(command, "\n")
sb.WriteString(lead + strings.TrimSpace(commandLines[0]) + lineBreak)
for i := 1; i < len(commandLines); i++ {
sb.WriteString(leadSpacing + strings.TrimSpace(commandLines[i]) + lineBreak)
}
return lineBreak + sb.String()
}
func (p *PostgresTranslator) generatePreparedStatementPrint(metadata events.Metadata, databaseMetadata events.DatabaseMetadata, portalName string) *events.SessionPrint {

View file

@ -344,6 +344,67 @@ func TestPostgresRecording(t *testing.T) {
nil,
},
},
"queries with line spacing": {
events: []events.AuditEvent{
&events.DatabaseSessionQuery{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
DatabaseQuery: "SELECT \n1;",
},
&events.DatabaseSessionCommandResult{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
Status: events.Status{Success: true},
AffectedRecords: 1,
},
&events.DatabaseSessionQuery{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
DatabaseQuery: "SELECT *\nfrom events\nlimit 1;",
},
&events.DatabaseSessionCommandResult{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
Status: events.Status{Success: true},
AffectedRecords: 1,
},
&events.DatabaseSessionQuery{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
DatabaseQuery: "SELECT *\r\nfrom events\r\nlimit 1;",
},
&events.DatabaseSessionCommandResult{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
Status: events.Status{Success: true},
AffectedRecords: 1,
},
&events.DatabaseSessionQuery{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
DatabaseQuery: "SELECT *\r\nfrom events\r\nlimit 1;",
},
&events.DatabaseSessionCommandResult{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
Status: events.Status{Success: true},
AffectedRecords: 1,
},
&events.DatabaseSessionQuery{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
DatabaseQuery: "SELECT *\t\r\n from events\t\r\n limit 1;",
},
&events.DatabaseSessionCommandResult{
DatabaseMetadata: events.DatabaseMetadata{DatabaseName: "test"},
Status: events.Status{Success: true},
AffectedRecords: 1,
},
},
expectedPrints: [][]byte{
[]byte("\r\ntest=> SELECT\r\n 1;\r\n"),
[]byte("SUCCESS\r\n(1 row affected)\r\n"),
[]byte("\r\ntest=> SELECT *\r\n from events\r\n limit 1;\r\n"),
[]byte("SUCCESS\r\n(1 row affected)\r\n"),
[]byte("\r\ntest=> SELECT *\r\n from events\r\n limit 1;\r\n"),
[]byte("SUCCESS\r\n(1 row affected)\r\n"),
[]byte("\r\ntest=> SELECT *\r\n from events\r\n limit 1;\r\n"),
[]byte("SUCCESS\r\n(1 row affected)\r\n"),
[]byte("\r\ntest=> SELECT *\r\n from events\r\n limit 1;\r\n"),
[]byte("SUCCESS\r\n(1 row affected)\r\n"),
},
},
} {
t.Run(name, func(t *testing.T) {
translator := NewPostgresTranslator()