Merge pull request #503 from EverythingMe/fix/python_result_set

Fix: if you change the result object, python runner wouldn't return any results

Makefile

@@ -13,7 +13,7 @@ deps:
 
 pack:
 	sed -ri "s/^__version__ = '([0-9.]*)'/__version__ = '$(FULL_VERSION)'/" redash/__init__.py
-	tar -zcv -f $(FILENAME) --exclude=".git*" --exclude="*.pyc" --exclude="*.pyo" --exclude="venv" --exclude="rd_ui/node_modules" --exclude="rd_ui/dist/bower_components" --exclude="rd_ui/app" *
+	tar -zcv -f $(FILENAME) --exclude="optipng*" --exclude=".git*" --exclude="*.pyc" --exclude="*.pyo" --exclude="venv" --exclude="rd_ui/node_modules" --exclude="rd_ui/dist/bower_components" --exclude="rd_ui/app" *
 
 upload:
 	python bin/release_manager.py $(CIRCLE_SHA1) $(BASE_VERSION) $(FILENAME)

circle.yml

@@ -10,8 +10,6 @@ dependencies:
     - wget http://downloads.sourceforge.net/project/optipng/OptiPNG/optipng-0.7.5/optipng-0.7.5.tar.gz
     - tar xvf optipng-0.7.5.tar.gz
     - cd optipng-0.7.5; ./configure; make; sudo checkinstall -y;
-    - rm -rf optipng-0.7.5
-    - rm optipng-0.7.5.tar.gz
     - make deps
     - pip install -r dev_requirements.txt
     - pip install -r requirements.txt

redash/query_runner/python.py

@@ -14,23 +14,30 @@
 from RestrictedPython import compile_restricted
 from RestrictedPython.Guards import safe_builtins
 
+
 class CustomPrint(object):
     """ CustomPrint redirect "print" calls to be sent as "log" on the result object """
-    def __init__(self, python_runner):
-        self._python_runner = python_runner
+    def __init__(self):
+        self.enabled = True
+        self.lines = []
 
     def write(self, text):
-        if self._python_runner()._enable_print_log:
+        if self.enabled:
             if text and text.strip():
                 log_line = "[{0}] {1}".format(datetime.datetime.utcnow().isoformat(), text)
-                self._python_runner()._result["log"].append(log_line)
+                self.lines.append(log_line)
+
+    def enable(self):
+        self.enabled = True
+
+    def disable(self):
+        self.enabled = False
 
     def __call__(self):
         return self
 
 
 class Python(BaseQueryRunner):
-
     @classmethod
     def configuration_schema(cls):
         return {
@@ -52,15 +59,14 @@ def annotate_query(cls):
         return False
 
     def __init__(self, configuration_json):
-        global ALLOWED_MODULES
-
         super(Python, self).__init__(configuration_json)
 
         self.syntax = "python"
 
         self._allowed_modules = {}
-        self._result = { "rows" : [], "columns" : [], "log" : [] }
+        self._script_locals = { "result" : { "rows" : [], "columns" : [], "log" : [] } }
         self._enable_print_log = True
+        self._custom_print = CustomPrint()
 
         if self.configuration.get("allowedImportModules", None):
             for item in self.configuration["allowedImportModules"].split(","):
@@ -92,12 +98,6 @@ def custom_get_item(self, obj, key):
     def custom_get_iter(self, obj):
         return iter(obj)
 
-    def disable_print_log(self):
-        self._enable_print_log = False
-
-    def enable_print_log(self):
-        self._enable_print_log = True
-
     def add_result_column(self, result, column_name, friendly_name, column_type):
         """ Helper function to add columns inside a Python script running in re:dash in an easier way """
         if column_type not in SUPPORTED_COLUMN_TYPES:
@@ -164,17 +164,15 @@ def run_query(self, query):
             safe_builtins["setattr"] = setattr
             safe_builtins["_getitem_"] = self.custom_get_item
             safe_builtins["_getiter_"] = self.custom_get_iter
-            safe_builtins["_print_"] = CustomPrint(weakref.ref(self))
-
-            script_locals = { "result" : self._result }
+            safe_builtins["_print_"] = self._custom_print
 
             restricted_globals = dict(__builtins__=safe_builtins)
             restricted_globals["get_query_result"] = self.get_query_result
             restricted_globals["execute_query"] = self.execute_query
             restricted_globals["add_result_column"] = self.add_result_column
             restricted_globals["add_result_row"] = self.add_result_row
-            restricted_globals["disable_print_log"] = self.disable_print_log
-            restricted_globals["enable_print_log"] = self.enable_print_log
+            restricted_globals["disable_print_log"] = self._custom_print.disable
+            restricted_globals["enable_print_log"] = self._custom_print.enable
 
             restricted_globals["TYPE_DATETIME"] = TYPE_DATETIME
             restricted_globals["TYPE_BOOLEAN"] = TYPE_BOOLEAN
@@ -187,9 +185,11 @@ def run_query(self, query):
             #       One option is to use ETA with Celery + timeouts on workers
             #       And replacement of worker process every X requests handled.
 
-            exec(code) in restricted_globals, script_locals
+            exec(code) in restricted_globals, self._script_locals
 
-            json_data = json.dumps(self._result)
+            result = self._script_locals['result']
+            result['log'] = self._custom_print.lines
+            json_data = json.dumps(result)
         except KeyboardInterrupt:
             error = "Query cancelled by user."
             json_data = None