From 010ed5765e32a55311c9865be4d75ba5fccfaea3 Mon Sep 17 00:00:00 2001
From: Mark Nuttall-Smith <nuttall@apolloagriculture.com>
Date: Thu, 27 Mar 2025 11:59:09 +0300
Subject: [PATCH] Configurable addition of Redash metadata as query tag

---
 redash/query_runner/snowflake.py | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/redash/query_runner/snowflake.py b/redash/query_runner/snowflake.py
index bb67c20d68..7df1634990 100644
--- a/redash/query_runner/snowflake.py
+++ b/redash/query_runner/snowflake.py
@@ -17,6 +17,7 @@
     BaseSQLQueryRunner,
     register,
 )
+from redash.utils import json_dumps
 
 TYPES_MAP = {
     0: TYPE_INTEGER,
@@ -33,8 +34,14 @@
 
 
 class Snowflake(BaseSQLQueryRunner):
+    add_query_tags = False
+    metadata = None
     noop_query = "SELECT 1"
 
+    def __init__(self, configuration):
+        super().__init__(configuration)
+        self.add_query_tags = configuration.get("add_query_tags", False)
+
     @classmethod
     def configuration_schema(cls):
         return {
@@ -52,6 +59,11 @@ def configuration_schema(cls):
                     "default": False,
                 },
                 "host": {"type": "string"},
+                "add_query_tags": {
+                    "type": "boolean",
+                    "title": "Add Redash metadata as query tag",
+                    "default": False,
+                },
             },
             "order": [
                 "account",
@@ -69,6 +81,11 @@ def configuration_schema(cls):
             ],
         }
 
+    def annotate_query(self, query, metadata):
+        annotated = super(Snowflake, self).annotate_query(query, metadata)
+        self.metadata = metadata
+        return annotated
+
     @classmethod
     def enabled(cls):
         return enabled
@@ -130,6 +147,10 @@ def run_query(self, query, user):
             cursor.execute("USE WAREHOUSE {}".format(self.configuration["warehouse"]))
             cursor.execute("USE {}".format(self.configuration["database"]))
 
+            if self.add_query_tags:
+                if self.metadata:
+                    cursor.execute(f"ALTER SESSION SET QUERY_TAG=$$'{json_dumps(self.metadata)}'$$")
+
             cursor.execute(query)
 
             data = self._parse_results(cursor)