Repository: aipotheosis-labs/aci
File: aipolabs/server/routes/linked_accounts.py (Python)

Code before:
        client_secret=app_default_oauth2_config.client_secret,
        scope=app_default_oauth2_config.scope,
        token_endpoint_auth_method=app_default_oauth2_config.token_endpoint_auth_method,
        authorize_url=app_default_oauth2_config.authorize_url,
        access_token_url=app_default_oauth2_config.access_token_url,
        refresh_token_url=app_default_oauth2_config.refresh_token_url,
        server_metadata_url=app_default_oauth2_config.server_metadata_url,
    )

    # get oauth2 account credentials
    # TODO: can each OAuth2 provider return different fields? if so, need to handle them accordingly. Maybe can
    # store the auth reponse schema in the App record in db. and cast the auth_response to the schema here.
    try:
        logger.info("retrieving oauth2 token")
        token_response = await oauth2.authorize_access_token_without_browser_session(
            oauth2_client, request, state.redirect_uri, state.code_verifier, state.nonce
        )
        # TODO: remove PII log
        logger.warning(f"oauth2 token requested successfully, token_response: {token_response}")
    except Exception:
        logger.exception("failed to retrieve oauth2 token")
        raise AuthenticationError("failed to retrieve oauth2 token during account linking")

    # TODO: we might want to verify scope authorized by end user (token_response["scope"]) is what we asked
    security_credentials = OAuth2SchemeCredentials(
        access_token=token_response["access_token"],
        token_type=token_response["token_type"],
        expires_at=(
            int(time.time()) + token_response["expires_in"]
            if "expires_in" in token_response
            else None
        ),
        refresh_token=(
            token_response["refresh_token"] if "refresh_token" in token_response else None
        ),
    )
    logger.info(
        f"security_credentials: \n {json.dumps(security_credentials.model_dump(mode='json'), indent=2)}"
    )

    # if the linked account already exists, update it, otherwise create a new one
    # TODO: consider separating the logic for updating and creating a linked account or give warning to clients
    # if the linked account already exists to avoid accidental overwriting the account
    # TODO: try/except, retry?
    linked_account = crud.linked_accounts.get_linked_account(
        db_session,
        state.project_id,
        state.app_name,
        state.linked_account_owner_id,
    )
    if linked_account:

Review comment (suggestion):
we can start using `token_response.get("token_type", None), token_response.get("refresh_token", None)` etc to avoid the verbose `if else`

Code after:
        client_secret=app_default_oauth2_config.client_secret,
        scope=app_default_oauth2_config.scope,
        token_endpoint_auth_method=app_default_oauth2_config.token_endpoint_auth_method,
        authorize_url=app_default_oauth2_config.authorize_url,
        access_token_url=app_default_oauth2_config.access_token_url,
        refresh_token_url=app_default_oauth2_config.refresh_token_url,
        server_metadata_url=app_default_oauth2_config.server_metadata_url,
    )

    # get oauth2 account credentials
    # TODO: can each OAuth2 provider return different fields? if so, need to handle them accordingly. Maybe can
    # store the auth reponse schema in the App record in db. and cast the auth_response to the schema here.
    try:
        logger.info("retrieving oauth2 token")
        token_response = await oauth2.authorize_access_token_without_browser_session(
            oauth2_client, request, state.redirect_uri, state.code_verifier, state.nonce
        )
        # TODO: remove PII log
        logger.warning(f"oauth2 token requested successfully, token_response: {token_response}")
    except Exception:
        logger.exception("failed to retrieve oauth2 token")
        raise AuthenticationError("failed to retrieve oauth2 token during account linking")

    # TODO: we might want to verify scope authorized by end user (token_response["scope"]) is what we asked
    security_credentials = OAuth2SchemeCredentials(
        access_token=token_response["access_token"],
        token_type=token_response.get("token_type", None),
        expires_at=(
            int(time.time()) + token_response["expires_in"]
            if "expires_in" in token_response
            else None
        ),
        refresh_token=token_response.get("refresh_token", None),
    )
    logger.info(
        f"security_credentials: \n {json.dumps(security_credentials.model_dump(mode='json'), indent=2)}"
    )

    # if the linked account already exists, update it, otherwise create a new one
    # TODO: consider separating the logic for updating and creating a linked account or give warning to clients
    # if the linked account already exists to avoid accidental overwriting the account
    # TODO: try/except, retry?
    linked_account = crud.linked_accounts.get_linked_account(
        db_session,
        state.project_id,
        state.app_name,
        state.linked_account_owner_id,
    )
    if linked_account:
        logger.info(f"updating oauth2 credentials for linked account={linked_account.id}")
        linked_account = crud.linked_accounts.update_linked_account(