From 95cb257a5fdac22e71d3ab175653a3d4cac4d511 Mon Sep 17 00:00:00 2001
From: Abhijeet Prasad <aprasad@sentry.io>
Date: Tue, 8 Aug 2023 11:32:23 -0400
Subject: [PATCH] feat(spans): Add `SAVEPOINT` and `SET` sql db actions (#2385)

Add support for parsing `SET`, `SAVEPOINT`, `RELEASE SAVEPOINT`, and
`ROLLBACK TO SAVEPOINT` db sql actions.

Adds table test to validate this.

Co-authored-by: Iker Barriocanal <32816711+iker-barriocanal@users.noreply.github.com>
---
 .../store/normalize/span/tag_extraction.rs    | 37 +++++++++++++++++--
 ..._spans__tests__extract_span_metrics-2.snap |  2 +
 ...n__spans__tests__extract_span_metrics.snap |  3 ++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/relay-general/src/store/normalize/span/tag_extraction.rs b/relay-general/src/store/normalize/span/tag_extraction.rs
index 40b3705777..107165daed 100644
--- a/relay-general/src/store/normalize/span/tag_extraction.rs
+++ b/relay-general/src/store/normalize/span/tag_extraction.rs
@@ -333,9 +333,10 @@ fn truncate_string(mut string: String, max_bytes: usize) -> String {
 
 /// Regex with a capture group to extract the database action from a query.
 ///
-/// Currently, we're only interested in `SELECT`, `INSERT`, `DELETE` and `UPDATE` statements.
-static SQL_ACTION_EXTRACTOR_REGEX: Lazy<Regex> =
-    Lazy::new(|| Regex::new(r#"(?i)(?P<action>(SELECT|INSERT|DELETE|UPDATE))"#).unwrap());
+/// Currently we have an explicit allow-list of database actions considered important.
+static SQL_ACTION_EXTRACTOR_REGEX: Lazy<Regex> = Lazy::new(|| {
+    Regex::new(r#"(?i)(?P<action>(SELECT|INSERT|DELETE|UPDATE|SET|SAVEPOINT|RELEASE SAVEPOINT|ROLLBACK TO SAVEPOINT))"#).unwrap()
+});
 
 fn sql_action_from_query(query: &str) -> Option<&str> {
     extract_captured_substring(query, &SQL_ACTION_EXTRACTOR_REGEX)
@@ -560,4 +561,34 @@ mod tests {
         let query = r#"UPDATE "a" SET "x" = %s, "y" = %s WHERE "z" = %s"#;
         assert_eq!(sql_table_from_query(query).unwrap(), "a");
     }
+
+    #[test]
+    fn extract_sql_action() {
+        let test_cases = vec![
+            (
+                r#"SELECT "sentry_organization"."id" FROM "sentry_organization" WHERE "sentry_organization"."id" = %s"#,
+                "SELECT",
+            ),
+            (
+                r#"INSERT INTO "sentry_groupseen" ("project_id", "group_id", "user_id", "last_seen") VALUES (%s, %s, %s, %s) RETURNING "sentry_groupseen"."id"#,
+                "INSERT",
+            ),
+            (
+                r#"UPDATE sentry_release SET date_released = %s WHERE id = %s"#,
+                "UPDATE",
+            ),
+            (
+                r#"DELETE FROM "sentry_groupinbox" WHERE "sentry_groupinbox"."id" IN (%s)"#,
+                "DELETE",
+            ),
+            (r#"SET search_path TO my_schema, public"#, "SET"),
+            (r#"SAVEPOINT %s"#, "SAVEPOINT"),
+            (r#"RELEASE SAVEPOINT %s"#, "RELEASE SAVEPOINT"),
+            (r#"ROLLBACK TO SAVEPOINT %s"#, "ROLLBACK TO SAVEPOINT"),
+        ];
+
+        for (query, expected) in test_cases {
+            assert_eq!(sql_action_from_query(query).unwrap(), expected)
+        }
+    }
 }
diff --git a/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics-2.snap b/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics-2.snap
index 414f8ccf0c..7a77037371 100644
--- a/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics-2.snap
+++ b/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics-2.snap
@@ -772,6 +772,7 @@ expression: metrics
         tags: {
             "environment": "fake_environment",
             "http.status_code": "500",
+            "span.action": "SAVEPOINT",
             "span.category": "db",
             "span.description": "SAVEPOINT %s",
             "span.group": "3f955cbde39e04b9",
@@ -793,6 +794,7 @@ expression: metrics
         tags: {
             "environment": "fake_environment",
             "http.status_code": "500",
+            "span.action": "SAVEPOINT",
             "span.category": "db",
             "span.description": "SAVEPOINT %s",
             "span.group": "3f955cbde39e04b9",
diff --git a/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics.snap b/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics.snap
index 3e7ab1168d..5f461c99ca 100644
--- a/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics.snap
+++ b/relay-server/src/metrics_extraction/spans/snapshots/relay_server__metrics_extraction__spans__tests__extract_span_metrics.snap
@@ -1417,6 +1417,9 @@ expression: event.value().unwrap().spans
             "release": String(
                 "1.2.3",
             ),
+            "span.action": String(
+                "SAVEPOINT",
+            ),
             "span.category": String(
                 "db",
             ),