pyalisa (#3107)

* pyalisa support py2

* add task.py

* implement task.py

* move from_env from client to config module

* fix response field name

* add env:SKYNET_TENANT_ID

* add test for task

* remove duplicate test case

* fix config en/decoding

* add test case for task

* fix unittest for alisa.config

Author: weiguoz

python/runtime/dbapi/pyalisa/client.py

@@ -11,14 +11,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
+import ast
 import json
-import os
 import random
 import string
 import time
 from enum import Enum
 
-from runtime.dbapi.pyalisa.config import Config
 from runtime.dbapi.pyalisa.pop import Pop
 
 
@@ -35,15 +34,15 @@ class AlisaTaksStatus(Enum):
     ALISA_TASK_ALLOCATE = 11
 
 
-# used to deal with too many logs.
+# used to deal with too many logs
 MAX_LOG_NUM = 2000
 
 
 class Client(object):
     """Client for building kinds of tasks and submitting them to alisa gateway
 
     Args:
-        config(Config): the config for build the client
+        config(Config): the Config(runtime.dbapi.pyalisa.config)
     """
     def __init__(self, config):
         self.config = config  # noqa F841
@@ -93,7 +92,6 @@ def create_pyodps_task(self, code, args):
         params["Exec"] = self.config.withs["Exec4PyODPS"]
         if len(args) > 0:
             params["Args"] = args
-
         return self._create_task(params)
 
     def _create_task(self, params):
@@ -168,12 +166,12 @@ def read_logs(self, task_id, offset, w):
             params["AlisaTaskId"] = task_id
             params["Offset"] = str(offset)
             log = self._requet_and_parse_response("GetAlisaTaskLog", params)
-            rlen = int(log["ReadLen"])
+            rlen = int(log["readLength"])
             if rlen == 0:
                 return offset
             offset += rlen
-            w.write(log["Content"])
-            if bool(log["End"]):
+            w.write(log["logMsg"])
+            if bool(log["isEnd"]):
                 return -1
         return offset
 
@@ -205,20 +203,23 @@ def get_results(self, task_id, batch):
         if batch <= 0:
             raise ValueError("batch should greater than 0")
         count = self.count_results(task_id)
-        result = []
+
+        columns, body = [], []
         for i in range(0, count, batch):
             params = self._base_params()
             params["AlisaTaskId"] = task_id
             params["Start"] = str(i)
             params["Limit"] = str(batch)
-            r = self._requet_and_parse_response("GetAlisaTaskResult", params)
-            # TODO(lhw): parse the result like:
-            # https://github.com/sql-machine-learning/goalisa/blob/68d3aad1344c9e5c0cd35c6556e1f3f2b6ca9db7/alisa.go#L190
-            result.append(r)
-        return result
+            val = self._requet_and_parse_response("GetAlisaTaskResult", params)
+            header, rows = self._parse_alisa_value(val)
+            if len(columns) == 0:
+                columns = header
+            body.extend(rows)
+        return {"columns": columns, "body": body}
 
     def stop(self, task_id):
-        """Stop given task
+        """Stop given task.
+        NOTE(weiguoz): need to be tested.
 
         Args:
             task_id(string): the task to stop
@@ -231,51 +232,38 @@ def stop(self, task_id):
         res = self._requet_and_parse_response("StopAlisaTask", params)
         return bool(res)
 
+    def _parse_alisa_value(self, val):
+        """Parse 'returnValue' in alisa response
+        https://github.com/sql-machine-learning/goalisa/blob/68d3aad1344c9e5c0cd35c6556e1f3f2b6ca9db7/alisa.go#L190
+
+        Args:
+            val: [{u'resultMsg': u'[["Alice","23.8","56000"]]',
+            u'dataHeader': u'["name::string","age::double","salary::bigint"]'}]
+        """
+        jsval = ast.literal_eval(json.dumps(val))
+        columns = []
+        for h in json.loads(jsval['dataHeader']):
+            nt = h.split("::")
+            name, typ = (nt[0], nt[1]) if len(nt) == 2 else (h, "string")
+            columns.append({"name": str(name), "typ": str(typ)})
+        body = []
+        for m in json.loads(jsval['resultMsg']):
+            row = []
+            for i in ast.literal_eval(json.dumps(m)):
+                row.append(i)
+            body.append(row)
+        return columns, body
+
     def _requet_and_parse_response(self, action, params):
         params["Action"] = action
         params["ProjectEnv"] = self.config.env["SKYNET_SYSTEM_ENV"]
         url = self.config.pop_scheme + "://" + self.config.pop_url
         code, buf = Pop.request(url, params, self.config.pop_access_secret)
         resp = json.loads(buf)
         if code != 200:
-            raise RuntimeError("%s got a bad result, response=%s" %
-                               (code, buf))
+            raise RuntimeError("%s got a bad result, request=%s, response=%s" %
+                               (code, params, buf))
+        if resp['returnCode'] != '0':
+            raise Exception("returned an error request={}, response={}".format(
+                params, resp))
         return resp["returnValue"]
-
-    @staticmethod
-    def from_env():
-        """Build a Client from environment variable
-
-        Returns:
-            a Client instance
-        """
-        if not os.getenv("POP_SECRET"):
-            return None
-        conf = Config()
-        conf.pop_url = os.getenv("POP_URL")
-        conf.pop_access_id = os.getenv("POP_ID")
-        conf.pop_access_secret = os.getenv("POP_SECRET")
-        conf.pop_scheme = "http"
-        conf.verbose = os.getenv("VERBOSE") == "true"
-        conf.env = {
-            "SKYNET_ONDUTY": os.getenv("SKYNET_ONDUTY"),
-            "SKYNET_ACCESSID": os.getenv("SKYNET_ACCESSID"),
-            "SKYNET_ACCESSKEY": os.getenv("SKYNET_ACCESSKEY"),
-            "SKYNET_ENDPOINT": os.getenv("SKYNET_ENDPOINT"),
-            "SKYNET_SYSTEMID": os.getenv("SKYNET_SYSTEMID"),
-            "SKYNET_PACKAGEID": os.getenv("SKYNET_PACKAGEID"),
-            "SKYNET_SYSTEM_ENV": os.getenv("SKYNET_SYSTEM_ENV"),
-            "SKYNET_BIZDATE": os.getenv("SKYNET_BIZDATE"),
-            "ALISA_TASK_EXEC_TARGET": os.getenv("ALISA_TASK_EXEC_TARGET"),
-        }
-        conf.withs = {
-            "CustomerId": os.getenv("CustomerId"),
-            "PluginName": os.getenv("PluginName"),
-            "Exec": os.getenv("Exec"),
-            "PluginName4PyODPS": os.getenv("PluginName4PyODPS"),
-            "Exec4PyODPS": os.getenv("Exec4PyODPS"),
-        }
-        conf.curr_project = conf.env["SKYNET_PACKAGEID"]
-        if len(conf.env["SKYNET_SYSTEMID"]) > 0:
-            conf.curr_project += "_" + conf.env["SKYNET_SYSTEMID"]
-        return Client(conf)

python/runtime/dbapi/pyalisa/client_test.py

@@ -16,12 +16,13 @@
 
 from runtime import testing
 from runtime.dbapi.pyalisa.client import AlisaTaksStatus, Client
+from runtime.dbapi.pyalisa.config import Config
 
 
 @unittest.skipUnless(testing.get_driver() == "alisa", "Skip non-alisa test")
 class TestClient(unittest.TestCase):
     def test_create_sql_task(self):
-        ali = Client.from_env()
+        ali = Client(Config.from_env())
         code = "SELECT 2;"
         task_id, _ = ali.create_sql_task(code)
         self.assertIsNotNone(task_id)
@@ -31,6 +32,7 @@ def test_create_sql_task(self):
 
         # try get result
         for _ in range(10):
+            # to avoid touching the flow-control
             time.sleep(1)
             status = ali.get_status(task_id)
             if ali.completed(status):
@@ -40,24 +42,6 @@ def test_create_sql_task(self):
                 self.assertGreater(len(result), 0)
                 break
 
-    def test_create_pyodps_task(self):
-        ali = Client.from_env()
-        code = """import argparse
-if __name__ == "__main__":
-    input_table_name = args['input_table']
-    output_table_name = args['output_table']
-    print(input_table_name)
-    print(output_table_name)
-    input_table = o.get_table(input_table_name)
-    print(input_table.schema)
-    output_table = o.create_table(output_table_name, input_table.schema)
-        """
-        args = "input_table=table_1 output_table=table_2"
-        task_id, _ = ali.create_pyodps_task(code, args)
-        # to avoid touching the flow-control
-        time.sleep(2)
-        self.assertIsNotNone(task_id)
-
 
 if __name__ == "__main__":
     unittest.main()

python/runtime/dbapi/pyalisa/config.py

@@ -13,6 +13,7 @@
 
 import base64
 import json
+import os
 from collections import OrderedDict
 
 from six.moves.urllib.parse import parse_qs, urlparse
@@ -66,7 +67,7 @@ def _encode_json_base64(env):
     @staticmethod
     def _decode_json_base64(b64env):
         padded = b64env + "=" * (len(b64env) % 4)
-        jstr = base64.urlsafe_b64decode(padded).decode("utf8")
+        jstr = base64.urlsafe_b64decode(padded.encode('utf8')).decode("utf8")
         return json.loads(jstr)
 
     def to_url(self):
@@ -87,3 +88,42 @@ def to_url(self):
         )
         return ("alisa://%s:%s@%s?scheme=%s&verbose"
                 "=%s&curr_project=%s&env=%s&with=%s") % parts
