allow only specific columns in superset dataset (#1567)

Co-authored-by: Ali Hassan <[email protected]>

questionnaires/superset/__init__.py

@@ -0,0 +1 @@
+ALLOWED_COLUMNS = ['id', 'form_data', 'submit_time', 'page_id', 'user_id']

questionnaires/superset/client.py

@@ -32,6 +32,11 @@ def _get_headers(self):
         headers.update(self._get_csrf_token())
         return headers
 
+    def _get_params(self):
+        return {
+            'override_columns': 'true'
+        }
+
     def _get_validated_response(self, response):
         if response.ok:
             return response.json()
@@ -82,7 +87,7 @@ def get_dataset(self, id):
         return self._api_caller(request)
 
     def update_dataset(self, id, data):
-        request = Request(method='PUT', url=f'{self.dataset_url}{id}', headers=self._get_headers(), json=data)
+        request = Request(method='PUT', url=f'{self.dataset_url}{id}', params=self._get_params(), headers=self._get_headers(), json=data)
         return self._api_caller(request)
 
     def create_chart(self, data):

questionnaires/superset/datasets.py

@@ -1,3 +1,6 @@
+from questionnaires.superset import ALLOWED_COLUMNS
+
+
 class Dataset:
     def __init__(self, database_id, owner_id, table_name, dataset_name, page_id):
         self.database_id = database_id
@@ -16,9 +19,10 @@ def post_body(self):
         }
 
     def put_body(self, columns, metrics):
-        sql = f"SELECT * " \
+        sql = f"SELECT {', '.join(ALLOWED_COLUMNS)} " \
               f"FROM {self.table_name} " \
               f"WHERE page_id = {self.page_id}"
+
         return {
             'sql': sql,
             'table_name': self.dataset_name,

questionnaires/superset/utils.py

@@ -16,6 +16,7 @@
 from questionnaires.superset.client import SupersetClient
 from questionnaires.superset.dashboard import Dashboard
 from questionnaires.superset.datasets import Dataset
+from questionnaires.superset import ALLOWED_COLUMNS
 
 CHART_TYPE_MAP = {
     'checkbox': PieChart,
@@ -85,12 +86,12 @@ def _create_dataset(self, database_id, owner_id):
         dataset_id = resp.get('id')
 
         dataset_detail = self.client.get_dataset(dataset_id)
-        columns = copy(dataset_detail.get('result', {}).get('columns'))
-        for column in columns:
-            column.pop('changed_on', None)
-            column.pop('created_on', None)
-            column.pop('type_generic', None)
-            column.pop('uuid', None)
+
+        columns = [
+            {'column_name': column.get('column_name')}
+            for column in dataset_detail.get('result', {}).get('columns')
+            if column.get('column_name') in ALLOWED_COLUMNS
+        ]
 
         for question in self.questions:
             calculated_column_expression = CALCULATED_COLUMN_EXPRESSION_MAP.get(question.field_type)