teradataml 20.0.0.4__py3-none-any.whl → 20.0.0.6__py3-none-any.whl

teradataml/scriptmgmt/lls_utils.py

@@ -10,36 +10,38 @@ teradataml load library service wrappers.
 All teradataml wrappers to provide interface to load library service stored procedures
 from Open Analytics Framework.
 """
+import base64
 import concurrent.futures
 import functools
 import json
 import operator
 import os
+import warnings
+from json.decoder import JSONDecodeError
+from time import sleep, time
+from urllib.parse import urlparse
 
 import pandas as pd
-import requests
 
-from json.decoder import JSONDecodeError
 from teradataml import configure
-from teradataml.context.context import _get_user, get_connection
-from teradataml.common.constants import HTTPRequest, AsyncStatusColumns
-from teradataml.common.deprecations import argument_deprecation
+from teradataml.clients.auth_client import _AuthWorkflow
+from teradataml.clients.keycloak_client import _KeycloakManager
+from teradataml.clients.pkce_client import _DAWorkflow
+from teradataml.common.constants import (AsyncOpStatus, AsyncStatusColumns,
+                                         AuthMechs, HTTPRequest, TDServices)
 from teradataml.common.exceptions import TeradataMlException
-from teradataml.common.messages import Messages
 from teradataml.common.messagecodes import MessageCodes
+from teradataml.common.messages import Messages
 from teradataml.common.utils import UtilFuncs
-from teradataml.clients.pkce_client import _DAWorkflow
-from teradataml.clients.auth_client import _AuthWorkflow
+from teradataml.context.context import _get_user, get_connection
+from teradataml.scriptmgmt.UserEnv import (UserEnv, _AuthToken,
+                                           _get_auth_token, _get_ccp_url,
+                                           _get_ues_url, _process_ues_response)
+from teradataml.telemetry_utils.queryband import collect_queryband
 from teradataml.utils.internal_buffer import _InternalBuffer
-from teradataml.scriptmgmt.UserEnv import UserEnv, _get_auth_token, \
-    _process_ues_response, _get_ues_url, _AuthToken
-from teradataml.utils.validators import _Validators
-from time import time, sleep
-import warnings
-import webbrowser
-from urllib.parse import parse_qs, urlparse
 from teradataml.utils.utils import _async_run_id_info
-from teradataml.telemetry_utils.queryband import collect_queryband
+from teradataml.utils.validators import _Validators
+
 
 @collect_queryband(queryband="LstBsEnv")
 def list_base_envs():
@@ -197,7 +199,7 @@ def list_user_envs(env_name=None, **kwargs):
         ...            'python_3.9',
         ...            'Sales team environment.',
         ...            conda_env=True)
-        Conda environment creation initiated
+        Conda environment creation initiated.
         User environment 'Sales_cond_env' created.
 
         # Example 1: List all available user environments.
@@ -303,7 +305,7 @@ def list_user_envs(env_name=None, **kwargs):
 
     try:
         response = UtilFuncs._http_request(_get_ues_url(), headers=_get_auth_token())
-        # Below condition is special case handeling when remove_all_envs() used by user, remove_all_envs()
+        # Below condition is special case handling when remove_all_envs() used by user, remove_all_envs()
         # removes all the envs which result in a status_code 404 and due to which warnings provided in
         # list_user_envs() not appears.
         if response.status_code == 404 and "No user environments found." in response.text:
@@ -339,7 +341,7 @@ def list_user_envs(env_name=None, **kwargs):
             # Return the DataFrame if not empty.
             if len(pandas_df) > 0:
                 return pandas_df
-        
+
         print("No user environment(s) found.")
     except (TeradataMlException, RuntimeError):
         raise
@@ -451,8 +453,8 @@ def __create_envs(template):
                     except Exception as lib_installation_failure:
                         error_code = MessageCodes.FUNC_EXECUTION_FAILED
                         error_msg = Messages.get_message(error_code,
-                                                         "'install_lib' request for enviornment: '{}'".format(env_name), 
-                                                         '\n'+str(lib_installation_failure))
+                                                         "'install_lib' request for enviornment: '{}'".format(env_name),
+                                                         '\n' + str(lib_installation_failure))
                         print(error_msg)
                         errored = errored or True
                         pass
@@ -747,9 +749,18 @@ def create_env(env_name=None, base_env=None, desc=None, template=None, conda_env
         >>> fraud_detection_env = create_env('Fraud_detection_conda',
         ...                                  'python_3.8',
         ...                                  'Fraud detection through time matching',
-                                              conda_env=True)
-            Conda environment creation initiated
-            User environment 'Fraud_detection_conda' created.
+        ...                                   conda_env=True)
+        Conda environment creation initiated.
+        User environment 'Fraud_detection_conda' created.
+
+        # Example 5: Create a Conda R 4.2 environment with given name and
+        #            description in the Vantage.
+        >>> conda_r_env = create_env('conda_r_env',
+        ...                          'r_4.2',
+        ...                          'Conda R environment',
+        ...                           conda_env=True)
+        Conda environment creation initiated.
+        User environment 'conda_r_env' created.
     """
 
     # Either env_name or template can be used.
