Merge pull request #41 from dataiku/task/list-tables-in-public-api

add calls to fetch SQL/Hive schemas/tables [ch34607]

Author: FChataigner

dataikuapi/dss/future.py

@@ -10,6 +10,15 @@ def __init__(self, client, job_id, state=None):
        self.state = state
        self.state_is_peek = True
 
+    @classmethod
+    def get_result_wait_if_needed(cls, client, ret):
+        if 'jobId' in ret:
+            future = DSSFuture(client, ret["jobId"], ret)
+            future.wait_for_result()
+            return future.get_result()
+        else:
+            return ret['result']
+
     def abort(self):
         """
         Abort the future

dataikuapi/dss/project.py

@@ -860,6 +860,53 @@ def init_tables_import(self):
         """
         return TablesImportDefinition(self.client, self.project_key)
 
+    def list_sql_schemas(self, connection_name):
+        """
+        Lists schemas from which tables can be imported in a SQL connection
+
+        :returns: an array of schemas names
+        """
+        return self._list_schemas(connection_name)
+
+    def list_hive_databases(self):
+        """
+        Lists Hive databases from which tables can be imported
+
+        :returns: an array of databases names
+        """
+        return self._list_schemas("@virtual(hive-jdbc):default")
+
+    def _list_schemas(self, connection_name):
+        return self.client._perform_json("GET", "/projects/%s/datasets/tables-import/actions/list-schemas" % (self.project_key),
+                params = {"connectionName": connection_name} )
+
+    def list_sql_tables(self, connection_name, schema_name=None):
+        """
+        Lists tables to import in a SQL connection
+
+        :returns: an array of tables
+        """
+        ret = self.client._perform_json("GET", "/projects/%s/datasets/tables-import/actions/list-tables" % (self.project_key),
+                params = {"connectionName": connection_name, "schemaName": schema_name} )
+
+        def to_schema_table_pair(x):
+            return {"schema":x.get("schema", None), "table":x["table"]}
+        return [to_schema_table_pair(x) for x in DSSFuture.get_result_wait_if_needed(self.client, ret)['tables']]
+
+    def list_hive_tables(self, hive_database):
+        """
+        Lists tables to import in a Hive database
+
+        :returns: an array of tables
+        """
+        connection_name = "@virtual(hive-jdbc):" + hive_database
+        ret = self.client._perform_json("GET", "/projects/%s/datasets/tables-import/actions/list-tables" % (self.project_key),
+                params = {"connectionName": connection_name} )
+
+        def to_schema_table_pair(x):
+            return {"schema":x.get("databaseName", None), "table":x["table"]}
+        return [to_schema_table_pair(x) for x in DSSFuture.get_result_wait_if_needed(self.client, ret)['tables']]
+
 class TablesImportDefinition(object):
     """
     Temporary structure holding the list of tables to import