Selaa lähdekoodia

[flink] Add default_catalog and default_database option (#4015)

- `default_catalog` and `default_database` parameters make it possible to switch to desired Flink catalog and database automatically after the session is created
Grzegorz Kołakowski 8 kuukautta sitten
vanhempi
commit
4781ba6990
1 muutettua tiedostoa jossa 22 lisäystä ja 0 poistoa
  1. 22 0
      desktop/libs/notebook/src/notebook/connectors/flink_sql.py

+ 22 - 0
desktop/libs/notebook/src/notebook/connectors/flink_sql.py

@@ -33,6 +33,8 @@ _JSON_CONTENT_TYPE = 'application/json'
 _API_VERSION = 'v3'
 SESSION_KEY = '%(username)s-%(connector_name)s'
 OPERATION_TOKEN = '%(username)s-%(connector_name)s' + '-operation-token'
+DEFAULT_CATALOG_PARAM = "default_catalog"
+DEFAULT_DATABASE_PARAM = "default_database"
 
 
 def query_error_handler(func):
@@ -70,6 +72,8 @@ class FlinkSqlApi(Api):
 
     self.options = interpreter['options']
     api_url = self.options['url']
+    self.default_catalog = self.options.get(DEFAULT_CATALOG_PARAM)
+    self.default_database = self.options.get(DEFAULT_DATABASE_PARAM)
 
     self.db = FlinkSqlClient(user=user, api_url=api_url)
 
@@ -159,6 +163,10 @@ class FlinkSqlApi(Api):
 
     if not session:
       session = self.db.create_session()
+      if self.default_database:
+        self._use_database(self.default_catalog, self.default_database)
+      elif self.default_catalog:
+        self._use_catalog(self.default_catalog)
 
     try:
       self.db.session_heartbeat(session_handle=session['sessionHandle'])
@@ -453,6 +461,20 @@ class FlinkSqlApi(Api):
 
     return [{'name': function[0]} for function in function_list]
 
+  def _use_catalog(self, catalog):
+    session = self._get_session()
+    self.db.configure_session(session_handle=(session['id']), statement="USE CATALOG `%s`" % catalog)
+
+  def _use_database(self, catalog, database):
+    session = self._get_session()
+    if catalog:
+      self.db.configure_session(session_handle=(session['id']),
+                                statement="USE `%(catalog)s`.`%(database)s`" % {'catalog': catalog,
+                                                                                'database': database})
+    else:
+      self.db.configure_session(session_handle=(session['id']),
+                                statement="USE `%(database)s`" % {'database': database})
+
 
 class FlinkSqlClient:
   """