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)