def get_composer_config():
"""Gets the configuration information of the Cloud Composer environment"""
logger.log(
logging.INFO,
f"Entered get_composer_config -- {API_BASE_PATH_V1}/composer/config api GET method"
)
req_data = request.get_json()
if req_data:
logger.log(logging.DEBUG,
f"Request contains a json payload, validating: {req_data}")
api_validator.validate_project_json(req_data)
project_id, location, composer_environment = api_service.get_gcp_composer_details(
req_data)
else:
logger.log(
logging.DEBUG,
f"Request does not contain a json payload, validating implicitly")
project_id, location, composer_environment = api_service.get_gcp_composer_details(
None)
authenticated_session = auth_service.get_authenticated_session()
airflow_config = airflow_service.AirflowService(
authenticated_session, project_id, location,
composer_environment).get_airflow_config()
next_actions = {
'list': f'{API_BASE_PATH_V1}/dag/list',
'validate': f'{API_BASE_PATH_V1}/dag/validate',
'deploy': f'{API_BASE_PATH_V1}/dag/deploy'
}
return jsonify(airflow_config=airflow_config, next_actions=next_actions)
def get_composer_experimental_apì():
"""Gets the configuration information of the Cloud Composer experimental API"""
logger.log(
logging.INFO,
f"Entered get_composer_experimental_apì -- {API_BASE_PATH_V1}/composer/api api GET method"
)
req_data = request.get_json()
if req_data:
logger.log(logging.DEBUG,
f"Request contains a json payload, validating: {req_data}")
api_validator.validate_project_json(req_data)
project_id, location, composer_environment = api_service.get_gcp_composer_details(
req_data)
else:
logger.log(
logging.DEBUG,
f"Request does not contain a json payload, validating implicitly")
project_id, location, composer_environment = api_service.get_gcp_composer_details(
None)
authenticated_session = auth_service.get_authenticated_session()
airflow_uri, client_id = airflow_service.AirflowService(
authenticated_session, project_id, location,
composer_environment).get_airflow_experimental_api()
return jsonify(airflow_uri=airflow_uri,
client_id=client_id,
next_actions=api_service.get_next_actions_experimental_api(
airflow_uri, client_id))
def test_get_airflow_experimental_api(self, mock_six_parse,
mock_requests_get,
mock_flask_response,
mock_auth_session):
with open(
os.path.join(os.path.dirname(Path(__file__)),
'airflow_config.json')) as json_file:
airflow_config = json.load(json_file)
# mock the objects
mock_flask_response.json.return_value = airflow_config
mock_six_parse.return_value = airflow_config['query_string']
mock_auth_session.request.return_value = mock_flask_response
headers = {
'location': airflow_config['query_string']['redirect_uri'][0]
}
mock_requests_get.return_value = MockResponse({}, 200, headers)
airflow_svc = airflow_service.AirflowService(
mock_auth_session, os.environ.get('PROJECT_ID'),
os.environ.get('GCP_LOCATION'),
os.environ.get('COMPOSER_ENVIRONMENT'))
airflow_ui, client_id = airflow_svc.get_airflow_experimental_api()
assert airflow_ui is not None
assert client_id is not None
assert airflow_ui == 'https://sde120c7fa68ea00ep-tp.appspot.com/api/experimental'
assert client_id == '401501771865-j04v42mav328ocngb267ts6mlh82j8uk.apps.googleusercontent.com'
def trigger_dag(dag_name):
"""Triggers a specific, existing dag within a Cloud Composer environment"""
logger.log(
logging.INFO,
f"Entered trigger_dag -- {API_BASE_PATH_V1}/dag/trigger/{dag_name} api PUT method"
)
req_data = request.get_json()
if req_data:
if 'conf' not in req_data or len(req_data['conf']) == 0:
return {
'error':
"JSON payload provided but conf element is missing or empty"
}, 500
if 'project_id' in req_data:
api_validator.validate_project_json(req_data)
project_id, location, composer_environment = api_service.get_gcp_composer_details(
req_data)
else:
project_id, location, composer_environment = api_service.get_gcp_composer_details(
None)
try:
res = api_service.trigger_dag(project_id, location,
composer_environment, dag_name)
except:
return {'error': traceback.print_exc()}, 500
authenticated_session = auth_service.get_authenticated_session()
airflow_uri, client_id = airflow_service.AirflowService(
authenticated_session, project_id, location,
composer_environment).get_airflow_experimental_api()
return jsonify(api_response=res,
next_actions=api_service.get_next_actions_experimental_api(
airflow_uri, client_id))
def test_get_airflow_config(self, mock_six_parse, mock_requests_get,
mock_flask_response, mock_auth_session):
with open(
os.path.join(os.path.dirname(Path(__file__)),
'airflow_config.json')) as json_file:
airflow_config = json.load(json_file)
# mock the objects
mock_flask_response.json.return_value = airflow_config
mock_six_parse.return_value = airflow_config['query_string']
mock_auth_session.request.return_value = mock_flask_response
headers = {
'location': airflow_config['query_string']['redirect_uri'][0]
}
mock_requests_get.return_value = MockResponse({}, 200, headers)
airflow_svc = airflow_service.AirflowService(
mock_auth_session, os.environ.get('PROJECT_ID'),
os.environ.get('GCP_LOCATION'),
os.environ.get('COMPOSER_ENVIRONMENT'))
composer_config = airflow_svc.get_airflow_config()
assert composer_config is not None
assert 'name' in composer_config
assert composer_config['name'] == airflow_config['name']
assert 'uuid' in composer_config
assert composer_config['uuid'] == airflow_config['uuid']
assert 'dagGcsPrefix' in composer_config['config']
assert composer_config['config']['dagGcsPrefix'] == airflow_config[
'config']['dagGcsPrefix']
assert 'client_id' in composer_config['query_string']
assert composer_config['query_string']['client_id'] == airflow_config[
'query_string']['client_id']
def test_trigger_dag(self, mock_request, mock_auth_service):
mock_auth_service.return_value = 'fake_id_token'
mock_request.return_value = MockResponse({}, 200, {},
"mock_response_text")
airflow_svc = airflow_service.AirflowService(
'fake_auth_session', os.environ.get('PROJECT_ID'),
os.environ.get('GCP_LOCATION'),
os.environ.get('COMPOSER_ENVIRONMENT'))
res_text = airflow_svc.trigger_dag('mock_dag_name', 'mock_airflow_ui',
',mock_client_id')
assert res_text is not None
assert res_text == "mock_response_text"
def __get_composer_environment(project_id, location, composer_environment):
"""
Gets the configuration information for a Cloud Composer environment.
Args:
project_id (string): GCP Project Id of the Cloud Composer instance
location (string): GCP Zone of the Cloud Composer instance
composer_environment (string): Name of the Cloud Composer instance
Returns:
an instance of composer.airflow.AirflowService
"""
logger.log(logging.DEBUG, "Getting the composer environment")
authenticated_session = auth_service.get_authenticated_session()
return airflow_service.AirflowService(authenticated_session, project_id,
location, composer_environment)
def test_get_next_actions_experimental_api():
authenticated_session = auth_service.get_authenticated_session()
airflow_uri, client_id = airflow_service.AirflowService(
authenticated_session,
os.environ.get('PROJECT_ID'),
os.environ.get('GCP_LOCATION'),
os.environ.get('COMPOSER_ENVIRONMENT')
).get_airflow_experimental_api()
next_actions = api_service.get_next_actions_experimental_api(airflow_uri, client_id)
assert 'endpoints' in next_actions
assert len(next_actions['endpoints']) > 0
assert 'description' in next_actions['endpoints'][0]
assert 'api_url' in next_actions['endpoints'][0]
assert 'http_method' in next_actions['endpoints'][0]
assert 'api_ref' in next_actions['endpoints'][0]
assert 'client_id' in next_actions['endpoints'][0]
def test_get_airflow_dag_gcs(self, mock_six_parse, mock_requests_get,
mock_flask_response, mock_auth_session):
with open(
os.path.join(os.path.dirname(Path(__file__)),
'airflow_config.json')) as json_file:
airflow_config = json.load(json_file)
# mock the objects
mock_flask_response.json.return_value = airflow_config
mock_six_parse.return_value = airflow_config['query_string']
mock_auth_session.request.return_value = mock_flask_response
headers = {
'location': airflow_config['query_string']['redirect_uri'][0]
}
mock_requests_get.return_value = MockResponse({}, 200, headers)
airflow_svc = airflow_service.AirflowService(
mock_auth_session, os.environ.get('PROJECT_ID'),
os.environ.get('GCP_LOCATION'),
os.environ.get('COMPOSER_ENVIRONMENT'))
dag_gcs_path = airflow_svc.get_airflow_dag_gcs()
assert dag_gcs_path is not None
assert dag_gcs_path == 'gs://europe-west3-composer-1b28efe1-bucket/dags'
