")
60 | else:
61 | yield MockServicer.message_response('bad request', 0)
62 |
63 | def Fetch(self, request, context):
64 | job_id = request.job.id
65 | yield MockServicer.message_fetch_response(request.job, "execute workflow")
66 |
67 | @staticmethod
68 | def get_test_table():
69 | return {"column_names": ['x', 'y'], "rows": [[1, 2], [3, 4]]}
70 |
71 | @staticmethod
72 | def wrap_value(value):
73 | if isinstance(value, bool):
74 | message = wrapper.BoolValue()
75 | message.value = value
76 | elif isinstance(value, int):
77 | message = wrapper.Int64Value()
78 | message.value = value
79 | elif isinstance(value, float):
80 | message = wrapper.DoubleValue()
81 | message.value = value
82 | else:
83 | raise Exception("Unsupported type {}".format(type(value)))
84 | return message
85 |
86 | @staticmethod
87 | def table_response(table):
88 | res = pb.Response()
89 | head = pb.Head()
90 | for name in table['column_names']:
91 | head.column_names.append(name)
92 | res.head.CopyFrom(head)
93 | yield res
94 |
95 | for row in table['rows']:
96 | res = pb.Response()
97 | row_message = pb.Row()
98 | for data in row:
99 | row_message.data.add().Pack(MockServicer.wrap_value(data))
100 | res.row.CopyFrom(row_message)
101 | yield res
102 |
103 | @staticmethod
104 | def message_response(message):
105 | pb_msg = pb.Message()
106 | pb_msg.message = message
107 |
108 | res = pb.Response()
109 | res.message.CopyFrom(pb_msg)
110 | return res
111 |
112 | @staticmethod
113 | def eoe_response():
114 | pb_eoe = pb.EndOfExecution()
115 | res = pb.Response()
116 | res.eoe.CopyFrom(pb_eoe)
117 | return res
118 |
119 |
120 | def _server(server_instance, port, event, ca_crt, ca_key):
121 | svr = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
122 | with open(ca_key, "rb") as f:
123 | private_key = f.read()
124 | with open(ca_crt, "rb") as f:
125 | certification_chain = f.read()
126 | server_credentials = grpc.ssl_server_credentials( ( (private_key, certification_chain), ) )
127 | pb_grpc.add_SQLFlowServicer_to_server(server_instance, svr)
128 | svr.add_secure_port('[::]:%d' % port, server_credentials)
129 | svr.start()
130 | try:
131 | event.wait()
132 | except KeyboardInterrupt:
133 | pass
134 | svr.stop(0)
135 |
--------------------------------------------------------------------------------
/tests/test_client.py:
--------------------------------------------------------------------------------
1 | import unittest
2 | import threading
3 | import time
4 | import os
5 | from unittest import mock
6 | import tempfile
7 | import subprocess
8 | import shutil
9 |
10 | from sqlflow.client import Client
11 | from tests.mock_servicer import _server, MockServicer, MockWorkflowServicer
12 |
13 | from google.protobuf.timestamp_pb2 import Timestamp
14 | from google.protobuf.any_pb2 import Any
15 | import sqlflow.proto.sqlflow_pb2 as pb
16 |
17 | def generateTempCA():
18 | tmp_dir = tempfile.mkdtemp(suffix="sqlflow_ssl", dir="/tmp")
19 | ca_key = os.path.join(tmp_dir, "ca.key")
20 | ca_csr = os.path.join(tmp_dir, "ca.csr")
21 | ca_crt = os.path.join(tmp_dir, "ca.crt")
22 |
23 | assert subprocess.call(["openssl", "genrsa", "-out", ca_key, "2048"]) == 0
24 | assert subprocess.call(["openssl", "req", "-nodes", "-new", "-key", ca_key, "-subj", "/CN=localhost", "-out", ca_csr]) == 0
25 | assert subprocess.call(["openssl", "x509", "-req", "-sha256", "-days", "365", "-in", ca_csr, "-signkey", ca_key, "-out", ca_crt]) == 0
26 |
27 | return tmp_dir, ca_crt, ca_key
28 |
29 |
30 | class ClientServerTest(unittest.TestCase):
31 | @classmethod
32 | def setUpClass(cls):
33 | # TODO: free port is better
34 | port = 8765
35 | cls.server_url = "localhost:%d" % port
36 | cls.event = threading.Event()
37 | cls.tmp_ca_dir, cls.ca_crt, ca_key = generateTempCA()
38 | threading.Thread(target=_server, args=[MockServicer(), port, cls.event, cls.ca_crt, ca_key]).start()
39 | # wait for start
40 | time.sleep(1)
41 | cls.client = Client(cls.server_url, cls.ca_crt)
42 |
43 | @classmethod
44 | def tearDownClass(cls):
45 | # shutdown server after this test
46 | cls.event.set()
47 | shutil.rmtree(cls.tmp_ca_dir, ignore_errors=True)
48 |
49 | def test_execute_stream(self):
50 | with mock.patch('sqlflow.client._LOGGER') as log_mock:
51 | res = self.client.execute("select * from galaxy train ..")
52 | log_mock.info.assert_called_with("extended sql")
53 |
54 | expected_table = MockServicer.get_test_table()
55 | rows = self.client.execute("select * from galaxy").get(0)
56 | assert expected_table["column_names"] == rows.column_names()
57 | assert expected_table["rows"] == [r for r in rows.rows()]
58 |
59 | def test_cmd(self):
60 | assert subprocess.call(["sqlflow", "--url", self.server_url,
61 | "--ca_crt", self.ca_crt,
62 | "select * from galaxy"]) == 0
63 |
64 | def test_decode_time(self):
65 | any_message = Any()
66 | timestamp_message = Timestamp()
67 | timestamp_message.GetCurrentTime()
68 | any_message.Pack(timestamp_message)
69 | assert timestamp_message.ToDatetime() == Client._decode_any(any_message)
70 |
71 | def test_decode_null(self):
72 | any_message = Any()
73 | null_message = pb.Row.Null()
74 | any_message.Pack(null_message)
75 | assert Client._decode_any(any_message) is None
76 |
77 | # TODO(typhoonzero): this test seems not useful, need to find a better way
78 | # def test_session(self):
79 | # token = "unittest-user"
80 | # ds = "maxcompute://AK:SK@host:port"
81 | # os.environ["SQLFLOW_USER_TOKEN"] = token
82 | # os.environ["SQLFLOW_DATASOURCE"] = ds
83 | # os.environ["SQLFLOW_EXIT_ON_SUBMIT"] = "True"
84 | # os.environ["SQLFLOW_USER_ID"] = "sqlflow_user"
85 | # os.environ["SQLFLOW_SUBMITTER"] = "pai"
86 | # with mock.patch('sqlflow.client._LOGGER') as log_mock:
87 | # self.client.execute("TEST VERIFY SESSION")
88 | # log_mock.debug.assert_called_with("|".join([token, ds, "True", "sqlflow_user"]))
89 |
90 | def test_draw_html(self):
91 | from IPython.core.display import display, HTML
92 | with mock.patch('IPython.core.display.HTML') as log_mock:
93 | self.client.execute("TEST RENDER HTML")
94 | log_mock.assert_called_with("
")
95 |
96 | class ClientWorkflowServerTest(unittest.TestCase):
97 | @classmethod
98 | def setUpClass(cls):
99 | # TODO: free port is better
100 | port = 8766
101 | cls.server_url = "localhost:%d" % port
102 | cls.event = threading.Event()
103 | cls.tmp_ca_dir, cls.ca_crt, ca_key = generateTempCA()
104 | threading.Thread(target=_server, args=[MockWorkflowServicer(), port, cls.event, cls.ca_crt, ca_key]).start()
105 | # wait for start
106 | time.sleep(1)
107 | cls.client = Client(cls.server_url, cls.ca_crt)
108 |
109 | @classmethod
110 | def tearDownClass(cls):
111 | # shutdown server after this test
112 | cls.event.set()
113 | shutil.rmtree(cls.tmp_ca_dir, ignore_errors=True)
114 |
115 | def test_execute_stream(self):
116 | expected_table = MockServicer.get_test_table()
117 | rows = self.client.execute("select * from galaxy train ..").get(0)
118 | assert expected_table["column_names"] == rows.column_names()
119 | assert expected_table["rows"] == [r for r in rows.rows()]
120 |
--------------------------------------------------------------------------------
/tests/test_env_expand.py:
--------------------------------------------------------------------------------
1 | import os
2 | import unittest
3 | from sqlflow.env_expand import EnvExpander, EnvExpanderError
4 |
5 | class EnvExpanderTest(unittest.TestCase):
6 | @classmethod
7 | def setUpClass(cls):
8 | os.environ["BIZDATE"] = "20190731"
9 | os.environ["t1"] = "tablename"
10 | cls.expander = EnvExpander(os.environ)
11 |
12 | def test_expand(self):
13 | sql = "SELECT * from ${t1} where pt=${yyyyMMdd}"
14 | new_sql = self.expander.expand(sql)
15 | expected_sql = "SELECT * from tablename where pt=20190731"
16 | assert new_sql == expected_sql
17 |
18 | def test_bizdate_delta(self):
19 | sql = "SELECT * from ${t1} where pt=${yyyyMMdd + 1d}"
20 | new_sql = self.expander.expand(sql)
21 | expected_sql = "SELECT * from tablename where pt=20190801"
22 | assert new_sql == expected_sql
23 |
24 | def test_bizdate_arbitrary_space(self):
25 | sql = "SELECT * from ${t1} where pt=${yyyyMMdd+ 1d}"
26 | new_sql = self.expander.expand(sql)
27 | expected_sql = "SELECT * from tablename where pt=20190801"
28 | assert new_sql == expected_sql
29 |
30 | def test_expand_error(self):
31 | sql = "SELECT * from ${no_exists} where pt=${yyyyMMdd + 1d}"
32 | with self.assertRaises(EnvExpanderError):
33 | self.expander.expand(sql)
34 |
35 |
--------------------------------------------------------------------------------
/tests/test_version.py:
--------------------------------------------------------------------------------
1 | import sqlflow
2 |
3 |
4 | def test_answer():
5 | assert sqlflow.__version__ == sqlflow._version.__version__
6 |
--------------------------------------------------------------------------------