[client] KUDU-3472 Python client API to import JWT

This patch adds JWT to the Python client builder and to the
kudu.connect() methods.

Change-Id: Icaef1bd28efe9d47252bd211fb937f2e6e48cea9
Reviewed-on: http://gerrit.cloudera.org:8080/19862
Tested-by: Kudu Jenkins
Reviewed-by: Alexey Serbin <alexey@apache.org>
Reviewed-by: Zoltan Chovan <zchovan@cloudera.com>
diff --git a/python/kudu/__init__.py b/python/kudu/__init__.py
index 200d5ce..7c90518 100644
--- a/python/kudu/__init__.py
+++ b/python/kudu/__init__.py
@@ -61,7 +61,8 @@
 
 
 def connect(host, port=7051, admin_timeout_ms=None, rpc_timeout_ms=None,
-            require_authentication=False, encryption_policy=ENCRYPTION_OPTIONAL):
+            require_authentication=False, encryption_policy=ENCRYPTION_OPTIONAL,
+            jwt=None):
     """
     Connect to a Kudu master server
 
@@ -80,6 +81,8 @@
       Whether to require authentication
     encryption_policy : enum, optional
       Whether to require encryption
+    jwt : string, optional
+      The JSON web token to set.
 
     Returns
     -------
@@ -105,7 +108,7 @@
     return Client(addresses, admin_timeout_ms=admin_timeout_ms,
                   rpc_timeout_ms=rpc_timeout_ms,
                   encryption_policy=encryption_policy,
-                  require_authentication=require_authentication)
+                  require_authentication=require_authentication, jwt=jwt)
 
 
 def timedelta(seconds=0, millis=0, micros=0, nanos=0):
diff --git a/python/kudu/client.pyx b/python/kudu/client.pyx
index 78c2092..2cdb539 100644
--- a/python/kudu/client.pyx
+++ b/python/kudu/client.pyx
@@ -293,7 +293,7 @@
     def __cinit__(self, addr_or_addrs, admin_timeout_ms=None,
                   rpc_timeout_ms=None, sasl_protocol_name=None,
                   require_authentication=False,
-                  encryption_policy=ENCRYPTION_OPTIONAL):
+                  encryption_policy=ENCRYPTION_OPTIONAL, jwt=None):
         cdef:
             string c_addr
             vector[string] c_addrs
@@ -341,6 +341,9 @@
         if require_authentication:
             builder.require_authentication(require_authentication)
 
+        if jwt is not None:
+            builder.jwt(tobytes(jwt))
+
         builder.encryption_policy(encryption_policy)
 
 
diff --git a/python/kudu/libkudu_client.pxd b/python/kudu/libkudu_client.pxd
index 671aca0..e9522b3 100644
--- a/python/kudu/libkudu_client.pxd
+++ b/python/kudu/libkudu_client.pxd
@@ -598,6 +598,8 @@
 
         KuduClientBuilder& encryption_policy(EncryptionPolicy encryption_policy)
 
+        KuduClientBuilder& jwt(const string& jwt)
+
         Status Build(shared_ptr[KuduClient]* client)
 
     cdef cppclass KuduTabletServer:
diff --git a/python/kudu/tests/common.py b/python/kudu/tests/common.py
index 4087832..a59922d 100644
--- a/python/kudu/tests/common.py
+++ b/python/kudu/tests/common.py
@@ -40,6 +40,9 @@
     NUM_MASTER_SERVERS = 3
     NUM_TABLET_SERVERS = 3
 
+    valid_account_id = "valid_account_id"
+    invalid_account_id = "invalid_account_id"
+
     @classmethod
     def send_and_receive(cls, proc, request):
         binary_req = (json.dumps(request) + "\n").encode("utf-8")
@@ -88,7 +91,15 @@
                        "--default_num_replicas=1",
                        "--ipki_ca_key_size=2048",
                        "--ipki_server_key_size=2048" ],
-                   "extraTserverFlags" : [ "--ipki_server_key_size=2048" ]}})
+                   "extraTserverFlags" : [ "--ipki_server_key_size=2048" ],
+                   "mini_oidc_options" :
+                   { "expiration_time" : "300000",
+                     "jwks_options" :
+                     [{ "account_id" : cls.valid_account_id,
+                       "is_valid_key" : "true" },
+                       { "account_id" : cls.invalid_account_id,
+                       "is_valid_key" : "false" },
+                       ]}}})
         cls.send_and_receive(p, { "start_cluster" : {}})
 
         # Get information about the cluster's masters.
@@ -137,3 +148,17 @@
     @classmethod
     def example_partitioning(cls):
         return Partitioning().set_range_partition_columns(['key'])
+
+    @classmethod
+    def get_jwt(cls, valid=True):
+        account_id = cls.valid_account_id
+        is_valid_key = valid
+        if not valid:
+            account_id = cls.invalid_account_id
+
+        resp = cls.send_and_receive(
+            cls.cluster_proc, { "create_jwt" : {
+                "account_id" : account_id,
+                "subject" : "test",
+                "is_valid_key" : is_valid_key}})
+        return resp['createJwt']['jwt']
diff --git a/python/kudu/tests/test_client.py b/python/kudu/tests/test_client.py
index 44d8e4a..02da29d 100755
--- a/python/kudu/tests/test_client.py
+++ b/python/kudu/tests/test_client.py
@@ -881,6 +881,7 @@
             except:
                 pass
 
+class TestAuthAndEncription(KuduTestBase, CompatUnitTest):
     def test_require_encryption(self):
         client = kudu.connect(self.master_hosts, self.master_ports,
                               encryption_policy=ENCRYPTION_REQUIRED)
@@ -893,6 +894,19 @@
             client = kudu.connect(self.master_hosts, self.master_ports,
                      require_authentication=True)
 
+class TestJwt(KuduTestBase, CompatUnitTest):
+    def test_jwt(self):
+        jwt = self.get_jwt(valid=True)
+        client = kudu.connect(self.master_hosts, self.master_ports,
+                              require_authentication=True, jwt=jwt)
+
+        jwt = self.get_jwt(valid=False)
+        error_msg = ('FATAL_INVALID_JWT: Not authorized: Verification failed, error: ' +
+        'failed to verify signature: VerifyFinal failed')
+        with self.assertRaisesRegex(kudu.KuduBadStatus, error_msg):
+            client = kudu.connect(self.master_hosts, self.master_ports,
+                              require_authentication=True, jwt=jwt)
+
 class TestMonoDelta(CompatUnitTest):
 
     def test_empty_ctor(self):