[Python] Add: get_source_schema method

redash/query_runner/python.py

@@ -40,8 +40,8 @@ def __call__(self):
 class Python(BaseQueryRunner):
     safe_builtins = (
         'sorted', 'reversed', 'map', 'reduce', 'any', 'all',
-        'slice', 'filter', 'len', 'next', 'enumerate'
-        'sum', 'abs', 'min', 'max', 'round', 'cmp', 'divmod'
+        'slice', 'filter', 'len', 'next', 'enumerate',
+        'sum', 'abs', 'min', 'max', 'round', 'cmp', 'divmod',
         'str', 'unicode', 'int', 'float', 'complex',
         'tuple', 'set', 'list', 'dict', 'bool',
     )
@@ -101,20 +101,24 @@ def custom_import(self, name, globals=None, locals=None, fromlist=(), level=0):
 
         raise Exception("'{0}' is not configured as a supported import module".format(name))
 
-    def custom_write(self, obj):
+    @staticmethod
+    def custom_write(obj):
         """
         Custom hooks which controls the way objects/lists/tuples/dicts behave in
         RestrictedPython
         """
         return obj
 
-    def custom_get_item(self, obj, key):
+    @staticmethod
+    def custom_get_item(obj, key):
         return obj[key]
 
-    def custom_get_iter(self, obj):
+    @staticmethod
+    def custom_get_iter(obj):
         return iter(obj)
 
-    def add_result_column(self, result, column_name, friendly_name, column_type):
+    @staticmethod
+    def add_result_column(result, column_name, friendly_name, column_type):
         """Helper function to add columns inside a Python script running in Redash in an easier way
 
         Parameters:
@@ -135,7 +139,8 @@ def add_result_column(self, result, column_name, friendly_name, column_type):
             "type": column_type
         })
 
-    def add_result_row(self, result, values):
+    @staticmethod
+    def add_result_row(result, values):
         """Helper function to add one row to results set.
 
         Parameters:
@@ -147,7 +152,8 @@ def add_result_row(self, result, values):
 
         result["rows"].append(values)
 
-    def execute_query(self, data_source_name_or_id, query):
+    @staticmethod
+    def execute_query(data_source_name_or_id, query):
         """Run query from specific data source.
 
         Parameters:
@@ -160,7 +166,7 @@ def execute_query(self, data_source_name_or_id, query):
             else:
                 data_source = models.DataSource.get_by_name(data_source_name_or_id)
         except models.NoResultFound:
-                raise Exception("Wrong data source name/id: %s." % data_source_name_or_id)
+            raise Exception("Wrong data source name/id: %s." % data_source_name_or_id)
 
         # TODO: pass the user here...
         data, error = data_source.query_runner.run_query(query, None)
@@ -170,7 +176,25 @@ def execute_query(self, data_source_name_or_id, query):
         # TODO: allow avoiding the json.dumps/loads in same process
         return json.loads(data)
 
-    def get_query_result(self, query_id):
+    @staticmethod
+    def get_source_schema(data_source_name_or_id):
+        """Get schema from specific data source.
+
+        :param data_source_name_or_id: string|integer: Name or ID of the data source
+        :return:
+        """
+        try:
+            if type(data_source_name_or_id) == int:
+                data_source = models.DataSource.get_by_id(data_source_name_or_id)
+            else:
+                data_source = models.DataSource.get_by_name(data_source_name_or_id)
+        except models.NoResultFound:
+            raise Exception("Wrong data source name/id: %s." % data_source_name_or_id)
+        schema = data_source.query_runner.get_schema()
+        return schema
+
+    @staticmethod
+    def get_query_result(query_id):
         """Get result of an existing query.
 
         Parameters:
@@ -216,6 +240,7 @@ def run_query(self, query, user):
 
             restricted_globals = dict(__builtins__=builtins)
             restricted_globals["get_query_result"] = self.get_query_result
+            restricted_globals["get_source_schema"] = self.get_source_schema
             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