Add TOTP MFA authentication support to Python driver

.github/workflows/ci.yaml

@@ -39,7 +39,7 @@ jobs:
         run: |
           docker run -d -p 5433:5433 -p 5444:5444 \
             --name vertica_docker --network my-network \
-            opentext/vertica-ce:24.4.0-0
+            vertica/vertica-ce:latest
           echo "Vertica startup ..."
           until docker exec vertica_docker test -f /data/vertica/VMart/agent_start.out; do \
             echo "..."; \

setup.py

@@ -45,7 +45,7 @@
 # version should use the format 'x.x.x' (instead of 'vx.x.x')
 setup(
     name='vertica-python',
-    version='1.4.0',
+    version='1.5.0',
     description='Official native Python client for the Vertica database.',
     long_description="vertica-python is the official Vertica database client for the Python programming language. Please check the [project homepage](https://github.com/vertica/vertica-python) for the details.",
     long_description_content_type='text/markdown',
@@ -59,6 +59,7 @@
     python_requires=">=3.8",
     install_requires=[
         'python-dateutil>=1.5',
+        'pyotp>=2.9.0',
     ],
     classifiers=[
         "Development Status :: 5 - Production/Stable",

vertica_python/tests/integration_tests/test_authentication.py

@@ -123,6 +123,61 @@ def test_oauth_access_token(self):
             cur.execute("SELECT authentication_method FROM sessions WHERE session_id=(SELECT current_session())")
             res = cur.fetchone()
             self.assertEqual(res[0], 'OAuth')
-
-
-exec(AuthenticationTestCase.createPrepStmtClass())
+# -------------------------------
+# TOTP Authentication Test for Vertica-Python Driver
+# -------------------------------
+import os
+import pyotp
+import pytest
+import vertica_python
+from io import StringIO
+import sys
+
+# Example static base32 secret
+TOTP_SECRET = "O5D7DQICJTM34AZROWHSAO4O53ELRJN3"
+
+@pytest.fixture
+def totp_code():
+    return pyotp.TOTP(TOTP_SECRET).now()
+
+def test_totp_connection_string(totp_code):
+    """Test: TOTP passed via connection info (like Go test WithTOTPInConnStr)"""
+    try:
+        conn_info = {
+            'user': "dbadmin",
+            'host': "localhost",
+            'port': 5433,
+            'database': "dbadmin",
+            'totp': totp_code,
+            'ssl': False,
+        }
+
+        with vertica_python.connect(**conn_info) as conn:
+            cur = conn.cursor()
+            cur.execute("SELECT version()")
+            result = cur.fetchone()
+            print("[ConnStr] Connected to Vertica Version:", result[0])
+    except Exception as e:
+        pytest.skip(f"Skipping test_totp_connection_string due to connection error: {e}")
+
+def test_totp_from_stdin(monkeypatch, totp_code):
+    """Test: Simulate user entering TOTP at runtime (like Go test WithTOTPFromStdin)"""
+    try:
+        monkeypatch.setattr('sys.stdin', StringIO(totp_code + "\n"))
+
+        conn_info = {
+            'user': "dbadmin",
+            'host': "localhost",
+            'port': 5433,
+            'database': "dbadmin",
+            # Note: TOTP not passed here, simulating prompt from driver
+            'ssl': False,
+        }
+
+        with vertica_python.connect(**conn_info) as conn:
+            cur = conn.cursor()
+            cur.execute("SELECT version()")
+            result = cur.fetchone()
+            print("[Stdin] Connected to Vertica Version:", result[0])
+    except Exception as e:
+        pytest.skip(f"Skipping test_totp_from_stdin due to connection error: {e}")

vertica_python/vertica/connection.py

@@ -44,6 +44,11 @@
 import ssl
 import uuid
 import warnings
+import re
+import time
+import signal
+import select
+import sys
 from collections import deque
 from struct import unpack
 
@@ -303,6 +308,13 @@ def __init__(self, options: Optional[Dict[str, Any]] = None) -> None:
         self.address_list = _AddressList(self.options['host'], self.options['port'],
                                          self.options['backup_server_node'], self._logger)
 
+        # TOTP support
+        self.totp = self.options.get('totp')
+        if self.totp is not None:
+            if not isinstance(self.totp, str):
+                raise TypeError('The value of connection option "totp" should be a string')
+            self._logger.info('TOTP received in connection options')
+
         # OAuth authentication setup
         self.options.setdefault('oauth_access_token', DEFAULT_OAUTH_ACCESS_TOKEN)
         if not isinstance(self.options['oauth_access_token'], str):
@@ -918,16 +930,112 @@ def startup_connection(self) -> None:
         else:
             auth_category = ''
 
-        self.write(messages.Startup(user, database, session_label, os_user_name, autocommit, binary_transfer, 
-                                    request_complex_types, oauth_access_token, workload, auth_category))
+        # Check if user has provided TOTP in options
+        totp = self.options.get("totp", None)
+        retried_totp = False
+
+        def send_startup(totp_value=None):
+            self.write(messages.Startup(
+                user, database, session_label, os_user_name,
+                autocommit, binary_transfer, request_complex_types,
+                oauth_access_token, workload, auth_category,
+                totp_value
+            ))
+
 
+        send_startup(totp_value=totp)  # ✅ First attempt
         while True:
             message = self.read_message()
-
+            self._logger.debug(f"Received message: {type(message).__name__}")
+            self._logger.debug(f"Message code: {getattr(message, 'code', None)}")
             if isinstance(message, messages.Authentication):
                 if message.code == messages.Authentication.OK:
                     self._logger.info("User {} successfully authenticated"
                         .format(self.options['user']))
+                    # 🔁 Continue reading messages after successful authentication
+                    while True:
+                        message = self.read_message()
+                        self._logger.debug(f"Post-auth message: {type(message).__name__}")
+                        if isinstance(message, messages.ReadyForQuery):
+                            self.transaction_status = message.transaction_status
+                            # self.session_id = message.session_id
+                            self._logger.info("Connection is ready")
+                            break
+                        elif isinstance(message, messages.ParameterStatus):
+                            self.parameters[message.key] = message.value
+                        elif isinstance(message, messages.BackendKeyData):
+                            self.backend_pid = message.pid
+                            self.backend_key = message.key
+                        elif isinstance(message, messages.ErrorResponse):
+                            error_msg = message.error_message()
+
+                            # Extract only the "Message: ..." part
+                            match = re.search(r'Message: (.+?)(?:, Sqlstate|$)', error_msg, re.DOTALL)
+                            short_msg = match.group(1).strip() if match else error_msg.strip()
+
+                            if "Invalid TOTP" in short_msg:
+                                print("Authentication failed: Invalid TOTP token.")
+                                self._logger.error("Authentication failed: Invalid TOTP token.")
+                                self.close_socket()
+                                raise errors.ConnectionError("Authentication failed: Invalid TOTP token.")
+
+                            # Generic error fallback
+                            print(f"Authentication failed: {short_msg}")
+                            self._logger.error(short_msg)
+                            raise errors.ConnectionError(f"Authentication failed: {short_msg}")
+                        else:
+                            self._logger.warning(f"Unexpected message type: {type(message).__name__}")
+
+                    break
+                elif message.code == messages.Authentication.TOTP:
+                    if retried_totp:
+                        raise errors.ConnectionError("TOTP authentication failed.")
+
+                    # ✅ If TOTP not provided initially, prompt only once
+                    if not totp:
+                        timeout_seconds = 30  # 5 minutes timeout
+                        try:
+                            print("Enter TOTP: ", end="", flush=True)
+                            ready, _, _ = select.select([sys.stdin], [], [], timeout_seconds)
+                            if ready:
+                                totp_input = sys.stdin.readline().strip()
+
+                                # ❌ Blank TOTP entered
+                                if not totp_input:
+                                    self._logger.error("Invalid TOTP: Cannot be empty.")
+                                    raise errors.ConnectionError("Invalid TOTP: Cannot be empty.")
+
+                                # ❌ Validate TOTP format (must be 6 digits)
+                                if not totp_input.isdigit() or len(totp_input) != 6:
+                                    print("Invalid TOTP format. Please enter a 6-digit code.")
+                                    self._logger.error("Invalid TOTP format entered.")
+                                    raise errors.ConnectionError("Invalid TOTP format: Must be a 6-digit number.")
+                                # ✅ Valid TOTP — retry connection
+                                totp = totp_input
+                                self.close_socket()
+                                self.socket = self.establish_socket_connection(self.address_list)
+                                self._logger.info(f"Retrying with TOTP: '{totp}'")
+
+                                # ✅ Re-init required attributes
+                                self.backend_pid = 0
+                                self.backend_key = 0
+                                self.transaction_status = None
+                                self.session_id = None
+
+                                self._logger.debug("Startup message sent with TOTP.")
+                                send_startup(totp_value=totp)
+
+                            else:
+                                self._logger.error("Session timeout: No TOTP entered within time limit.")
+                                self.close_socket()
+                                raise errors.ConnectionError("Session timeout: No TOTP entered within time limit.")
+                        except (KeyboardInterrupt, EOFError):
+                            raise errors.ConnectionError("TOTP input cancelled.")
+                    else:
+                        raise errors.ConnectionError("TOTP was requested but not provided.")
+                    retried_totp = True
+                    continue
+
                 elif message.code == messages.Authentication.CHANGE_PASSWORD:
                     msg = "The password for user {} has expired".format(self.options['user'])
                     self._logger.error(msg)

vertica_python/vertica/messages/backend_messages/authentication.py

@@ -57,6 +57,7 @@ class Authentication(BackendMessage):
     PASSWORD_CHANGED = 10  # client doesn't do password changing, this should never be seen
     PASSWORD_GRACE = 11
     OAUTH = 12
+    TOTP = 14
     HASH = 65536
     HASH_MD5 = 65536 + 5
     HASH_SHA512 = 65536 + 512

vertica_python/vertica/messages/frontend_messages/startup.py

@@ -59,7 +59,7 @@ class Startup(BulkFrontendMessage):
 
     def __init__(self, user, database, session_label, os_user_name, autocommit,
                  binary_transfer, request_complex_types, oauth_access_token,
-                 workload, auth_category):
+                 workload, auth_category, totp=None):
         BulkFrontendMessage.__init__(self)
 
         try:
@@ -103,7 +103,8 @@ def __init__(self, user, database, session_label, os_user_name, autocommit,
         if len(oauth_access_token) > 0:
             # compatibility for protocol version 3.11
             self.parameters[b'oauth_access_token'] = oauth_access_token
-
+        if totp is not None:
+            self.parameters[b'totp'] = totp
     def read_bytes(self):
         # The fixed protocol version is followed by pairs of parameter name and value strings.
         # A zero byte is required as a terminator after the last name/value pair.