[Data Sources] Add: MemSQL query runner

redash/query_runner/memsql_ds.py

@@ -0,0 +1,155 @@
+import json
+import logging
+import sys
+
+from redash.query_runner import *
+from redash.utils import JSONEncoder
+
+logger = logging.getLogger(__name__)
+
+try:
+    from memsql.common import database
+
+    enabled = True
+except ImportError, e:
+    logger.warning(e)
+    enabled = False
+
+COLUMN_NAME = 0
+COLUMN_TYPE = 1
+
+types_map = {
+    'BIGINT': TYPE_INTEGER,
+    'TINYINT': TYPE_INTEGER,
+    'SMALLINT': TYPE_INTEGER,
+    'MEDIUMINT': TYPE_INTEGER,
+    'INT': TYPE_INTEGER,
+    'DOUBLE': TYPE_FLOAT,
+    'DECIMAL': TYPE_FLOAT,
+    'FLOAT': TYPE_FLOAT,
+    'REAL': TYPE_FLOAT,
+    'BOOL': TYPE_BOOLEAN,
+    'BOOLEAN': TYPE_BOOLEAN,
+    'TIMESTAMP': TYPE_DATETIME,
+    'DATETIME': TYPE_DATETIME,
+    'DATE': TYPE_DATETIME,
+    'JSON': TYPE_STRING,
+    'CHAR': TYPE_STRING,
+    'VARCHAR': TYPE_STRING
+}
+
+
+class MemSQL(BaseSQLQueryRunner):
+    noop_query = 'SELECT 1'
+
+    @classmethod
+    def configuration_schema(cls):
+        return {
+            "type": "object",
+            "properties": {
+                "host": {
+                    "type": "string"
+                },
+                "port": {
+                    "type": "number"
+                },
+                "user": {
+                    "type": "string"
+                },
+                "password": {
+                    "type": "string"
+                }
+
+            },
+            "required": ["host", "port"],
+            "secret": ["password"]
+        }
+
+    @classmethod
+    def annotate_query(cls):
+        return False
+
+    @classmethod
+    def type(cls):
+        return "memsql"
+
+    @classmethod
+    def enabled(cls):
+        return enabled
+
+    def __init__(self, configuration):
+        super(MemSQL, self).__init__(configuration)
+
+    def _get_tables(self, schema):
+        schemas_query = "show schemas"
+
+        tables_query = "show tables in %s"
+
+        columns_query = "show columns in %s"
+
+        for schema_name in filter(lambda a: len(a) > 0,
+                                  map(lambda a: str(a['Database']), self._run_query_internal(schemas_query))):
+            for table_name in filter(lambda a: len(a) > 0, map(lambda a: str(a['Tables_in_%s' % schema_name]),
+                                                               self._run_query_internal(
+                                                                       tables_query % schema_name))):
+                table_name = '.'.join((schema_name, table_name))
+                columns = filter(lambda a: len(a) > 0, map(lambda a: str(a['Field']),
+                                                           self._run_query_internal(columns_query % table_name)))
+
+                schema[table_name] = {'name': table_name, 'columns': columns}
+        return schema.values()
+
+    def run_query(self, query, user):
+
+        cursor = None
+        try:
+            cursor = database.connect(**self.configuration.to_dict())
+
+            res = cursor.query(query)
+            # column_names = []
+            # columns = []
+            #
+            # for column in cursor.description:
+            #     column_name = column[COLUMN_NAME]
+            #     column_names.append(column_name)
+            #
+            #     columns.append({
+            #         'name': column_name,
+            #         'friendly_name': column_name,
+            #         'type': types_map.get(column[COLUMN_TYPE], None)
+            #     })
+
+            rows = [dict(zip(list(row.keys()), list(row.values()))) for row in res]
+
+            # ====================================================================================================
+            # temporary - until https://github.com/memsql/memsql-python/pull/8 gets merged
+            # ====================================================================================================
+            columns = []
+            column_names = rows[0].keys() if rows else None
+
+            if column_names:
+                for column in column_names:
+                    columns.append({
+                        'name': column,
+                        'friendly_name': column,
+                        'type': TYPE_STRING
+                    })
+
+            data = {'columns': columns, 'rows': rows}
+            json_data = json.dumps(data, cls=JSONEncoder)
+            error = None
+        except KeyboardInterrupt:
+            cursor.close()
+            error = "Query cancelled by user."
+            json_data = None
+        except Exception as e:
+            logging.exception(e)
+            raise sys.exc_info()[1], None, sys.exc_info()[2]
+        finally:
+            if cursor:
+                cursor.close()
+
+        return json_data, error
+
+
+register(MemSQL)

redash/settings.py

@@ -183,6 +183,7 @@ def all_settings():
     'redash.query_runner.sqlite',
     'redash.query_runner.dynamodb_sql',
     'redash.query_runner.mssql',
+    'redash.query_runner.memsql_ds',
     'redash.query_runner.jql',
     'redash.query_runner.google_analytics',
     'redash.query_runner.snowflake',

requirements_all_ds.txt

@@ -17,6 +17,7 @@ sasl>=0.1.3
 thrift>=0.8.0
 thrift_sasl>=0.1.0
 cassandra-driver==3.1.1
+memsql==2.16.0
 snowflake_connector_python==1.3.7
 atsd_client==2.0.12
 simple_salesforce==0.72.2