def test_missing_file(self, mock_exists):
with self.assertRaisesRegex(
AirflowException,
re.escape(
"File a.json was not found. Check the configuration of your Secrets backend."
),
):
local_filesystem.load_connections("a.json")
def test_yaml_file_should_load_connection_extras(self, file_content, expected_extras):
connections_by_conn_id = local_filesystem.load_connections(
TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.yaml"))
connection_uris_by_conn_id = {
conn_id: [connection.extra_dejson for connection in connections]
for conn_id, connections in connections_by_conn_id.items()
}
self.assertEqual(expected_extras, connection_uris_by_conn_id)
def test_env_file_should_load_connection(self, file_content, expected_connection_uris):
connections_by_conn_id = local_filesystem.load_connections(
TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.env"))
connection_uris_by_conn_id = {
conn_id: [connection.get_uri() for connection in connections]
for conn_id, connections in connections_by_conn_id.items()
}
self.assertEqual(expected_connection_uris, connection_uris_by_conn_id)
def test_yaml_file_should_load_connection(self, file_content, expected_connection_uris):
with mock_local_file(yaml.dump(yaml.load(json.dumps(file_content)), default_flow_style=False)):
connections_by_conn_id = local_filesystem.load_connections("a.yaml")
connection_uris_by_conn_id = {
conn_id: [connection.get_uri() for connection in connections]
for conn_id, connections in connections_by_conn_id.items()
}
self.assertEqual(expected_connection_uris, connection_uris_by_conn_id)
def test_json_file_should_load_connection(self, file_content, expected_connection_uris):
with mock_local_file(json.dumps(file_content)):
connections_by_conn_id = local_filesystem.load_connections("a.json")
connection_uris_by_conn_id = {
conn_id: [connection.get_uri() for connection in connections]
for conn_id, connections in connections_by_conn_id.items()
}
self.assertEqual(expected_connection_uris, connection_uris_by_conn_id)
def test_yaml_file_should_load_connection_extras(self, file_content,
expected_extras):
with mock_local_file(file_content):
connections_by_conn_id = local_filesystem.load_connections(
"a.yaml")
connection_uris_by_conn_id = {
conn_id:
[connection.extra_dejson for connection in connections]
for conn_id, connections in connections_by_conn_id.items()
}
self.assertEqual(expected_extras, connection_uris_by_conn_id)
def test_env_file_should_load_connection(self, file_content,
expected_connection_uris):
with mock_local_file(file_content):
connections_by_conn_id = local_filesystem.load_connections("a.env")
connection_uris_by_conn_id = {
conn_id: [connection.get_uri() for connection in connections]
for conn_id, connections in six.iteritems(
connections_by_conn_id)
}
self.assertEqual(expected_connection_uris,
connection_uris_by_conn_id)
def connections_import(args):
"""Import new connections"""
try:
conns_map = load_connections(args.file)
except AirflowException:
sys.exit(1)
if not args.conflict_disposition:
disposition = DIS_RESTRICT
else:
disposition = args.conflict_disposition
conn_status_map = {DIS_OVERWRITE: [], DIS_IGNORE: [], CREATED: []}
try:
with create_session() as session:
for _, conn_list in conns_map.items():
for conn in conn_list:
conn_row = (session.query(Connection).filter(
Connection.conn_id == conn.conn_id).first())
if not conn_row:
session.add(conn)
conn_status_map[CREATED].append(conn)
elif disposition == DIS_OVERWRITE:
session.delete(conn_row)
session.flush()
session.add(conn)
conn_status_map[DIS_OVERWRITE].append(conn)
elif disposition == DIS_IGNORE:
conn_status_map[DIS_IGNORE].append(conn)
else:
msg = "\nConnection with `conn_id`={conn_ids} already exists"
msg = msg.format(conn_ids=conn.conn_id)
raise AirflowException(msg)
print(_prep_import_status_message(conn_status_map))
except (SQLAlchemyError, AirflowException) as e:
print(e)
session.rollback()
finally:
session.close()
def test_yaml_invalid_extra(self, file_content, expected_message):
with mock_local_file(file_content):
with self.assertRaisesRegex(AirflowException,
re.escape(expected_message)):
local_filesystem.load_connections("a.yaml")
def test_env_file_invalid_input(self, file_content,
expected_connection_uris):
with mock_local_file(json.dumps(file_content)):
with self.assertRaisesRegex(AirflowException,
re.escape(expected_connection_uris)):
local_filesystem.load_connections("a.json")
def test_env_file_invalid_format(self, content, expected_message):
with mock_local_file(content):
with self.assertRaisesRegex(AirflowFileParseException,
re.escape(expected_message)):
local_filesystem.load_connections("a.env")
def test_ensure_unique_connection_yaml(self, file_content):
with self.assertRaises(ConnectionNotUnique):
local_filesystem.load_connections(
TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.yaml"))
def test_yaml_invalid_extra(self, file_content, expected_message):
with self.assertRaisesRegex(AirflowException, re.escape(expected_message)):
local_filesystem.load_connections(
TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.yaml"))
def test_env_file_invalid_input(self, file_content, expected_connection_uris):
with self.assertRaisesRegex(AirflowException, re.escape(expected_connection_uris)):
local_filesystem.load_connections(
TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.json"))
def test_ensure_unique_connection_yaml(self, file_content):
with mock_local_file(file_content):
with self.assertRaises(ConnectionNotUnique):
local_filesystem.load_connections("a.yaml")
def test_ensure_unique_connection_json(self, file_content):
with mock_local_file(json.dumps(file_content)):
with self.assertRaises(ConnectionNotUnique):
local_filesystem.load_connections("a.json")
