add session struct in proto

.gitignore

@@ -17,3 +17,4 @@ __pycache__/
 *.log
 
 .DS_Store
+.vscode

proto/sqlflow/proto/sqlflow.proto

@@ -19,13 +19,19 @@ service SQLFlow {
     rpc Run (Request) returns (stream Response);
 }
 
+message Session {
+    string token = 1;
+    string db_conn_str = 2;
+}
+
 // SQL statements to run
 // e.g.
 //      1. `SELECT ...`
 //      2. `USE ...`, `DELETE ...`
 //      3. `SELECT ... TRAIN/PREDICT ...`
 message Request {
-    string sql = 1;		// The SQL statement to be executed.
+    string sql = 1;		    // The SQL statement to be executed.
+    Session session = 2;    // The Session struct including user credentical message. 
 }
 
 message Response {

sqlflow/client.py

@@ -92,9 +92,9 @@ def __init__(self, server_url=None, ca_crt=None):
                 raise ValueError("Can't find environment variable SQLFLOW_SERVER")
             server_url = os.environ["SQLFLOW_SERVER"]
 
-        self._stub = pb_grpc.SQLFlowStub(self.newRPCChannel(server_url, ca_crt))
+        self._stub = pb_grpc.SQLFlowStub(self.new_rpc_channel(server_url, ca_crt))
 
-    def newRPCChannel(self, server_url, ca_crt):
+    def new_rpc_channel(self, server_url, ca_crt):
         if ca_crt is None and "SQLFLOW_CA_CRT" not in os.environ:
             # client would connect SQLFLow gRPC server with insecure mode.
             channel = grpc.insecure_channel(server_url) 
@@ -106,6 +106,12 @@ def newRPCChannel(self, server_url, ca_crt):
             channel = grpc.secure_channel(server_url, creds)
         return channel
 
+    def sql_request(self, sql):
+        token = os.getenv("SQLFLOW_USER_TOKEN", "")
+        db_conn_str = os.getenv("SQLFLOW_DATASOURCE", "")
+        se = pb.Session(token=token, db_conn_str=db_conn_str)
+        return pb.Request(sql=sql, session=se)
+
     def execute(self, operation):
         """Run a SQL statement
 
@@ -120,7 +126,7 @@ def execute(self, operation):
 
         """
         try:
-            stream_response = self._stub.Run(pb.Request(sql=operation))
+            stream_response = self._stub.Run(self.sql_request(operation))
             return self.display(stream_response)
         except grpc.RpcError as e:
             _LOGGER.error("%s\n%s", e.code(), e.details())

tests/mock_servicer.py

@@ -20,6 +20,9 @@ def Run(self, request, context):
             else:
                 for res in MockServicer.table_response(MockServicer.get_test_table()):
                     yield res
+        elif SQL == "TEST VERIFY SESSION":
+            # TODO(Yancey1989): using a elegant way to test the session instead of the trick.
+            yield MockServicer.message_response("|".join([request.session.token, request.session.db_conn_str]))
         else:
             yield MockServicer.message_response('bad request', 0)
 

tests/test_client.py

@@ -67,3 +67,12 @@ def test_decode_null(self):
         null_message = pb.Row.Null()
         any_message.Pack(null_message)
         assert Client._decode_any(any_message) is None
+
+    def test_session(self):
+        token = "unittest-user"
+        ds = "maxcompute://AK:SK@host:port" 
+        os.environ["SQLFLOW_USER_TOKEN"] = token
+        os.environ["SQLFLOW_DATASOURCE"] = ds
+        with mock.patch('sqlflow.client._LOGGER') as log_mock:
+            self.client.execute("TEST VERIFY SESSION")
+            log_mock.info.assert_called_with("|".join([token, ds]))