Spike max row option for postgres datasource

docker-compose.yml

@@ -32,6 +32,7 @@ services:
       REDASH_DATABASE_URL: "postgresql://postgres@postgres/postgres"
       QUEUES: "queries,scheduled_queries,celery,schemas"
       WORKERS_COUNT: 2
+      QUERY_RESULTS_MAX_ROWS: 50000
   redis:
     image: redis:3-alpine
     restart: unless-stopped

redash/query_runner/pg.py

@@ -6,6 +6,7 @@
 from psycopg2.extras import Range
 
 from redash.query_runner import *
+from redash.settings import QUERY_RESULTS_MAX_ROWS
 from redash.utils import JSONEncoder, json_dumps, json_loads
 
 logger = logging.getLogger(__name__)
@@ -181,6 +182,12 @@ def run_query(self, query, user):
             _wait(connection)
 
             if cursor.description is not None:
+                if QUERY_RESULTS_MAX_ROWS != -1 and cursor.rowcount > QUERY_RESULTS_MAX_ROWS:
+                    json_data = None
+                    error = "Query returned too many rows ({0}) - maximum allowed rows is {1}".format(
+                        cursor.rowcount, QUERY_RESULTS_MAX_ROWS
+                    )
+
                 columns = self.fetch_columns([(i[0], types_map.get(i[1], None)) for i in cursor.description])
                 rows = [dict(zip((c['name'] for c in columns), row)) for row in cursor]
 

redash/settings/__init__.py

@@ -46,6 +46,8 @@ def all_settings():
 QUERY_RESULTS_CLEANUP_COUNT = int(os.environ.get("REDASH_QUERY_RESULTS_CLEANUP_COUNT", "100"))
 QUERY_RESULTS_CLEANUP_MAX_AGE = int(os.environ.get("REDASH_QUERY_RESULTS_CLEANUP_MAX_AGE", "7"))
 
+QUERY_RESULTS_MAX_ROWS = int(os.environ.get("QUERY_RESULTS_MAX_ROWS", "-1"))
+
 SCHEMAS_REFRESH_SCHEDULE = int(os.environ.get("REDASH_SCHEMAS_REFRESH_SCHEDULE", 30))
 SCHEMAS_REFRESH_QUEUE = os.environ.get("REDASH_SCHEMAS_REFRESH_QUEUE", "celery")