@@ -774,7 +785,7 @@ def create_env(env_name=None, base_env=None, desc=None, template=None, conda_env
         # Or if base_env is provided and not in the list of base envs.
         # Note: By default python base env is obtained.
         if configure.ues_url is not None and \
-        get_connection() is not None:
+                get_connection() is not None:
             # Check if base_env is provided or not in the list of base envs.
 
             # Check if user requested for conda environment but do not specify the base_env.
@@ -787,11 +798,11 @@ def create_env(env_name=None, base_env=None, desc=None, template=None, conda_env
                 # Check if base_env provided or not. If provided, check if it is available in
                 # the list of base envs. If not available, set base_env to the default python base env.
                 if not base_env or \
-                    base_env.lower() not in list_base_envs()['base_name'].str.lower().to_list():
+                        base_env.lower() not in list_base_envs()['base_name'].str.lower().to_list():
                     # Print warning message if base_env provided is not available.
                     if base_env:
-                        print(f"Note: The specified base environment '{base_env}' is unavailable. "\
-                               "Using the default base environment as specified in the documentation.")
+                        print(f"Note: The specified base environment '{base_env}' is unavailable. " \
+                              "Using the default base environment as specified in the documentation.")
                     # Set base_env to the default
                     base_env = __get_default_base_env()
         if not desc:
@@ -804,21 +815,25 @@ def create_env(env_name=None, base_env=None, desc=None, template=None, conda_env
             response = UtilFuncs._http_request(
                 _get_ues_url(conda_env=conda_env), HTTPRequest.POST, headers=_get_auth_token(), json=data)
 
-            # UES reponse.
-            resp = _process_ues_response(api_name="create_env", response=response)
+            # Validate UES response.
+            _process_ues_response(api_name="create_env", response=response)
 
             msg = "User environment '{}' created."
 
             if conda_env:
-                print("Conda environment creation initiated")
+                print("Conda environment creation initiated.")
                 # Get claim_id.
                 claim_id = response.json().get("claim_id", "")
-                # Poll the claim_id status.
+
+                # Since create_env() for conda environment is internally
+                # asynchronous but exposed as synchronous API, keep polling
+                # the status of underlying asynchronous operation until
+                # it is either successful or errored.
                 __poll_claim_id_status(claim_id, "create_env")
             print(msg.format(env_name))
 
             # Return an instance of class UserEnv.
-            return UserEnv(env_name, base_env, desc)
+            return UserEnv(env_name, base_env, desc, conda_env)
 
         except (TeradataMlException, RuntimeError):
             raise
@@ -855,7 +870,7 @@ def _async_run_status_open_af(claim_id):
         __get_claim_id_status('278381bf-e3b3-47ff-9ba5-c3b5d9007363')
     """
     # Get the claim id status.
-    resp_data = __get_status(claim_id)
+    resp_data = _get_status(claim_id)
 
     desc = _async_run_id_info.get(claim_id, {}).get("description", "Unknown")
     get_details = lambda data: {AsyncStatusColumns.ADDITIONAL_DETAILS.value:
@@ -871,7 +886,7 @@ def _async_run_status_open_af(claim_id):
     return [get_details(sub_step) for sub_step in resp_data]
 
 
-def __get_status(claim_id):
+def _get_status(claim_id):
     """
     DESCRIPTION:
         Internal function to get the status of a claim_id using
@@ -891,7 +906,7 @@ def __get_status(claim_id):
         None
 
     EXAMPLES:
-        __get_status('278381bf-e3b3-47ff-9ba5-c3b5d9007363')
+        _get_status('278381bf-e3b3-47ff-9ba5-c3b5d9007363')
     """
     # Get the claim id status
     response = UtilFuncs._http_request(_get_ues_url(env_type="fm",
@@ -1030,6 +1045,9 @@ def __manage_envs(env_name=None, api_name="remove_env", **kwargs):
 
     try:
         # Get the ues url for corresponding API.
+        # While deleting environment, endpoint UES URL for deleting
+        # normal and conda environment is same, unlike creating
+        # normal and conda environment.
         ues_url = _get_ues_url(env_name=env_name, api_name=api_name) if api_name == "remove_env" \
             else _get_ues_url(remove_all_envs=True, api_name=api_name)
 
@@ -1048,7 +1066,7 @@ def __manage_envs(env_name=None, api_name="remove_env", **kwargs):
                 if api_name == "remove_env":
                     msg = "{2}list_user_envs(). If environment is not removed, " \
                           "check the status of asynchronous call using" \
-                          " async_run_status('{1}') or get_env('{0}').status('{1}')".\
+                          " async_run_status('{1}') or get_env('{0}').status('{1}')". \
                         format(env_name, claim_id, msg)
                 else:
                     msg = "{0}async_run_status('{1}')".format(msg, claim_id)
@@ -1093,30 +1111,37 @@ def __poll_claim_id_status(claim_id, api_name="remove_env"):
             Default Value: remove_env
             Types: str
 
-
-
     RETURNS:
         None.
 
     RAISES:
-        None.
+        TeradataMlException
 
     EXAMPLES:
         __poll_claim_id_status('cf7245f0-e962-4451-addf-efa7e123998d')
     """
+    err_details = None
     while True:
         sleep(2)
 
         # Poll the claim id to get the status.
-        resp_data = __get_status(claim_id)
+        resp_data = _get_status(claim_id)
 
         # Breaking condition -
         # For create_env and remove_env: Check for the 'Finished' stage in the list of resp.
         # For remove_all_envs: above cond. and No user envs condition should break it .
         for data in resp_data:
-            if ("Finished" in data["stage"]) or \
-                    (api_name in ["create_env", "remove_all_envs"] and "Errored" in data["stage"]):
+            if AsyncOpStatus.FINISHED.value in data["stage"]:
                 return
+            elif AsyncOpStatus.ERRED.value in data["stage"]:
+                err_details = data["details"]
+                break
+        if err_details:
+            break
+
+    raise TeradataMlException(Messages.get_message(MessageCodes.FUNC_EXECUTION_FAILED,
+                                                   api_name, err_details),
+                              MessageCodes.FUNC_EXECUTION_FAILED)
 
 
 @collect_queryband(queryband="GtEnv")
@@ -1191,7 +1216,9 @@ def get_env(env_name):
         # Return an instance of class UserEnv.
         return UserEnv(userenv_row.env_name.values[0],
                        userenv_row.base_env_name.values[0],
-                       userenv_row.env_description.values[0])
+                       userenv_row.env_description.values[0],
+                       userenv_row.conda.values[0]
+                       )
     except (TeradataMlException, RuntimeError) as tdemsg:
         # TeradataMlException and RuntimeError are raised by list_user_envs.
         # list_user_envs should be replaced with get_env in the error 
@@ -1446,7 +1473,7 @@ def _remove_all_envs(env_type, **kwargs):
     if env_type.capitalize() == "Py":
         env_type = ["Python", "python"]
     else:
-        env_type = ["R"]
+        env_type = ["R", "r"]
         env_type_message = "R"
     asynchronous = kwargs.get("asynchronous", False)
 
@@ -1485,7 +1512,7 @@ def _remove_all_envs(env_type, **kwargs):
             if len(failed_envs) > 0:
                 emsg = ""
                 for env, tdemsg in failed_envs.items():
-                    emsg += "\nUser environment '{0}' failed to remove. Reason: {1}"\
+                    emsg += "\nUser environment '{0}' failed to remove. Reason: {1}" \
                         .format(env, tdemsg.args[0])
                 msg_code = MessageCodes.FUNC_EXECUTION_FAILED
                 error_msg = Messages.get_message(msg_code, "remove_all_envs()", emsg)
@@ -1598,23 +1625,94 @@ def get_user_env():
     return configure._default_user_env
 
 
+def _validate_jwt_token(base_url, token_data):
+    """
+    DESCRIPTION:
+        Function to validate the authentication token generated using PAT and PEM file.
+
+    PARAMETERS:
+        base_url:
+            Required Argument.
+            Specifies the endpoint URL for a given environment on VantageCloud Lake.
+            Types: str
+
+        token_data:
+            Required Argument.
+            Specifies the JWT token to be authenticated.
+
+    RETURNS:
+        Boolan flag representing validation status.
+            * True: Indicates that token is valid.
+            * None: Indicates that token is not validated.
+
+    RAISES:
+        TeradataMlException
+
+    EXAMPLES:
+        Example 1: Validate JWT token.
+        >>> _validate_jwt_token(base_url, token_data)
+
+    """
+    # Extract environment id from base_url.
+    try:
+        url_parser = urlparse(base_url)
+        env_id = url_parser.path.split("accounts/")[1].split("/")[0]
+        if not env_id:
+            raise
+    except Exception:
+        raise TeradataMlException(Messages.get_message(MessageCodes.FUNC_EXECUTION_FAILED,
+                                                       "set_auth_token",
+                                                       "Use valid value for 'base_url'"),
+                                  MessageCodes.FUNC_EXECUTION_FAILED)
+
+    valid_token = None
+    try:
+        response = UtilFuncs._http_request(url="{}/{}/{}/{}".format(_get_ccp_url(base_url),
+                                                                    "api", "accounts", env_id),
+                                           method_type=HTTPRequest.GET,
+                                           headers={"Authorization": "Bearer {}".format(token_data)})
+        if 200 <= response.status_code < 300:  # Authorized access.
+            valid_token = True
+        elif 400 <= response.status_code < 500:  # Unauthorized access.
+            valid_token = False
+    except:
+        pass
+
+    if valid_token is False:
+        raise TeradataMlException(Messages.get_message(MessageCodes.FUNC_EXECUTION_FAILED,
+                                                       "set_auth_token",
+                                                       "Use valid values for input arguments ['base_url',"
+                                                       " 'pat_token', 'pem_file']."),
+                                  MessageCodes.FUNC_EXECUTION_FAILED)
+    return valid_token
+
+
 @collect_queryband(queryband="StAthTkn")
 def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None, **kwargs):
     """
     DESCRIPTION:
-        Function to set the Authentication token to connect to User Environment Service
-        in VantageCloud Lake.
-        Note:
-            User must have a privilege to login with a NULL password to use set_auth_token().
-            Please refer to GRANT LOGON section in Teradata Documentation for more details.
-            If base_url and client_id are specified then authentication is through OAuth.
-            If base_url, pat_token, pem_file are specified then authentication is through PAT.
-            Refresh token still works but only for OAuth authentication.
+        Function to set the authentication token required to access services running on
+        Teradata Vantage.
+        Notes:
+            * User must have a privilege to login with a NULL password to use set_auth_token().
+              Refer to GRANT LOGON section in Teradata Documentation for more details.
+            * When "auth_mech" is not specified, arguments are used in the following combination
+              to derive authentication mechanism.
+                * If "base_url" and "client_id" are specified then token generation is done through OAuth.
+                * If "base_url", "pat_token", "pem_file" are specified then token generation is done using PAT.
+                * If "base_url", "username" and "password" are specified then authentication is done via
+                  Basic authentication mechanism using user credentials.
+                * If "base_url" and "auth_token" are specified then readily available token is used.
+                * If only "base_url" is specified then token generation is done through OAuth.
+            * Refresh token works only for OAuth authentication.
+            * Use the argument "kid" only when key used during the pem file generation is different
+              from pem file name. For example, if you use the key as 'key1' while generating pem file
+              and the name of the pem file is `key1(1).pem`, then pass value 'key1' to the argument "kid".
 
     PARAMETERS:
         base_url:
             Required Argument.
-            Specifies the CCP endpoint URL.
+            Specifies the endpoint URL for a given environment on Teradata Vantage system.
             Types: str
 
         client_id:
@@ -1631,35 +1729,99 @@ def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None,
         pem_file:
             Required, if PAT authentication is to be used, optional otherwise.
             Specifies the path to private key file which is generated from VantageCloud Lake Console.
+            Note:
+                Teradata recommends not to change the name of the file generated from VantageCloud Lake
+                Console. If the name of the file is changed, then authentication token generated from
+                this function will not work.
             Types: str
 
         **kwargs:
             username:
+                Optional Argument.
                 Specifies the user for which authentication is to be requested.
                 If not specified, then user associated with current connection is used.
-                Note:
-                    1. Use this option only if name of the database username has lower case letters.
-                    2. This option is used only for PAT and not for OAuth.
+                Notes:
+                    * Use this option only if name of the database username has lowercase letters.
+                    * This option is used only for PAT and not for OAuth.
                 Types: str
 
             expiration_time:
-                Specifies the expiration time of the token in seconds. After expiry time JWT token expires and
-                UserEnv methods does not work, user should regenerate the token.
+                Optional Argument.
+                Specifies the expiration time of the token in seconds. After expiry time, JWT
+                token expires and UserEnv methods does not work, user should regenerate the token.
                 Note:
-                    This option is used only for PAT and not for OAuth.
+                    * This option is used only for PAT and not for OAuth.
                 Default Value: 31536000
                 Types: int
 
             auth_token:
-                Optional Parameter.
-                Specifies the authentication token to connect to VantageCloud Lake.
+                Optional Argument.
+                Specifies the authentication token required to access services running
+                on Teradata Vantage.
                 Notes:
-                     * if "auth_token" is set through this function, then this function
-                       should always be used only after create_context.
-                     * use this option only if user has got JWT token and wants to set the same
-                       instead of generating it again from this function.
+                    * If "auth_token" is set through this function, then this function
+                      should always be used only after create_context().
+                    * Use this option only if user has got JWT token and wants to set
+                      the same instead of generating it again from this function.
+
+                Types: str
+
+            kid:
+                Optional Argument.
+                Specifies the name of the key which is used while generating 'pem_file'.
+                Types: str
+
+            password:
+                Optional Argument.
+                Specifies the password for database user to be used for Basic authentication.
+                Types: str
+
+            auth_url:
+                Optional Argument.
+                Specifies the endpoint URL for a keycloak server.
                 Types: str
 
+            rest_client:
+                Optional Argument.
+                Specifies the service for which keycloak token is to be generated.
+                Permitted values: "VECTORSTORE"
+                Default value: "VECTORSTORE"
+                Types: str
+
+            auth_mech:
+                Optional Argument.
+                Specifies the mechanism to be used for generating authentication token.
+                Note:
+                    * When "auth_mech" is provided, other arguments are used in the following
+                      combination as per value of "auth_mech":
+                        * OAuth: Token generation is done through OAuth by using client id
+                                 which can be sepcified by user in "client_id" argument or
+                                 can be derived internally from "base_url".
+                        * PAT: Token generation is done using "pat_token" and "pem_file".
+                        * BASIC: Authentication is done via Basic authentication mechanism
+                                 using user credentials passed in "username" and "password"
+                                 arguments.
+                        * JWT: Readily available token in "auth_token" argument is used.
+                        * KEYCLOAK: Token generation is done using keycloak.
+                Permitted Values: "OAuth", "PAT", "BASIC", "JWT", "KEYCLOAK".
+                Types: str
+
+            validate_jwt:
+                Optional Argument.
+                Specifies whether to validate generated JWT token or not.
+                Note:
+                    * Applicable only when "auth_mech" is "PAT".
+                Default value: True
+                Types: boolean
+
+            valid_from:
+                Optional Argument.
+                Specifies epoch seconds representing time from which JWT token will be valid.
+                Note:
+                    * Applicable only when "auth_mech" is "PAT".
+                Default value: None
+                Types: int
+
     RETURNS:
         True, if the operation is successful.
 
@@ -1671,10 +1833,14 @@ def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None,
         # Example 1: Set the Authentication token using default client_id.
         >>> import getpass
         >>> set_auth_token(base_url=getpass.getpass("ues_url : "))
+        Authentication token is generated and set for the session.
+        True
 
         # Example 2: Set the Authentication token by specifying the client_id.
         >>> set_auth_token(base_url=getpass.getpass("base_url : "),
         ...                client_id=getpass.getpass("client_id : "))
+        Authentication token is generated and set for the session.
+        True
 
         # Example 3: Set the Authentication token by specifying the "pem_file" and "pat_token"
         #            without specifying "username".
@@ -1682,6 +1848,7 @@ def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None,
         >>> set_auth_token(base_url=getpass.getpass("base_url : "),
         ...                pat_token=getpass.getpass("pat_token : "),
         ...                pem_file=getpass.getpass("pem_file : "))
+        Authentication token is generated, authenticated and set for the session.
         True
 
         # Example 4: Set the Authentication token by specifying the "pem_file" and "pat_token"
@@ -1689,85 +1856,201 @@ def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None,
         >>> import getpass
         >>> set_auth_token(base_url=getpass.getpass("base_url : "),
         ...                pat_token=getpass.getpass("pat_token : "),
-        ...                pem_file=getpass.getpass("pem_file : "))
-        ...                username = "alice")
+        ...                pem_file=getpass.getpass("pem_file : "),
+        ...                username=getpass.getpass("username : "))
+        Authentication token is generated, authenticated and set for the session.
+        True
+
+        # Example 5: Set the Authentication token by specifying the "pem_file" and "pat_token"
+        #            and "kid".
+        >>> import getpass
+        >>> set_auth_token(base_url=getpass.getpass("base_url : "),
+        ...                pat_token=getpass.getpass("pat_token : "),
+        ...                pem_file=getpass.getpass("pem_file : ")
+        ...                kid="key1")
+        Authentication token is generated, authenticated and set for the session.
+        True
+
+        # Example 6: Set the authentication token via Basic Authentication mechanism by
+        #            specifying the "base_url", "username" and "password".
+        >>> import getpass
+        >>> set_auth_token(base_url=getpass.getpass("base_url : "),
+        ...                username=getpass.getpass("username : "),
+        ...                password=getpass.getpass("password : "))
+        Authentication token is generated and set for the session.
+        True
+
+        # Example 7: Set the authentication token for by specifying "base_url" and
+        #            "auth_mech" as "OAuth".
+        >>> import getpass
+        >>> set_auth_token(base_url=getpass.getpass("base_url : "),
+        ...                auth_mech="OAuth")
+        Authentication token is generated and set for the session.
+        True
+
+        # Example 8: Set the authentication token for by specifying "base_url", "auth_url"
+        #            "password" and "rest_client" and generating keycloak token internally.
+        >>> import getpass
+        >>> set_auth_token(base_url=getpass.getpass("base_url : "),
+        ...                auth_url=getpass.getpass("auth_url : "),
+        ...                password=getpass.getpass("password : "),
+        ...                rest_client=getpass.getpass("rest_client : "))
+        Authentication token is generated and set for the session.
         True
+
     """
+
     # Deriving global connection using get_connection().
-    con = get_connection()
-    if con is None:
-        raise TeradataMlException(Messages.get_message(MessageCodes.INVALID_CONTEXT_CONNECTION), 
+    if get_connection() is None:
+        raise TeradataMlException(Messages.get_message(MessageCodes.INVALID_CONTEXT_CONNECTION),
                                   MessageCodes.INVALID_CONTEXT_CONNECTION)
 
-    # Getting the ues_url.
-    ues_url = kwargs.get("ues_url", None)
+    # Remove keys from _InternalBuffer which are interrelated to base_url and authentication token.
+    _InternalBuffer.remove_keys(['list_base_envs', 'default_base_env',
+                                 'vs_session_id', 'vs_header'])
 
+    # ---------------------------------ARGUMENT VALIDATION------------------------------------------------------
+    # STEP 1: Validate arguments for allowed types.
+    # ----------------------------------------------------------------------------------------------------------
     __arg_info_matrix = []
     __arg_info_matrix.append(["base_url", base_url, True, (str), True])
-    __arg_info_matrix.append(["ues_url", ues_url, True, (str), True])
     __arg_info_matrix.append(["client_id", client_id, True, (str), True])
     __arg_info_matrix.append(["pat_token", pat_token, True, (str), True])
     __arg_info_matrix.append(["pem_file", pem_file, True, (str), True])
 
-    username = kwargs.get("username", None)
-    __arg_info_matrix.append((["username", username, True, (str), True]))
+    # Get keyword arguments.
+    ues_url = kwargs.get("ues_url", None)
+    __arg_info_matrix.append(["ues_url", ues_url, True, (str), True])
+
+    username = kwargs.get("username", _get_user())
+    __arg_info_matrix.append(["username", username, True, (str), True])
+
+    password = kwargs.get("password", None)
+    __arg_info_matrix.append(["password", password, True, (str), True])
 
     auth_token = kwargs.get("auth_token")
-    __arg_info_matrix.append((["auth_token", auth_token, True, (str), True]))
+    __arg_info_matrix.append(["auth_token", auth_token, True, (str), True])
 
-    expiration_time = kwargs.get("expiration_time", 31536000)
-    __arg_info_matrix.append((["expiration_time", expiration_time, True, (int), True]))
+    expiration_time = kwargs.get("expiration_time", 31536000) # 31536000 seconds meaning 365 days.
+    __arg_info_matrix.append(["expiration_time", expiration_time, True, (int), True])
+
+    kid = kwargs.get("kid")
+    __arg_info_matrix.append(["kid", kid, True, (str), True])
+
+    auth_url = kwargs.get("auth_url", None)
+    __arg_info_matrix.append(["auth_url", auth_url, True, (str), True])
+
+    rest_client = kwargs.get("rest_client", "VECTORSTORE")
+    __arg_info_matrix.append(["rest_client", rest_client, True, (str), True, [svc.name for svc in TDServices]])
+
+    auth_mech = kwargs.get("auth_mech", None)
+    __arg_info_matrix.append(["auth_mech", auth_mech, True, (str), True, [mech.name for mech in AuthMechs]])
+
+    validate_jwt = kwargs.get("validate_jwt", True)
+    __arg_info_matrix.append(["validate_jwt", validate_jwt, True, (bool)])
+
+    valid_from = kwargs.get("valid_from", None)
+    __arg_info_matrix.append(["valid_from", valid_from, True, int])
 
     # Validate arguments.
     _Validators._validate_function_arguments(__arg_info_matrix)
 
+    # ---------------------------------BASE_URL PROCESSING------------------------------------------------------
+    # STEP 2: Process base_url/ues_url and set applicable config options.
+    # ----------------------------------------------------------------------------------------------------------
+
     # base_url should not end with 'open-analytics' or 'data-insights'
     if base_url:
         if base_url.endswith('open-analytics') or base_url.endswith('data-insights'):
             message = Messages.get_message(MessageCodes.ARG_NONE,
-                                           "base_url", "ending with 'data-insights' or 'open-analytics", None)
+                                           "base_url", "ending with 'data-insights' or 'open-analytics", "")
             raise TeradataMlException(message, MessageCodes.ARG_NONE)
 
         # Set the vector_store_base_url. This should only be done if base_url is set.
         # In case ues_url is set, vector_store_base_url should not be set.
         configure._vector_store_base_url = f'{base_url}/data-insights'
 
+    if ues_url:
+        # If incorrectly formatted UES service URL is passed, set it to None
+        # and let further validation raise error.
+        if not (ues_url.endswith('open-analytics') or ues_url.endswith('user-environment-service/api/v1/')):
+            ues_url = None
+
     # If ues_url is provided, then use it as base_url.
-    base_url = kwargs.get("ues_url", base_url)
+    base_url = ues_url if ues_url else base_url
+
+    if not (base_url or ues_url):
+        raise TeradataMlException(Messages.get_message(MessageCodes.MISSING_ARGS, ["base_url"]),
+                                  MessageCodes.MISSING_ARGS)
 
     # Set the OpenAF url.
-    # If ues_url is present use that otherwise generate it from base_url.
+    # If ues_url is present, then use that otherwise generate it from base_url.
     configure.ues_url = ues_url if ues_url else f'{base_url}/open-analytics'
 
-    # If user pass Auth token, set it.
-    if auth_token:
-        _InternalBuffer.add(auth_token=_AuthToken(token=auth_token))
-        return True
-
-    if client_id and any([pat_token, pem_file]):
-        message = Messages.get_message(MessageCodes.EITHER_THIS_OR_THAT_ARGUMENT,
-                                       "client_id", "pat_token' and 'pem_file")
-        raise TeradataMlException(message, MessageCodes.EITHER_THIS_OR_THAT_ARGUMENT)
-
-    if client_id is None:
-        if (pat_token and pem_file is None) or (pem_file and pat_token is None):
-            message = Messages.get_message(MessageCodes.MUST_PASS_ARGUMENT,
-                                           "pat_token", "pem_file")
-            raise TeradataMlException(message, MessageCodes.MUST_PASS_ARGUMENT)
-
-    # Check if pem file exists.
-    if pem_file is not None:
-        _Validators._validate_file_exists(pem_file)
-
-    # Extract the base URL.
+    # Extract the base URL and org id.
     url_parser = urlparse(base_url)
     parsed_base_url = "{}://{}".format(url_parser.scheme, url_parser.netloc)
-    netloc = url_parser.netloc.split('.')[0]
-
-    # Check if the authentication is PAT based or OAuth.
-    if all(arg is None for arg in [pat_token, pem_file]):
+    org_id = url_parser.netloc.split('.')[0]
+
+    # ---------------------------------TOKEN GENERATION------------------------------------------------------
+    # STEP 3: Based on auth_mech, generate authentication token data and store in _InternalBuffer.
+    # Note: auth_mech can be user-provided or can be derived from valid combination of supporting parameters.
+    # --------------------------------------------------------------------------------------------------------
+    if auth_mech:
+        auth_mech = auth_mech.lower()
+        if auth_mech == 'oauth':
+            pat_token = pem_file = password = auth_token = auth_url = None
+        elif auth_mech == 'jwt':
+            pat_token = pem_file = password = client_id = auth_url = None
+        elif auth_mech == 'basic':
+            pat_token = pem_file = auth_token = client_id = auth_url = None
+        elif auth_mech == 'pat':
+            password = client_id = auth_token = auth_url = None
+        elif auth_mech == 'keycloak':
+            pat_token = pem_file = auth_token = client_id = None
+
+    # Validate arguments for mutual exclusiveness.
+    all_groups_none = \
+        _Validators._validate_mutually_exclusive_argument_groups({"client_id": client_id},
+                                                                 {"auth_token": auth_token},
+                                                                 {"pat_token": pat_token,
+                                                                  "pem_file": pem_file},
+                                                                 {"password": password} if not auth_url else
+                                                                 {"password": password, "auth_url": auth_url},
+                                                                 return_all_falsy_status=True)
+
+    # Determine authentication mechanism from availability of supportive arguments.
+    if auth_mech is None:
+        if auth_token:
+            auth_mech = 'jwt'
+        elif any([pat_token, pem_file]):
+            auth_mech = 'pat'
+        elif auth_url:
+            auth_mech = 'keycloak'
+        elif password:
+            # Authentication is done via Basic authentication mechanism
+            # by passing 'basic' field in header.
+            auth_mech = 'basic'
+        # When all supporting arguments are None, default mechanism is OAuth.
+        elif client_id or all_groups_none:
+            auth_mech = 'oauth'
+
+    token_validated = False
+    # Generate and use authentication data as per authentication mechanism.
+    if auth_mech == 'jwt':
+        if not auth_token:
+            raise TeradataMlException(Messages.get_message(MessageCodes.MISSING_ARGS, ["auth_token"]),
+                                      MessageCodes.MISSING_ARGS)
+        # Validate JWT token if base_url points to CCP environment.
+        # TODO: Uncomment when mechanism to validate JWT for AI-On-prem system is available.
+        # if not ues_url:
+        #     token_validated = _validate_jwt_token(base_url, auth_token)
+
+        _InternalBuffer.add(auth_token=_AuthToken(token=auth_token,
+                                                  auth_type='bearer'))
+    elif auth_mech == 'oauth':
         configure._oauth = True
-        client_id = "{}-oaf-device".format(netloc) if client_id is None else client_id
+        client_id = "{}-oaf-device".format(org_id) if client_id is None else client_id
         da_wf = _DAWorkflow(parsed_base_url, client_id)
         token_data = da_wf._get_token_data()
 
@@ -1777,34 +2060,72 @@ def set_auth_token(base_url=None, client_id=None, pat_token=None, pem_file=None,
         configure._auth_token_expiry_time = time() + token_data["expires_in"] - 15
 
         # Store the jwt token in internal class attribute.
-        _InternalBuffer.add(auth_token=_AuthToken(token=token_data["access_token"]))
-
-    else:
-        configure._oauth = False
-
-        if username is None:
-            # If username is not specified then the database username associated with the current context will be
-            # considered.
-            username = _get_user()
-
-        org_id = netloc
+        _InternalBuffer.add(auth_token=_AuthToken(token=token_data["access_token"],
+                                                  auth_type='bearer'))
+    elif auth_mech == 'pat':
+        if any([pat_token, pem_file]):
+            _Validators._validate_mutually_inclusive_n_arguments(pat_token=pat_token,
+                                                                 pem_file=pem_file)
+        else:
+            raise TeradataMlException(Messages.get_message(MessageCodes.MISSING_ARGS, ["pat_token", "pem_file"]),
+                                      MessageCodes.MISSING_ARGS)
+
+        # Check if pem file exists.
+        if pem_file is not None:
+            _Validators._validate_file_exists(pem_file)
+
+        # Generate JWT token.
+        auth_wf = _AuthWorkflow({"base_url": parsed_base_url,
+                                 "org_id": org_id,
+                                 "pat_token": pat_token,
+                                 "pem_file": pem_file,
+                                 "username": username,
+                                 "expiration_time": expiration_time,
+                                 "kid": kid,
+                                 "valid_from": valid_from})
+        token_data = auth_wf._proxy_jwt()
 
-        # Construct a dictionary to be passed to _AuthWorkflow().
-        state_dict = {}
-        state_dict["base_url"] = parsed_base_url
-        state_dict["org_id"] = org_id
-        state_dict["pat_token"] = pat_token
-        state_dict["pem_file"] = pem_file
-        state_dict["username"] = username
-        state_dict["expiration_time"] = expiration_time
+        if validate_jwt:
+            # Validate generated JWT token.
+            token_validated = _validate_jwt_token(base_url, token_data)
 
-        auth_wf = _AuthWorkflow(state_dict)
-        token_data = auth_wf._proxy_jwt()
         # Store the jwt token in internal class attribute.
-        _InternalBuffer.add(auth_token=_AuthToken(token=token_data))
-    # If set_auth_token is triggered then it will be ccp_enabled = True.
-    # The function returns if we have just passed the auth_token, thus
-    # having ccp_enabled = False.
-    configure._ccp_enabled = True
+        _InternalBuffer.add(auth_token=_AuthToken(token=token_data,
+                                                  auth_type='bearer'))
+    elif auth_mech == 'basic':
+        if not password:
+            raise TeradataMlException(Messages.get_message(MessageCodes.MISSING_ARGS, ["password"]),
+                                      MessageCodes.MISSING_ARGS)
+        credentials = f"{username}:{password}"
+        # Encode the credentials string using Base64.
+        encoded_credentials = base64.b64encode(credentials.encode('utf-8')).decode('utf-8')
+        # Store the header data in internal class attribute.
+        _InternalBuffer.add(auth_token=_AuthToken(token=encoded_credentials,
+                                                  auth_type='basic'))
+    elif auth_mech == 'keycloak':
+        _Validators._validate_missing_required_arguments([["password", password, False, (str), True],
+                                                          ["auth_url", auth_url, False, (str), True]
+                                                          ])
+        token_generator = _KeycloakManager(auth_url=auth_url,
+                                           client_id=TDServices[rest_client].value)
+
+        # Store manager object in _InternalBuffer in order to generate token after expiry time.
+        _InternalBuffer.add(keycloak_manager=token_generator)
+        try:
+            token_data = token_generator.generate_token(username=username,
+                                                        password=password)
+        except:
+            raise TeradataMlException(Messages.get_message(MessageCodes.FUNC_EXECUTION_FAILED,
+                                                           "set_auth_token",
+                                                           "Failed to generate keycloak token."),
+                                      MessageCodes.FUNC_EXECUTION_FAILED)
+
+        _InternalBuffer.add(auth_token=_AuthToken(token=token_data,
+                                                  auth_type='keycloak'))
+
+    if token_validated:
+        print("Authentication token is generated, authenticated and set for the session.")
+    else:
+        print("Authentication token is generated and set for the session.")
 
     return True