[GitHub] timifasubaa closed pull request #6015: Allow customizable cursor init params

[GitBox <gi...@apache.org>]

timifasubaa closed pull request #6015: Allow customizable cursor init params
URL: https://github.com/apache/incubator-superset/pull/6015
 
 
   

This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:

As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):

diff --git a/superset/config.py b/superset/config.py
index 80f6f85e23..369880ba5b 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -291,6 +291,12 @@
 # not the production version of the site.
 WARNING_MSG = None
 
+# The source label for sqllab pyhive connection initialization.
+SQLLAB_PYHIVE_SOURCE = None
+
+# The source label for chart pyhive connection initialization.
+CHART_PYHIVE_SOURCE = None
+
 # Default celery config is to use SQLA as a broker, in a production setting
 # you'll want to use a proper broker as specified here:
 # http://docs.celeryproject.org/en/latest/getting-started/brokers/index.html
diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 6c8e07061f..50b99d5910 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -375,6 +375,14 @@ def execute(cls, cursor, query, **kwargs):
             cursor.arraysize = cls.arraysize
         cursor.execute(query)
 
+    @classmethod
+    def custom_sqllab_cursor_params(cls):
+        pass
+
+    @classmethod
+    def custom_chart_cursor_params(cls):
+        pass
+
     @classmethod
     def make_label_compatible(cls, label):
         """
@@ -889,6 +897,14 @@ def latest_partition(cls, table_name, schema, database, show_first=False):
         df = database.get_df(sql, schema)
         return part_field, cls._latest_partition_from_df(df)
 
+    @classmethod
+    def custom_sqllab_cursor_params(cls):
+        return {'source': config['SQLLAB_PYHIVE_SOURCE']}
+
+    @classmethod
+    def custom_chart_cursor_params(cls):
+        return {'source': config['CHART_PYHIVE_SOURCE']}
+
     @classmethod
     def latest_sub_partition(cls, table_name, schema, database, **kwargs):
         """Returns the latest (max) partition value for a table
diff --git a/superset/models/core.py b/superset/models/core.py
index 52a005b478..091f616d8a 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -733,7 +733,7 @@ def get_effective_user(self, url, user_name=None):
 
     @utils.memoized(
         watch=('impersonate_user', 'sqlalchemy_uri_decrypted', 'extra'))
-    def get_sqla_engine(self, schema=None, nullpool=True, user_name=None):
+    def get_sqla_engine(self, schema=None, nullpool=True, user_name=None, sqllab=False):
         extra = self.get_extra()
         url = make_url(self.sqlalchemy_uri_decrypted)
         url = self.db_engine_spec.adjust_database_uri(url, schema)
@@ -767,6 +767,20 @@ def get_sqla_engine(self, schema=None, nullpool=True, user_name=None):
         if DB_CONNECTION_MUTATOR:
             url, params = DB_CONNECTION_MUTATOR(
                 url, params, effective_username, security_manager)
+
+        custom_url_params = None
+
+        if (sqllab and self.db_engine_spec.custom_sqllab_cursor_params()):
+            custom_url_params = parse.urlencode(
+                self.db_engine_spec.custom_sqllab_cursor_params())
+
+        if ((not sqllab) and self.db_engine_spec.custom_chart_cursor_params()):
+            custom_url_params = parse.urlencode(
+                self.db_engine_spec.custom_chart_cursor_params())
+
+        if custom_url_params:
+            url = '{}?{}'.format(url, custom_url_params)
+
         return create_engine(url, **params)
 
     def get_reserved_words(self):
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index a2732d17c5..8f235cdcda 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -176,6 +176,7 @@ def handle_error(msg):
             schema=query.schema,
             nullpool=True,
             user_name=user_name,
+            sqllab=True,
         )
         conn = engine.raw_connection()
         cursor = conn.cursor()


 

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org