OAuth token refresh (Part 1) (#547)

Author: Siting Ren

.github/workflows/ci.yaml

@@ -8,13 +8,20 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.7', '3.8', '3.9', '3.10', '3.11', '3.12', 'pypy3.9']
+        python-version: ['3.7', '3.8', '3.9', '3.10', '3.11', '3.12', 'pypy3.10']
+
+    env:
+      REALM: test
+      USER: oauth_user
+      PASSWORD: password
+      CLIENT_ID: vertica
+      CLIENT_SECRET: P9f8350QQIUhFfK1GF5sMhq4Dm3P6Sbs
 
     steps:
       - name: Check out repository
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Set up a Keycloak docker container
@@ -47,12 +54,6 @@ jobs:
           echo "Wait for keycloak ready ..."
           bash -c 'while true; do curl -s localhost:8080 &>/dev/null; ret=$?; [[ $ret -eq 0 ]] && break; echo "..."; sleep 3; done'
 
-          REALM="test"
-          USER="oauth_user"
-          PASSWORD="password"
-          CLIENT_ID="vertica"
-          CLIENT_SECRET="P9f8350QQIUhFfK1GF5sMhq4Dm3P6Sbs"
-
           docker exec -i keycloak /bin/bash <<EOF
             /opt/keycloak/bin/kcadm.sh config credentials --server http://localhost:8080 --realm master --user admin --password admin
             /opt/keycloak/bin/kcadm.sh create realms -s realm=${REALM} -s enabled=true
@@ -74,6 +75,7 @@ jobs:
               --data-urlencode "client_secret=${CLIENT_SECRET}" \
               --data-urlencode 'grant_type=password' -o oauth.json
           cat oauth.json | python3 -c 'import json,sys;obj=json.load(sys.stdin);print(obj["access_token"])' > access_token.txt
+          cat oauth.json | python3 -c 'import json,sys;obj=json.load(sys.stdin);print(obj["refresh_token"])' > refresh_token.txt
 
           docker exec -u dbadmin vertica_docker /opt/vertica/bin/vsql -c "CREATE AUTHENTICATION v_oauth METHOD 'oauth' HOST '0.0.0.0/0';"
           docker exec -u dbadmin vertica_docker /opt/vertica/bin/vsql -c "ALTER AUTHENTICATION v_oauth SET client_id = '${CLIENT_ID}';"
@@ -95,5 +97,10 @@ jobs:
         run: |
           export VP_TEST_USER=dbadmin
           export VP_TEST_OAUTH_ACCESS_TOKEN=`cat access_token.txt`
-          export VP_TEST_OAUTH_USER=oauth_user
+          export VP_TEST_OAUTH_REFRESH_TOKEN=`cat refresh_token.txt`
+          export VP_TEST_OAUTH_USER=${USER}
+          export VP_TEST_OAUTH_CLIENT_ID=${CLIENT_ID}
+          export VP_TEST_OAUTH_CLIENT_SECRET=${CLIENT_SECRET}
+          export VP_TEST_OAUTH_TOKEN_URL="http://`hostname`:8080/realms/${REALM}/protocol/openid-connect/token"
+          export VP_TEST_OAUTH_DISCOVERY_URL="http://`hostname`:8080/realms/${REALM}/.well-known/openid-configuration"
           tox -e py

README.md

@@ -90,7 +90,7 @@ with vertica_python.connect(**conn_info) as connection:
 | ------------- | ------------- |
 | host     | The server host of the connection. This can be a host name or an IP address. <br>**_Default_**: "localhost" |
 | port     | The port of the connection. <br>**_Default_**: 5433 |
-| user     | The database user name to use to connect to the database. <br>**_Default_**: OS login user name |
+| user     | The database user name to use to connect to the database. <br>**_Default_**:<br>&nbsp;&nbsp;&nbsp;&nbsp;(for non-OAuth connections) OS login user name  <br>&nbsp;&nbsp;&nbsp;&nbsp;(for OAuth connections) "" |
 | password | The password to use to log into the database. <br>**_Default_**: "" |
 | database | The database name. <br>**_Default_**: "" |
 | autocommit | See [Autocommit](#autocommit). <br>**_Default_**: False |
@@ -103,7 +103,9 @@ with vertica_python.connect(**conn_info) as connection:
 | kerberos_service_name | See [Kerberos Authentication](#kerberos-authentication). <br>**_Default_**: "vertica" |
 | log_level | See [Logging](#logging). |
 | log_path | See [Logging](#logging). |
-| oauth_access_token | To authenticate via OAuth, provide an OAuth Access Token that authorizes a user to the database. <br>**_Default_**: "" |
+| oauth_access_token | See [OAuth Authentication](#oauth-authentication). <br>**_Default_**: "" |
+| oauth_refresh_token | See [OAuth Authentication](#oauth-authentication). <br>**_Default_**: "" |
+| oauth_config | See [OAuth Authentication](#oauth-authentication). <br>**_Default_**: {} |
 | request_complex_types | See [SQL Data conversion to Python objects](#sql-data-conversion-to-python-objects). <br>**_Default_**: True |
 | session_label | Sets a label for the connection on the server. This value appears in the client_label column of the _v_monitor.sessions_ system table. <br>**_Default_**: an auto-generated label with format of `vertica-python-{version}-{random_uuid}` |
 | ssl | See [TLS/SSL](#tlsssl). <br>**_Default_**: False (disabled) |
@@ -141,7 +143,7 @@ with vertica_python.connect(dsn=connection_str, **additional_info) as conn:
 ```
 
 #### TLS/SSL
-You can pass `True` to `ssl` to enable TLS/SSL connection (Internally [ssl.wrap_socket(sock)](https://docs.python.org/3/library/ssl.html#ssl.wrap_socket) is called).
+You can pass `True` to `ssl` to enable TLS/SSL connection (equivalent to TLSMode=require).
 
 ```python
 import vertica_python
@@ -258,6 +260,50 @@ with vertica_python.connect(**conn_info) as conn:
     # do things
 ```
 
+#### OAuth Authentication
+To authenticate via OAuth, one way is to provide an `oauth_access_token` that authorizes a user to the database. 
+```python
+import vertica_python
+
+conn_info = {'host': '127.0.0.1',
+             'port': 5433,
+             'database': 'a_database',
+             # valid OAuth access token
+             'oauth_access_token': 'xxxxxx'}
+
+with vertica_python.connect(**conn_info) as conn:
+    # do things
+```
+In cases where `oauth_access_token` is not set or introspection fails (e.g. when the access token expires), the client can do a token refresh when both `oauth_refresh_token` and `oauth_config` are set. The client will retrieve a new access token from the identity provider and use it to connect with the database.
+```python
+import vertica_python
+
+conn_info = {'host': '127.0.0.1',
+             'port': 5433,
+             'database': 'a_database',
+             # OAuth refresh token and configurations
+             'oauth_refresh_token': 'xxxxxx',
+             'oauth_config': {
+                 'client_secret': 'wK3SqFbExDS',
+                 'client_id': 'vertica',
+                 'token_url': 'https://example.com:8443/realms/master/protocol/openid-connect/token',
+             }
+}
+
+with vertica_python.connect(**conn_info) as conn:
+    # do things
+```
+The following table lists the `oauth_config` parameters used to configure OAuth token refresh:
+
+| Parameter     |  Description |
+| ------------- | ------------- |
+| client_id     | The client ID of the client application registered in the identity provider. |
+| client_secret | The client secret of the client application registered in the identity provider.|
+| token_url     | The endpoint to which token refresh requests are sent. The format for this depends on your provider. For examples, see the [Keycloak](https://www.keycloak.org/docs/latest/securing_apps/#token-endpoint) and [Okta](https://developer.okta.com/docs/reference/api/oidc/#token) documentation.|
+| discovery_url | Also known as the [OpenID Provider Configuration Document](https://openid.net/specs/openid-connect-discovery-1_0.html#ProviderConfigurationRequest), this endpoint contains a list of all other endpoints supported by the identity provider. If set, *token_url* do not need to be specified.<br>If you set both *discovery_url* and *token_url*, then *token_url* takes precedence.|
+| scope         | The requested OAuth scopes, delimited with spaces. These scopes define the extent of access to the resource server (in this case, Vertica) granted to the client by the access token. For details, see the [OAuth documentation](https://www.oauth.com/oauth2-servers/scope/defining-scopes/). |
+
+
 #### Logging
 Logging is disabled by default if neither ```log_level``` or ```log_path``` are set. Passing value to at least one of those options to enable logging.
 

requirements-dev.txt

@@ -2,5 +2,6 @@ pytest
 pytest-timeout
 python-dateutil
 six
+requests
 tox
 #kerberos

setup.py

@@ -59,7 +59,8 @@
     python_requires=">=3.7",
     install_requires=[
         'python-dateutil>=1.5',
-        'six>=1.10.0'
+        'six>=1.10.0',
+        'requests>=2.26.0'
     ],
     classifiers=[
         "Development Status :: 5 - Production/Stable",

vertica_python/__init__.py

@@ -59,8 +59,8 @@
 version_info = (1, 3, 8)
 __version__ = '.'.join(map(str, version_info))
 
-# The protocol version (3.15) implemented in this library.
-PROTOCOL_VERSION = 3 << 16 | 15
+# The protocol version (3.16) implemented in this library.
+PROTOCOL_VERSION = 3 << 16 | 16
 
 apilevel = 2.0
 threadsafety = 1  # Threads may share the module, but not connections!

vertica_python/errors.py

@@ -103,6 +103,14 @@ class KerberosError(ConnectionError):
 class SSLNotSupported(ConnectionError):
     pass
 
+class OAuthConfigurationError(ConnectionError):
+    pass
+
+class OAuthEndpointDiscoveryError(ConnectionError):
+    pass
+
+class OAuthTokenRefreshError(ConnectionError):
+    pass
 
 class MessageError(InternalError):
     pass

vertica_python/tests/common/base.py

@@ -56,6 +56,11 @@
     'password': '',
     'database': '',
     'oauth_access_token': '',
+    'oauth_refresh_token': '',
+    'oauth_client_id': '',
+    'oauth_client_secret': '',
+    'oauth_token_url': '',
+    'oauth_discovery_url': '',
     'oauth_user': '',
 }
 

vertica_python/tests/common/vp_test.conf.example

@@ -18,4 +18,9 @@ VP_TEST_LOG_DIR=mylog/vp_tox_tests_log
 # OAuth authentication information
 #VP_TEST_OAUTH_USER=<can be ignored if VP_TEST_DATABASE is set>
 #VP_TEST_OAUTH_ACCESS_TOKEN=******
+#VP_TEST_OAUTH_REFRESH_TOKEN=******
+#VP_TEST_OAUTH_CLIENT_ID=vertica
+#VP_TEST_OAUTH_CLIENT_SECRET=******
+#VP_TEST_OAUTH_TOKEN_URL=http://hostname:8080/realms/test/protocol/openid-connect/token
+#VP_TEST_OAUTH_DISCOVERY_URL=http://hostname:8080/realms/test/.well-known/openid-configuration
 

vertica_python/tests/integration_tests/base.py

@@ -55,7 +55,10 @@ class VerticaPythonIntegrationTestCase(VerticaPythonTestCase):
     def setUpClass(cls):
         config_list = ['log_dir', 'log_level', 'host', 'port',
                        'user', 'password', 'database',
-                       'oauth_access_token', 'oauth_user',]
+                       'oauth_access_token', 'oauth_refresh_token',
+                       'oauth_client_secret', 'oauth_client_id',
+                       'oauth_token_url', 'oauth_discovery_url',
+                       'oauth_user',]
         cls.test_config = cls._load_test_config(config_list)
 
         # Test logger
@@ -76,6 +79,11 @@ def setUpClass(cls):
         }
         cls._oauth_info = {
             'access_token': cls.test_config['oauth_access_token'],
+            'refresh_token': cls.test_config['oauth_refresh_token'],
+            'client_secret': cls.test_config['oauth_client_secret'],
+            'client_id': cls.test_config['oauth_client_id'],
+            'token_url': cls.test_config['oauth_token_url'],
+            'discovery_url': cls.test_config['oauth_discovery_url'],
             'user': cls.test_config['oauth_user'],
         }
         cls.db_node_num = cls._get_node_num()

vertica_python/tests/integration_tests/test_authentication.py

@@ -15,6 +15,7 @@
 from __future__ import print_function, division, absolute_import, annotations
 
 from .base import VerticaPythonIntegrationTestCase
+from ...errors import OAuthTokenRefreshError
 
 
 class AuthenticationTestCase(VerticaPythonIntegrationTestCase):
@@ -28,6 +29,10 @@ def tearDown(self):
         self._conn_info['password'] = self._password
         if 'oauth_access_token' in self._conn_info:
             del self._conn_info['oauth_access_token']
+        if 'oauth_refresh_token' in self._conn_info:
+            del self._conn_info['oauth_refresh_token']
+        if 'oauth_config' in self._conn_info:
+            del self._conn_info['oauth_config']
         super(AuthenticationTestCase, self).tearDown()
 
     def test_SHA512(self):
@@ -109,10 +114,12 @@ def test_password_expire(self):
                 cur.execute("DROP AUTHENTICATION IF EXISTS testIPv6hostHash CASCADE")
                 cur.execute("DROP AUTHENTICATION IF EXISTS testlocalHash CASCADE")
 
-    def test_oauth(self):
+    def test_oauth_access_token(self):
         self.require_protocol_at_least(3 << 16 | 11)
         if not self._oauth_info['access_token']:
-            self.skipTest('OAuth not set')
+            self.skipTest('OAuth access token not set')
+        if not self._oauth_info['user'] and not self._conn_info['database']:
+            self.skipTest('Both database and oauth_user are not set')
 
         self._conn_info['user'] = self._oauth_info['user']
         self._conn_info['oauth_access_token'] = self._oauth_info['access_token']
@@ -122,5 +129,67 @@ def test_oauth(self):
             res = cur.fetchone()
             self.assertEqual(res[0], 'OAuth')
 
+    def _test_oauth_refresh(self, access_token):
+        self.require_protocol_at_least(3 << 16 | 11)
+        if not self._oauth_info['refresh_token']:
+            self.skipTest('OAuth refresh token not set')
+        if not (self._oauth_info['client_secret'] and self._oauth_info['client_id'] and self._oauth_info['token_url']):
+            self.skipTest('One or more OAuth config (client_id, client_secret, token_url) not set')
+        if not self._oauth_info['user'] and not self._conn_info['database']:
+            self.skipTest('Both database and oauth_user are not set')
+
+        if access_token is not None:
+            self._conn_info['oauth_access_token'] = access_token
+        self._conn_info['user'] = self._oauth_info['user']
+        self._conn_info['oauth_refresh_token'] = self._oauth_info['refresh_token']
+        self._conn_info['oauth_config'] = {
+                'client_secret': self._oauth_info['client_secret'],
+                'client_id': self._oauth_info['client_id'],
+                'token_url': self._oauth_info['token_url'],
+        }
+        with self._connect() as conn:
+            cur = conn.cursor()
+            cur.execute("SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
+            res = cur.fetchone()
+            self.assertEqual(res[0], 'OAuth')
+
+    def test_oauth_token_refresh_with_access_token_not_set(self):
+        self._test_oauth_refresh(access_token=None)
+
+    def test_oauth_token_refresh_with_invalid_access_token(self):
+        self._test_oauth_refresh(access_token='invalid_value')
+
+    def test_oauth_token_refresh_with_empty_access_token(self):
+        self._test_oauth_refresh(access_token='')
+
+    def test_oauth_token_refresh_with_discovery_url(self):
+        self.require_protocol_at_least(3 << 16 | 11)
+        if not self._oauth_info['refresh_token']:
+            self.skipTest('OAuth refresh token not set')
+        if not (self._oauth_info['client_secret'] and self._oauth_info['client_id'] and self._oauth_info['discovery_url']):
+            self.skipTest('One or more OAuth config (client_id, client_secret, discovery_url) not set')
+        if not self._oauth_info['user'] and not self._conn_info['database']:
+            self.skipTest('Both database and oauth_user are not set')
+
+        self._conn_info['user'] = self._oauth_info['user']
+        self._conn_info['oauth_refresh_token'] = self._oauth_info['refresh_token']
+        msg = 'Token URL or Discovery URL must be set.'
+        self.assertConnectionFail(err_type=OAuthTokenRefreshError, err_msg=msg)
+
+        self._conn_info['oauth_config'] = {
+                'client_secret': self._oauth_info['client_secret'],
+                'client_id': self._oauth_info['client_id'],
+                'discovery_url': self._oauth_info['discovery_url'],
+        }
+        with self._connect() as conn:
+            cur = conn.cursor()
+            cur.execute("SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
+            res = cur.fetchone()
+            self.assertEqual(res[0], 'OAuth')
+
+        # Token URL takes precedence over Discovery URL
+        self._conn_info['oauth_config']['token_url'] = 'invalid_value'
+        self.assertConnectionFail(err_type=OAuthTokenRefreshError, err_msg='Failed getting OAuth access token from a refresh token.')
+
 
 exec(AuthenticationTestCase.createPrepStmtClass())