Is there a way to setup Airflow 2 with Google OAuth2 authentication

I have Airflow 1.10.14 nicely configured to allow organisation-level Google OAuth2 authentication.

I have recently tried to upgrade to Airflow 2 but I could not find a way to make Google OAuth2 work.

This is how my webserver_config.py looks like:

import os

from flask_appbuilder.security.manager import AUTH_OAUTH

basedir = os.path.abspath(os.path.dirname(__file__))

# Flask-WTF flag for CSRF
WTF_CSRF_ENABLED = True

AUTH_TYPE = AUTH_OAUTH

OAUTH_PROVIDERS = [{
   'name':'google',
     'token_key':'access_token',
     'icon':'fa-google',
         'remote_app': {
             'api_base_url':'https://www.googleapis.com/oauth2/v2/',
             'client_kwargs':{
                 'scope': 'email profile'
             },
             'access_token_url':'https://accounts.google.com/o/oauth2/token',
             'authorize_url':'https://accounts.google.com/o/oauth2/auth',
             'request_token_url': None,
             'client_id': '<client_id>',
             'client_secret': '<client_secret>',
         }
}]

As redirect URL I have https://airflow.mydomain.tld/oauth-authorized/google I have configured scope and enabled the .../auth/userinfo.email and .../auth/userinfo.profile.

With this setup I get the Google consent page, and once I pick my user (created in Airflow 2 with "Admin" privilege) I get:

Something bad has happened.
Please consider letting us know by creating a bug report using GitHub.

Python version: 3.8.5
Airflow version: 2.1.4
Node: ip-x-x-x-x.ec2.internal
-------------------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/app.py", line 2447, in wsgi_app
    response = self.full_dispatch_request()
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/app.py", line 1952, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/app.py", line 1821, in handle_user_exception
    reraise(exc_type, exc_value, tb)
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/_compat.py", line 39, in reraise
    raise value
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/app.py", line 1950, in full_dispatch_request
    rv = self.dispatch_request()
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask/app.py", line 1936, in dispatch_request
    return self.view_functions[rule.endpoint](**req.view_args)
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/flask_appbuilder/security/views.py", line 655, in oauth_authorized
    resp = self.appbuilder.sm.oauth_remotes[provider].authorize_access_token()
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/authlib/integrations/flask_client/remote_app.py", line 74, in authorize_access_token
    params = self.retrieve_access_token_params(flask_req, request_token)
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/authlib/integrations/base_client/base_app.py", line 145, in retrieve_access_token_params
    params = self._retrieve_oauth2_access_token_params(request, params)
  File "/home/airflow/py/airflow2/lib64/python3.8/site-packages/authlib/integrations/base_client/base_app.py", line 126, in _retrieve_oauth2_access_token_params
    raise MismatchingStateError()
authlib.integrations.base_client.errors.MismatchingStateError: mismatching_state: CSRF Warning! State not equal in request and response.

I tried with different users, associated with different roles, with the same result.

You can resolve using the webserver_config.py like:

import os
from flask_appbuilder.security.manager import AUTH_OAUTH
from airflow.configuration import conf

basedir = os.path.abspath(os.path.dirname(__file__))

# The SQLAlchemy connection string.
AUTH_TYPE = AUTH_OAUTH

AUTH_USER_REGISTRATION = True

AUTH_USER_REGISTRATION_ROLE = "Admin"

CSRF_ENABLED = True

# The SQLAlchemy connection string.
SQLALCHEMY_DATABASE_URI = conf.get('core', 'SQL_ALCHEMY_CONN')

OAUTH_PROVIDERS = [
    {'name': 'google', 'icon': 'fa-google', 'token_key': 'access_token',
     'remote_app': {
         'client_id': '<client_id>',
         'client_secret': '<client_secret>',
         'api_base_url': 'https://www.googleapis.com/oauth2/v2/',
         'client_kwargs': {
             'scope': 'email profile'
         },
         'request_token_url': None,
         'access_token_url': 'https://accounts.google.com/o/oauth2/token',
         'authorize_url': 'https://accounts.google.com/o/oauth2/auth'}
     },
]

And in airflow.cfg you have to specify cookie_samesite as:

# Set samesite policy on session cookie
cookie_samesite = Lax

For organisation level security you could also add in the OAUTH_PROVIDERS:

'whitelist': ['@<your-organsation-email.com>']