+
+    @staticmethod
+    def from_env():
+        """Build a Client from environment variable
+
+        Returns:
+            a Client instance
+        """
+        if not os.getenv("POP_SECRET"):
+            return None
+        conf = Config()
+        conf.pop_url = os.getenv("POP_URL")
+        conf.pop_access_id = os.getenv("POP_ID")
+        conf.pop_access_secret = os.getenv("POP_SECRET")
+        conf.pop_scheme = "http"
+        conf.verbose = os.getenv("VERBOSE") == "true"
+        conf.env = {
+            "SKYNET_ONDUTY": os.getenv("SKYNET_ONDUTY"),
+            "SKYNET_ACCESSID": os.getenv("SKYNET_ACCESSID"),
+            "SKYNET_ACCESSKEY": os.getenv("SKYNET_ACCESSKEY"),
+            "SKYNET_ENDPOINT": os.getenv("SKYNET_ENDPOINT"),
+            "SKYNET_SYSTEMID": os.getenv("SKYNET_SYSTEMID"),
+            "SKYNET_PACKAGEID": os.getenv("SKYNET_PACKAGEID"),
+            "SKYNET_SYSTEM_ENV": os.getenv("SKYNET_SYSTEM_ENV"),
+            "SKYNET_BIZDATE": os.getenv("SKYNET_BIZDATE"),
+            "SKYNET_TENANT_ID": os.getenv("SKYNET_TENANT_ID"),
+            "ALISA_TASK_EXEC_TARGET": os.getenv("ALISA_TASK_EXEC_TARGET"),
+        }
+        conf.withs = {
+            "CustomerId": os.getenv("CustomerId"),
+            "PluginName": os.getenv("PluginName"),
+            "Exec": os.getenv("Exec"),
+            "PluginName4PyODPS": os.getenv("PluginName4PyODPS"),
+            "Exec4PyODPS": os.getenv("Exec4PyODPS"),
+        }
+        conf.curr_project = conf.env["SKYNET_PACKAGEID"]
+        if len(conf.env["SKYNET_SYSTEMID"]) > 0:
+            conf.curr_project += "_" + conf.env["SKYNET_SYSTEMID"]
+        return conf

