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_variables("a.json")
def variables_import(args): """Imports variables from a given file""" try: vars_map = load_variables(args.file) except AirflowException as e: print(e) sys.exit(1) if not args.conflict_disposition: disposition = DIS_RESTRICT else: disposition = args.conflict_disposition var_status_map = {DIS_OVERWRITE: [], DIS_IGNORE: [], CREATED: []} try: with create_session() as session: for key, val in vars_map.items(): serialize = isinstance(val, (list, dict)) vars_row = Variable.get_variable_from_secrets(key) if not vars_row: Variable.set(key=key, value=val, session=session, serialize_json=serialize) session.flush() var_status_map[CREATED].append(key) elif disposition == DIS_OVERWRITE: Variable.set(key=key, value=val, session=session, serialize_json=serialize) session.flush() var_status_map[DIS_OVERWRITE].append(key) elif disposition == DIS_IGNORE: var_status_map[DIS_IGNORE].append(key) else: msg = "\nVariable with `key`={key} already exists" msg = msg.format(key=key) raise AirflowException(msg) print(_prep_import_status_msgs(var_status_map)) except (SQLAlchemyError, AirflowException) as e: print(e) session.rollback() finally: session.close()
def test_env_file_invalid_logic(self, content, expected_message): with mock_local_file(content): with self.assertRaisesRegex(AirflowException, re.escape(expected_message)): local_filesystem.load_variables("a.env")
def test_env_file_should_load_variables(self, file_content, expected_variables): with mock_local_file(file_content): variables = local_filesystem.load_variables("a.env") self.assertEqual(expected_variables, variables)
def test_json_file_invalid_format(self, content, expected_message): with mock_local_file(content): with self.assertRaisesRegex(AirflowFileParseException, re.escape(expected_message)): local_filesystem.load_variables("a.json")
def test_json_file_should_load_variables(self, file_content, expected_variables): with mock_local_file(json.dumps(file_content)): variables = local_filesystem.load_variables("a.json") self.assertEqual(expected_variables, variables)
def test_env_file_invalid_logic(self, content, expected_message): with mock_local_file(content): with pytest.raises(AirflowException, match=re.escape(expected_message)): local_filesystem.load_variables("a.env")
def test_yaml_file_should_load_variables(self, file_content, expected_variables): with mock_local_file(file_content): variables = local_filesystem.load_variables('a.yaml') assert expected_variables == variables
def test_json_file_should_load_variables(self, file_content, expected_variables): variables = local_filesystem.load_variables( TextIOWrapper(io.BytesIO(file_content.encode('ascii')), name="a.json")) self.assertEqual(expected_variables, variables)
def test_env_file_invalid_logic(self, content, expected_message): with self.assertRaisesRegex(AirflowException, re.escape(expected_message)): local_filesystem.load_variables(TextIOWrapper(io.BytesIO(content.encode('ascii')), name="a.env"))
def test_yaml_file_should_load_variables(self, file_content, expected_variables): with mock_local_file(yaml.dump(yaml.load(json.dumps(file_content)), default_flow_style=False)): variables = local_filesystem.load_variables('a.yaml') self.assertEqual(expected_variables, variables)