def test_get_credentials_and_project_id_with_mutually_exclusive_configuration(
self,
):
with self.assertRaisesRegex(
AirflowException, re.escape('The `keyfile_dict` and `key_path` fields are mutually exclusive.')
):
get_credentials_and_project_id(key_path='KEY.json', keyfile_dict={'private_key': 'PRIVATE_KEY'})
def _get_credentials_and_project_id(
self) -> Tuple[google.auth.credentials.Credentials, Optional[str]]:
"""Returns the Credentials object for Google API and the associated project_id"""
if self._cached_credentials is not None:
return self._cached_credentials, self._cached_project_id
key_path: Optional[str] = self._get_field('key_path', None)
try:
keyfile_dict: Optional[str] = self._get_field('keyfile_dict', None)
keyfile_dict_json: Optional[Dict[str, str]] = None
if keyfile_dict:
keyfile_dict_json = json.loads(keyfile_dict)
except json.decoder.JSONDecodeError:
raise AirflowException('Invalid key JSON.')
target_principal, delegates = _get_target_principal_and_delegates(
self.impersonation_chain)
credentials, project_id = get_credentials_and_project_id(
key_path=key_path,
keyfile_dict=keyfile_dict_json,
scopes=self.scopes,
delegate_to=self.delegate_to,
target_principal=target_principal,
delegates=delegates,
)
overridden_project_id = self._get_field('project')
if overridden_project_id:
project_id = overridden_project_id
self._cached_credentials = credentials
self._cached_project_id = project_id
return credentials, project_id
def test_get_credentials_and_project_id_with_default_auth_and_scopes(self, scopes, mock_auth_default):
mock_credentials = mock.MagicMock()
mock_auth_default.return_value = (mock_credentials, self.test_project_id)
result = get_credentials_and_project_id(scopes=scopes)
mock_auth_default.assert_called_once_with(scopes=scopes)
self.assertEqual(mock_auth_default.return_value, result)
def __init__(
self,
connections_prefix: str = "airflow-connections",
variables_prefix: str = "airflow-variables",
config_prefix: str = "airflow-config",
gcp_keyfile_dict: Optional[dict] = None,
gcp_key_path: Optional[str] = None,
gcp_scopes: Optional[str] = None,
project_id: Optional[str] = None,
sep: str = "-",
**kwargs,
) -> None:
super().__init__(**kwargs)
self.connections_prefix = connections_prefix
self.variables_prefix = variables_prefix
self.config_prefix = config_prefix
self.sep = sep
if connections_prefix is not None:
if not self._is_valid_prefix_and_sep():
raise AirflowException(
"`connections_prefix`, `variables_prefix` and `sep` should "
f"follows that pattern {SECRET_ID_PATTERN}")
self.credentials, self.project_id = get_credentials_and_project_id(
keyfile_dict=gcp_keyfile_dict,
key_path=gcp_key_path,
scopes=gcp_scopes)
# In case project id provided
if project_id:
self.project_id = project_id
def client(self) -> hvac.Client:
"""
Return an authenticated Hashicorp Vault client
"""
_client = hvac.Client(url=self.url, **self.kwargs)
if self.auth_type == "token":
if not self.token:
raise VaultError("token cannot be None for auth_type='token'")
_client.token = self.token
elif self.auth_type == "ldap":
_client.auth.ldap.login(
username=self.username, password=self.password)
elif self.auth_type == "userpass":
_client.auth_userpass(username=self.username, password=self.password)
elif self.auth_type == "approle":
_client.auth_approle(role_id=self.role_id, secret_id=self.secret_id)
elif self.auth_type == "github":
_client.auth.github.login(token=self.token)
elif self.auth_type == "gcp":
from airflow.providers.google.cloud.utils.credentials_provider import (
get_credentials_and_project_id,
_get_scopes
)
scopes = _get_scopes(self.gcp_scopes)
credentials, _ = get_credentials_and_project_id(key_path=self.gcp_key_path, scopes=scopes)
_client.auth.gcp.configure(credentials=credentials)
else:
raise AirflowException(f"Authentication type '{self.auth_type}' not supported")
if _client.is_authenticated():
return _client
else:
raise VaultError("Vault Authentication Error!")
def test_get_credentials_and_project_id_with_default_auth_and_delegate(self, mock_auth_default):
mock_credentials = mock.MagicMock()
mock_auth_default.return_value = (mock_credentials, self.test_project_id)
result = get_credentials_and_project_id(delegate_to="USER")
mock_auth_default.assert_called_once_with(scopes=None)
mock_credentials.with_subject.assert_called_once_with("USER")
self.assertEqual((mock_credentials.with_subject.return_value, self.test_project_id), result)
def test_get_credentials_and_project_id_with_service_account_info(self, mock_from_service_account_info):
mock_from_service_account_info.return_value.project_id = self.test_project_id
service_account = {
'private_key': "PRIVATE_KEY"
}
result = get_credentials_and_project_id(keyfile_dict=service_account)
mock_from_service_account_info.assert_called_once_with(service_account, scopes=None)
self.assertEqual((mock_from_service_account_info.return_value, self.test_project_id), result)
def _auth_gcp(self, _client: hvac.Client) -> None:
# noinspection PyProtectedMember
from airflow.providers.google.cloud.utils.credentials_provider import (
get_credentials_and_project_id, _get_scopes)
scopes = _get_scopes(self.gcp_scopes)
credentials, _ = get_credentials_and_project_id(
key_path=self.gcp_key_path, scopes=scopes)
_client.auth.gcp.configure(credentials=credentials)
def test_get_credentials_and_project_id_with_service_account_file(self, mock_from_service_account_file):
mock_from_service_account_file.return_value.project_id = self.test_project_id
with self.assertLogs(level="DEBUG") as cm:
result = get_credentials_and_project_id(key_path=self.test_key_file)
mock_from_service_account_file.assert_called_once_with(self.test_key_file, scopes=None)
self.assertEqual((mock_from_service_account_file.return_value, self.test_project_id), result)
self.assertEqual([
'DEBUG:airflow.providers.google.cloud.utils.credentials_provider._CredentialProvider:Getting '
'connection using JSON key file KEY_PATH.json'
], cm.output)
def test_get_credentials_and_project_id_with_default_auth(self, mock_auth_default):
with self.assertLogs() as cm:
result = get_credentials_and_project_id()
mock_auth_default.assert_called_once_with(scopes=None)
self.assertEqual(("CREDENTIALS", "PROJECT_ID"), result)
self.assertEqual([
'INFO:airflow.providers.google.cloud.utils.credentials_provider._CredentialProvider:Getting '
'connection using `google.auth.default()` since no key file is defined for '
'hook.'
], cm.output)
def _client(self) -> gcp_logging.Client:
"""Google Cloud Library API client"""
credentials = get_credentials_and_project_id(
key_path=self.gcp_key_path,
scopes=self.scopes,
)
client = gcp_logging.Client(
credentials=credentials,
client_info=ClientInfo(client_library_version='airflow_v' +
version.version))
return client
def _auth_gcp(self, _client: hvac.Client) -> None:
from airflow.providers.google.cloud.utils.credentials_provider import ( # noqa
_get_scopes, get_credentials_and_project_id,
)
scopes = _get_scopes(self.gcp_scopes)
credentials, _ = get_credentials_and_project_id(key_path=self.gcp_key_path,
keyfile_dict=self.gcp_keyfile_dict,
scopes=scopes)
if self.auth_mount_point:
_client.auth.gcp.configure(credentials=credentials, mount_point=self.auth_mount_point)
else:
_client.auth.gcp.configure(credentials=credentials)
def client(self) -> SecretManagerServiceClient:
"""
Create an authenticated KMS client
"""
scopes = _get_scopes(self.gcp_scopes)
self.credentials, self.project_id = get_credentials_and_project_id(
key_path=self.gcp_key_path, scopes=scopes)
_client = SecretManagerServiceClient(
credentials=self.credentials,
client_info=ClientInfo(client_library_version='airflow_v' +
version.version))
return _client
def client(self) -> storage.Client:
"""Returns GCS Client."""
credentials, project_id = get_credentials_and_project_id(
key_path=self.gcp_key_path,
keyfile_dict=self.gcp_keyfile_dict,
scopes=self.scopes,
disable_logging=True)
return storage.Client(
credentials=credentials,
client_info=ClientInfo(client_library_version='airflow_v' +
version.version),
project=self.project_id if self.project_id else project_id)
def client(self) -> storage.Client:
"""Returns GCS Client."""
credentials, project_id = get_credentials_and_project_id(
key_path=self.gcp_key_path,
keyfile_dict=self.gcp_keyfile_dict,
scopes=self.scopes,
disable_logging=True,
)
return storage.Client(
credentials=credentials,
client_info=CLIENT_INFO,
project=self.project_id if self.project_id else project_id,
)
def test_get_credentials_and_project_id_with_service_account_info(self, mock_from_service_account_info):
mock_from_service_account_info.return_value.project_id = self.test_project_id
service_account = {
'private_key': "PRIVATE_KEY"
}
with self.assertLogs(level="DEBUG") as cm:
result = get_credentials_and_project_id(keyfile_dict=service_account)
mock_from_service_account_info.assert_called_once_with(service_account, scopes=None)
self.assertEqual((mock_from_service_account_info.return_value, self.test_project_id), result)
self.assertEqual([
'DEBUG:airflow.providers.google.cloud.utils.credentials_provider._CredentialProvider:Getting '
'connection using JSON Dict'
], cm.output)
def test_get_credentials_and_project_id_with_default_auth_and_target_principal(
self, mock_auth_default, mock_impersonated_credentials
):
mock_credentials = mock.MagicMock()
mock_auth_default.return_value = (mock_credentials, self.test_project_id)
result = get_credentials_and_project_id(target_principal="ACCOUNT_1")
mock_auth_default.assert_called_once_with(scopes=None)
mock_impersonated_credentials.assert_called_once_with(
source_credentials=mock_credentials,
target_principal="ACCOUNT_1",
delegates=None,
target_scopes=None,
)
self.assertEqual((mock_impersonated_credentials.return_value, self.test_project_id), result)
def test_get_credentials_and_project_id_with_default_auth_and_target_principal(
self, mock_auth_default, mock_impersonated_credentials):
mock_credentials = mock.MagicMock()
mock_auth_default.return_value = (mock_credentials,
self.test_project_id)
result = get_credentials_and_project_id(
target_principal=ACCOUNT_3_ANOTHER_PROJECT, )
mock_auth_default.assert_called_once_with(scopes=None)
mock_impersonated_credentials.assert_called_once_with(
source_credentials=mock_credentials,
target_principal=ACCOUNT_3_ANOTHER_PROJECT,
delegates=None,
target_scopes=None,
)
assert (mock_impersonated_credentials.return_value,
ANOTHER_PROJECT_ID) == result
def __init__(
self,
connections_prefix: str = "airflow-connections",
variables_prefix: str = "airflow-variables",
config_prefix: str = "airflow-config",
gcp_keyfile_dict: Optional[dict] = None,
gcp_key_path: Optional[str] = None,
gcp_scopes: Optional[str] = None,
project_id: Optional[str] = None,
sep: str = "-",
**kwargs,
) -> None:
super().__init__(**kwargs)
self.connections_prefix = connections_prefix
self.variables_prefix = variables_prefix
self.config_prefix = config_prefix
self.sep = sep
if connections_prefix is not None:
if not self._is_valid_prefix_and_sep():
raise AirflowException(
"`connections_prefix`, `variables_prefix` and `sep` should "
f"follows that pattern {SECRET_ID_PATTERN}"
)
try:
self.credentials, self.project_id = get_credentials_and_project_id(
keyfile_dict=gcp_keyfile_dict, key_path=gcp_key_path, scopes=gcp_scopes
)
except (DefaultCredentialsError, FileNotFoundError):
log.exception(
'Unable to load credentials for GCP Secret Manager. '
'Make sure that the keyfile path, dictionary, or GOOGLE_APPLICATION_CREDENTIALS '
'environment variable is correct and properly configured.'
)
# In case project id provided
if project_id:
self.project_id = project_id
def test_disable_logging(self, mock_default, mock_info, mock_file):
# assert not logs
with self.assertRaises(AssertionError), self.assertLogs(level="DEBUG"):
get_credentials_and_project_id(disable_logging=True)
# assert not logs
with self.assertRaises(AssertionError), self.assertLogs(level="DEBUG"):
get_credentials_and_project_id(
keyfile_dict={'private_key': 'PRIVATE_KEY'},
disable_logging=True,
)
# assert not logs
with self.assertRaises(AssertionError), self.assertLogs(level="DEBUG"):
get_credentials_and_project_id(
key_path='KEY.json',
disable_logging=True,
)
def test_get_credentials_and_project_id_with_service_account_file_and_non_valid_key(
self, _, file):
with self.assertRaises(AirflowException):
get_credentials_and_project_id(key_path=file)
def test_get_credentials_and_project_id_with_service_account_file(self, mock_from_service_account_file):
mock_from_service_account_file.return_value.project_id = self.test_project_id
result = get_credentials_and_project_id(key_path=self.test_key_file)
mock_from_service_account_file.assert_called_once_with(self.test_key_file, scopes=None)
self.assertEqual((mock_from_service_account_file.return_value, self.test_project_id), result)
def _credentials_and_project(self) -> Tuple[Credentials, str]:
credentials, project = get_credentials_and_project_id(
key_path=self.gcp_key_path, scopes=self.scopes, disable_logging=True
)
return credentials, project
def test_get_credentials_and_project_id_with_default_auth(self, mock_auth_default):
result = get_credentials_and_project_id()
mock_auth_default.assert_called_once_with(scopes=None)
self.assertEqual(("CREDENTIALS", "PROJECT_ID"), result)