python/runtime/dbapi/pyalisa/config_test.py

@@ -13,6 +13,7 @@
 
 import unittest
 
+from runtime import testing
 from runtime.dbapi.pyalisa.config import Config
 
 test_url = ("alisa://pid:[email protected]/?scheme=http&verbose=true&"
@@ -27,13 +28,16 @@
             "sICJQbHVnaW5OYW1lIjogIndwZSIsICJDdXN0b21lcklkIjogIndjZCJ9")
 
 
+@unittest.skipUnless(testing.get_driver() == "alisa", "Skip non-alisa test")
 class TestConfig(unittest.TestCase):
+    """We use python2 in alisa, so let's skip the tests in the other drivers.
+    """
     def test_encode_json_base64(self):
         params = dict()
         params["key1"] = "val1"
         params["key2"] = "val2"
         b64 = Config._encode_json_base64(params)
-        self.assertEqual("eyJrZXkxIjogInZhbDEiLCAia2V5MiI6ICJ2YWwyIn0", b64)
+        self.assertEqual("eyJrZXkyIjogInZhbDIiLCAia2V5MSI6ICJ2YWwxIn0", b64)
 
         params = Config._decode_json_base64(b64)
         self.assertEqual(2, len(params))
@@ -52,9 +56,18 @@ def test_dsn_parsing(self):
         self.assertEqual("SKY", cfg.env["SKYNET_ACCESSKEY"])
 
     def test_to_dsn(self):
