Add callback for oauth (@stevenylai, confluentinc#960)

stevenylai · web-flow · commit b7f8dce998ec · 2021-03-15T14:37:31.000+01:00
diff --git a/docs/index.rst b/docs/index.rst
@@ -459,6 +459,14 @@ The Python bindings also provide some additional configuration properties:
   This callback is served upon calling ``client.poll()`` or ``producer.flush()``. See
   https://github.com/edenhill/librdkafka/wiki/Statistics" for more information.
 
+* ``oauth_cb(config_str)``: Callback for retrieving OAuth Bearer token.
+  Function argument ``config_str`` is a str from config: ``sasl.oauthbearer.config``.
+  Return value of this callback is expected to be ``(token_str, expiry_time)`` tuple
+  where ``expiry_time`` is the time in seconds since the epoch as a floating point number.
+  This callback is useful only when ``sasl.mechanisms=OAUTHBEARER`` is set and
+  is served to get the initial token before a successful broker connection can be made.
+  The callback can be triggered by calling ``client.poll()`` or ``producer.flush()``.
+
 * ``on_delivery(kafka.KafkaError, kafka.Message)`` (**Producer**): value is a Python function reference
   that is called once for each produced message to indicate the final
   delivery result (success or failure).
diff --git a/examples/README.md b/examples/README.md
@@ -14,6 +14,7 @@ The scripts in this directory provide code examples using Confluent's Python cli
 * [protobuf_consumer.py](protobuf_consumer.py): DeserializingConsumer with ProtobufDeserializer
 * [sasl_producer.py](sasl_producer.py): SerializingProducer with SASL Authentication
 * [list_offsets.py](list_offsets.py): List committed offsets and consumer lag for group and topics
+* [oauth_producer.py](oauth_producer.py): SerializingProducer with OAuth Authentication (client credentials)
 
 Additional examples for [Confluent Cloud](https://www.confluent.io/confluent-cloud/):
 
diff --git a/examples/oauth_producer.py b/examples/oauth_producer.py
@@ -0,0 +1,138 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright 2020 Confluent Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+#
+# This uses OAuth client credentials grant:
+# https://www.oauth.com/oauth2-servers/access-tokens/client-credentials/
+# where client_id and client_secret are passed as HTTP Authorization header
+#
+
+import logging
+import functools
+import argparse
+import time
+from confluent_kafka import SerializingProducer
+from confluent_kafka.serialization import StringSerializer
+import requests
+
+
+def _get_token(args, config):
+    """Note here value of config comes from sasl.oauthbearer.config below.
+    It is not used in this example but you can put arbitrary values to
+    configure how you can get the token (e.g. which token URL to use)
+    """
+    payload = {
+        'grant_type': 'client_credentials',
+        'scope': ' '.join(args.scopes)
+    }
+    resp = requests.post(args.token_url,
+                         auth=(args.client_id, args.client_secret),
+                         data=payload)
+    token = resp.json()
+    return token['access_token'], time.time() + float(token['expires_in'])
+
+
+def producer_config(args):
+    logger = logging.getLogger(__name__)
+    return {
+        'bootstrap.servers': args.bootstrap_servers,
+        'key.serializer': StringSerializer('utf_8'),
+        'value.serializer': StringSerializer('utf_8'),
+        'security.protocol': 'sasl_plaintext',
+        'sasl.mechanisms': 'OAUTHBEARER',
+        # sasl.oauthbearer.config can be used to pass argument to your oauth_cb
+        # It is not used in this example since we are passing all the arguments
+        # from command line
+        # 'sasl.oauthbearer.config': 'not-used',
+        'oauth_cb': functools.partial(_get_token, args),
+        'logger': logger,
+    }
+
+
+def delivery_report(err, msg):
+    """
+    Reports the failure or success of a message delivery.
+
+    Args:
+        err (KafkaError): The error that occurred on None on success.
+
+        msg (Message): The message that was produced or failed.
+
+    Note:
+        In the delivery report callback the Message.key() and Message.value()
+        will be the binary format as encoded by any configured Serializers and
+        not the same object that was passed to produce().
+        If you wish to pass the original object(s) for key and value to delivery
+        report callback we recommend a bound callback or lambda where you pass
+        the objects along.
+
+    """
+    if err is not None:
+        print('Delivery failed for User record {}: {}'.format(msg.key(), err))
+        return
+    print('User record {} successfully produced to {} [{}] at offset {}'.format(
+        msg.key(), msg.topic(), msg.partition(), msg.offset()))
+
+
+def main(args):
+    topic = args.topic
+    delimiter = args.delimiter
+
+    producer_conf = producer_config(args)
+
+    producer = SerializingProducer(producer_conf)
+
+    print('Producing records to topic {}. ^C to exit.'.format(topic))
+    while True:
+        # Serve on_delivery callbacks from previous calls to produce()
+        producer.poll(0.0)
+        try:
+            msg_data = input(">")
+            msg = msg_data.split(delimiter)
+            if len(msg) == 2:
+                producer.produce(topic=topic, key=msg[0], value=msg[1],
+                                 on_delivery=delivery_report)
+            else:
+                producer.produce(topic=topic, value=msg[0],
+                                 on_delivery=delivery_report)
+        except KeyboardInterrupt:
+            break
+
+    print('\nFlushing {} records...'.format(len(producer)))
+    producer.flush()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="SerializingProducer OAUTH Example"
+                                                 " with client credentials grant")
+    parser.add_argument('-b', dest="bootstrap_servers", required=True,
+                        help="Bootstrap broker(s) (host[:port])")
+    parser.add_argument('-t', dest="topic", default="example_producer_oauth",
+                        help="Topic name")
+    parser.add_argument('-d', dest="delimiter", default="|",
+                        help="Key-Value delimiter. Defaults to '|'"),
+    parser.add_argument('--client', dest="client_id", required=True,
+                        help="Client ID for client credentials flow")
+    parser.add_argument('--secret', dest="client_secret", required=True,
+                        help="Client secret for client credentials flow.")
+    parser.add_argument('--token-url', dest="token_url", required=True,
+                        help="Token URL.")
+    parser.add_argument('--scopes', dest="scopes", required=True, nargs='+',
+                        help="Scopes requested from OAuth server.")
+
+    main(parser.parse_args())
diff --git a/examples/requirements.txt b/examples/requirements.txt
@@ -7,3 +7,4 @@ pyrsistent==0.16.1;python_version<"3.0"
 pyrsistent;python_version>"3.0"
 jsonschema
 protobuf
+requests
diff --git a/src/confluent_kafka/src/confluent_kafka.c b/src/confluent_kafka/src/confluent_kafka.c
@@ -1521,6 +1521,50 @@ static void log_cb (const rd_kafka_t *rk, int level,
         CallState_resume(cs);
 }
 
+static void oauth_cb (rd_kafka_t *rk, const char *oauthbearer_config,
+                      void *opaque) {
+        Handle *h = opaque;
+        PyObject *eo, *result;
+        CallState *cs;
+        const char *token;
+        double expiry;
+        char err_msg[2048];
+        rd_kafka_resp_err_t err_code;
+
+        cs = CallState_get(h);
+
+        eo = Py_BuildValue("s", oauthbearer_config);
+        result = PyObject_CallFunctionObjArgs(h->oauth_cb, eo, NULL);
+        Py_DECREF(eo);
+
+        if (!result) {
+                goto err;
+        }
+        if (!PyArg_ParseTuple(result, "sd", &token, &expiry)) {
+                Py_DECREF(result);
+                PyErr_Format(PyExc_TypeError,
+                             "expect returned value from oauth_cb "
+                             "to be (token_str, expiry_time) tuple");
+                goto err;
+        }
+        err_code = rd_kafka_oauthbearer_set_token(h->rk, token,
+                                                  (int64_t)(expiry * 1000),
+                                                  "", NULL, 0, err_msg,
+                                                  sizeof(err_msg));
+        Py_DECREF(result);
+        if (err_code) {
+                PyErr_Format(PyExc_ValueError, "%s", err_msg);
+                goto err;
+        }
+        goto done;
+
+ err:
+        CallState_crash(cs);
+        rd_kafka_yield(h->rk);
+ done:
+        CallState_resume(cs);
+}
+
 /****************************************************************************
  *
  *
@@ -1949,6 +1993,25 @@ rd_kafka_conf_t *common_conf_setup (rd_kafka_type_t ktype,
                         Py_XDECREF(ks8);
                         Py_DECREF(ks);
                         continue;
+                } else if (!strcmp(k, "oauth_cb")) {
+                        if (!PyCallable_Check(vo)) {
+                                PyErr_SetString(PyExc_TypeError,
+                                                "expected oauth_cb property "
+                                                "as a callable function");
+                                goto inner_err;
+                        }
+                        if (h->oauth_cb) {
+                                Py_DECREF(h->oauth_cb);
+                                h->oauth_cb = NULL;
+                        }
+
+                        if (vo != Py_None) {
+                                h->oauth_cb = vo;
+                                Py_INCREF(h->oauth_cb);
+                        }
+                        Py_XDECREF(ks8);
+                        Py_DECREF(ks);
+                        continue;
                 }
 
 		/* Special handling for certain config keys. */
@@ -2019,6 +2082,9 @@ rd_kafka_conf_t *common_conf_setup (rd_kafka_type_t ktype,
                 rd_kafka_conf_set_log_cb(conf, log_cb);
         }
 
+        if (h->oauth_cb)
+                rd_kafka_conf_set_oauthbearer_token_refresh_cb(conf, oauth_cb);
+
 	rd_kafka_conf_set_opaque(conf, h);
 
 #ifdef WITH_PY_TSS
diff --git a/src/confluent_kafka/src/confluent_kafka.h b/src/confluent_kafka/src/confluent_kafka.h
@@ -236,6 +236,7 @@ typedef struct {
         rd_kafka_type_t type; /* Producer or consumer */
 
         PyObject *logger;
+        PyObject *oauth_cb;
 
 	union {
 		/**
diff --git a/tests/test_misc.py b/tests/test_misc.py
@@ -130,6 +130,35 @@ def test_throttle_event_types():
     assert str(throttle_event) == "broker/0 throttled for 10000 ms"
 
 
+# global variable for oauth_cb call back function
+seen_oauth_cb = False
+
+
+def test_oauth_cb():
+    """ Tests oauth_cb. """
+
+    def oauth_cb(oauth_config):
+        global seen_oauth_cb
+        seen_oauth_cb = True
+        assert oauth_config == 'oauth_cb'
+        return 'token', time.time() + 300.0
+
+    conf = {'group.id': 'test',
+            'security.protocol': 'sasl_plaintext',
+            'sasl.mechanisms': 'OAUTHBEARER',
+            'socket.timeout.ms': '100',
+            'session.timeout.ms': 1000,  # Avoid close() blocking too long
+            'sasl.oauthbearer.config': 'oauth_cb',
+            'oauth_cb': oauth_cb
+            }
+
+    kc = confluent_kafka.Consumer(**conf)
+
+    while not seen_oauth_cb:
+        kc.poll(timeout=1)
+    kc.close()
+
+
 def skip_interceptors():
     # Run interceptor test if monitoring-interceptor is found
     for path in ["/usr/lib", "/usr/local/lib", "staging/libs", "."]:
