Improve Script Query Runner

This commit is contained in:
ariarijp 2018-06-29 00:03:17 +09:00
parent 75df88a8ff
commit 389c6cb3c5
2 changed files with 29 additions and 8 deletions

View File

@ -16,6 +16,16 @@ def query_to_script_path(path, query):
return query
def run_script(script, shell):
output = subprocess.check_output(script, shell=shell)
if output is not None:
output = output.strip()
if output != "":
return output, None
return None, "Error reading output"
class Script(BaseQueryRunner):
@classmethod
def annotate_query(cls):
@ -63,13 +73,7 @@ class Script(BaseQueryRunner):
def run_query(self, query, user):
try:
script = query_to_script_path(self.configuration["path"], query)
output = subprocess.check_output(script, shell=self.configuration['shell'])
if output is not None:
output = output.strip()
if output != "":
return output, None
return None, "Error reading output"
return run_script(script, self.configuration['shell'])
except IOError as e:
return None, e.message
except subprocess.CalledProcessError as e:

View File

@ -1,8 +1,9 @@
import os
import subprocess
from _pytest.monkeypatch import MonkeyPatch
from redash.query_runner.script import query_to_script_path
from redash.query_runner.script import query_to_script_path, run_script
from tests import BaseTestCase
@ -17,3 +18,19 @@ class TestQueryToScript(BaseTestCase):
self.monkeypatch.setattr(os.path, "exists", lambda x: True)
self.assertEqual(["/foo/bar/baz.sh"], query_to_script_path("/foo/bar", "baz.sh"))
class TestRunScript(BaseTestCase):
monkeypatch = MonkeyPatch()
def test_success(self):
self.monkeypatch.setattr(subprocess, "check_output", lambda script, shell: "test")
self.assertEqual(("test", None), run_script("/foo/bar/baz.sh", True))
def test_failure(self):
self.monkeypatch.setattr(subprocess, "check_output", lambda script, shell: None)
self.assertEqual((None, "Error reading output"), run_script("/foo/bar/baz.sh", True))
self.monkeypatch.setattr(subprocess, "check_output", lambda script, shell: "")
self.assertEqual((None, "Error reading output"), run_script("/foo/bar/baz.sh", True))
self.monkeypatch.setattr(subprocess, "check_output", lambda script, shell: " ")
self.assertEqual((None, "Error reading output"), run_script("/foo/bar/baz.sh", True))