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")