From a54171f2c28775f08a8cfc8ad6569f93d22a954d Mon Sep 17 00:00:00 2001
From: Andrii Chubatiuk <andrew.chubatiuk@gmail.com>
Date: Thu, 18 Jan 2024 03:13:47 +0200
Subject: [PATCH] cast query_results data back to text (#6713)

Co-authored-by: Andrew Chubatiuk <andrew.chubatiuk@motional.com>
---
 ..._change_type_of_json_fields_from_varchar_.py | 11 -----------
 redash/models/__init__.py                       |  3 ++-
 redash/models/types.py                          | 17 +++++++++++++++++
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/migrations/versions/7205816877ec_change_type_of_json_fields_from_varchar_.py b/migrations/versions/7205816877ec_change_type_of_json_fields_from_varchar_.py
index 1b4dd55dfe..0e1965e166 100644
--- a/migrations/versions/7205816877ec_change_type_of_json_fields_from_varchar_.py
+++ b/migrations/versions/7205816877ec_change_type_of_json_fields_from_varchar_.py
@@ -59,12 +59,6 @@ def upgrade():
         type_=JSONB(astext_type=sa.Text()),
         postgresql_using='layout::jsonb',
         server_default=sa.text("'{}'::jsonb"))
-    op.alter_column('query_results', 'data',
-        existing_type=sa.Text(),
-        type_=JSONB(astext_type=sa.Text()),
-        nullable=True,
-        postgresql_using='data::jsonb',
-        server_default=sa.text("'{}'::jsonb"))
     op.alter_column('changes', 'change',
         existing_type=JSON(astext_type=sa.Text()),
         type_=JSONB(astext_type=sa.Text()),
@@ -124,11 +118,6 @@ def downgrade():
         type_=sa.Text(),
         postgresql_using='layout::text',
         server_default=sa.text("'{}'::text"))
-    op.alter_column('query_results', 'data',
-        existing_type=JSONB(astext_type=sa.Text()),
-        type_=sa.Text(),
-        postgresql_using='data::text',
-        server_default=sa.text("'{}'::text"))
     op.alter_column('changes', 'change',
         existing_type=JSONB(astext_type=sa.Text()),
         type_=JSON(astext_type=sa.Text()),
diff --git a/redash/models/__init__.py b/redash/models/__init__.py
index 5beefac06b..1a27570c26 100644
--- a/redash/models/__init__.py
+++ b/redash/models/__init__.py
@@ -48,6 +48,7 @@
 from redash.models.types import (
     Configuration,
     EncryptedConfiguration,
+    JSONText,
     MutableDict,
     MutableList,
     json_cast_property,
@@ -315,7 +316,7 @@ class QueryResult(db.Model, BelongsToOrgMixin):
     data_source = db.relationship(DataSource, backref=backref("query_results"))
     query_hash = Column(db.String(32), index=True)
     query_text = Column("query", db.Text)
-    data = Column(MutableDict.as_mutable(JSONB), nullable=True)
+    data = Column(JSONText, nullable=True)
     runtime = Column(DOUBLE_PRECISION)
     retrieved_at = Column(db.DateTime(True))
 
diff --git a/redash/models/types.py b/redash/models/types.py
index 77c1cb4965..b3aa467dcf 100644
--- a/redash/models/types.py
+++ b/redash/models/types.py
@@ -3,6 +3,7 @@
 from sqlalchemy.types import TypeDecorator
 from sqlalchemy_utils import EncryptedType
 
+from redash.utils import json_dumps, json_loads
 from redash.utils.configuration import ConfigurationContainer
 
 from .base import db
@@ -28,6 +29,22 @@ def process_result_value(self, value, dialect):
         )
 
 
+# Utilized for cases when JSON size is bigger than JSONB (255MB) or JSON (10MB) limit
+class JSONText(TypeDecorator):
+    impl = db.Text
+
+    def process_bind_param(self, value, dialect):
+        if value is None:
+            return value
+
+        return json_dumps(value)
+
+    def process_result_value(self, value, dialect):
+        if not value:
+            return value
+        return json_loads(value)
+
+
 class MutableDict(Mutable, dict):
     @classmethod
     def coerce(cls, key, value):