-        cfg = Config(test_url)
-        url = cfg.to_url()
-        self.assertEqual(test_url, url)
+        c1 = Config(test_url)
+        u1 = c1.to_url()
+        c2 = Config(u1)
+        self.assertEqual(c1.pop_access_id, c2.pop_access_id)
+        self.assertEqual(c1.pop_access_secret, c2.pop_access_secret)
+        self.assertEqual(c1.curr_project, c2.curr_project)
+        self.assertEqual(c1.scheme, c2.scheme)
+        self.assertEqual(c1.withs["CustomerId"], c2.withs["CustomerId"])
+        self.assertEqual(c1.withs["PluginName"], c2.withs["PluginName"])
+        self.assertEqual(c1.withs["Exec"], c2.withs["Exec"])
+        self.assertEqual(c1.env["SKYNET_ACCESSKEY"],
+                         c2.env["SKYNET_ACCESSKEY"])
 
     def test_parse_error(self):
         # no env and with

python/runtime/dbapi/pyalisa/pop.py

@@ -14,10 +14,9 @@
 import base64
 import hashlib
 import hmac
-from encodings import utf_8
+import urllib
 
 import requests
-from six.moves import urllib
 
 
 class Pop(object):
@@ -59,10 +58,9 @@ def signature(params, http_method, secret):
             ek = Pop.percent_encode(k)
             ev = Pop.percent_encode(v)
             qry += '&{}={}'.format(ek, ev)
-        to_sign = '{}&%2F&{}'.format(http_method, Pop.percent_encode(qry[1:]))
-        bf = bytearray(utf_8.encode(to_sign)[0])
-        dig = hmac.digest(utf_8.encode(secret + "&")[0], bf, hashlib.sha1)
-        return utf_8.decode(base64.standard_b64encode(dig))[0]
+        str = '{}&%2F&{}'.format(http_method, Pop.percent_encode(qry[1:]))
+        dig = hmac.new(secret + '&', str, hashlib.sha1).digest()
+        return base64.standard_b64encode(dig)
 
     @staticmethod
     def percent_encode(str):
@@ -75,7 +73,7 @@ def percent_encode(str):
         Returns:
             (string) the encoded param
         """
-        es = urllib.parse.quote_plus(str)
-        es = es.replace('+', '%20')
-        es = es.replace('*', '%2A')
-        return es.replace('%7E', '~')
+        s = urllib.quote_plus(str)
+        s = s.replace('+', '%20')
+        s = s.replace('*', '%2A')
+        return s.replace('%7E', '~')