def upload_output(self, output_dir_path): """Receives the tmp_dir_path where the files to upload are stored and uploads files whose name matches the prefixes and suffixes specified in 'output'.""" get_logger().info('Searching for files to upload in folder \'%s\'', output_dir_path) output_files = FileUtils.get_all_files_in_dir(output_dir_path) stg_providers = {} # Filter files by prefix and suffix for output in self.output: get_logger().info( 'Checking files for uploading to \'%s\' on path: \'%s\'', output['storage_provider'], output['path']) provider_type = StrUtils.get_storage_type( output['storage_provider']) provider_id = StrUtils.get_storage_id(output['storage_provider']) for file_path in output_files: file_name = file_path.replace(f'{output_dir_path}/', '') prefix_ok = False suffix_ok = False # Check prefixes if ('prefix' not in output or len(output['prefix']) == 0): prefix_ok = True else: for pref in output['prefix']: if file_name.startswith(pref): prefix_ok = True break if prefix_ok: # Check suffixes if ('suffix' not in output or len(output['suffix']) == 0): suffix_ok = True else: for suff in output['suffix']: if file_name.endswith(suff): suffix_ok = True break # Only upload file if name matches the prefixes and suffixes if suffix_ok: if provider_type not in stg_providers: stg_providers[provider_type] = {} if provider_id not in stg_providers[provider_type]: auth_data = self._get_auth_data( provider_type, provider_id) stg_providers[provider_type][ provider_id] = create_provider(auth_data) stg_providers[provider_type][provider_id].upload_file( file_path, file_name, output['path'])
def test_parse_config_no_output(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_NO_OUTPUT)}, clear=True): StorageConfig() self.assertLogs('There is no output defined for this function.', level='WARNING')
def test_parse_config_no_storage_provider(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_NO_STORAGE_PROVIDER)}, clear=True): StorageConfig() self.assertLogs('There is no storage provider defined for this function.', level='WARNING')
def execute_function(self): if SysUtils.is_var_in_env('SCRIPT'): script_path = SysUtils.join_paths( SysUtils.get_env_var("TMP_INPUT_DIR"), self._SCRIPT_FILE_NAME) script_content = StrUtils.base64_to_str( SysUtils.get_env_var('SCRIPT')) FileUtils.create_file_with_content(script_path, script_content) get_logger().info("Script file created in '%s'", script_path) FileUtils.set_file_execution_rights(script_path) get_logger().info("Executing user defined script: '%s'", script_path) try: pyinstaller_library_path = SysUtils.get_env_var( 'LD_LIBRARY_PATH') orig_library_path = SysUtils.get_env_var( 'LD_LIBRARY_PATH_ORIG') if orig_library_path: SysUtils.set_env_var('LD_LIBRARY_PATH', orig_library_path) self.output = subprocess.check_output( ['/bin/sh', script_path], stderr=subprocess.STDOUT).decode("latin-1") SysUtils.set_env_var('LD_LIBRARY_PATH', pyinstaller_library_path) get_logger().debug("CONTAINER OUTPUT:\n %s", self.output) except subprocess.CalledProcessError as cpe: # Exit with user script return code if an # error occurs (Kubernetes handles the error) get_logger().error(cpe.output.decode('latin-1')) sys.exit(cpe.returncode) else: get_logger().error('No user script found!')
def test_parse_config_valid(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_OK)}, clear=True): config = StorageConfig() expected_output = [ { 'storage_provider': 's3', 'path': 'bucket/folder' }, { 'storage_provider': 'minio', 'path': 'bucket', 'suffix': ['txt', 'jpg'], 'prefix': ['result-'] } ] self.assertEqual(config.output, expected_output) self.assertEqual(config.minio_auth.type, 'MINIO') self.assertEqual(config.minio_auth.get_credential('access_key'), 'test_minio_access') self.assertEqual(config.minio_auth.get_credential('secret_key'), 'test_minio_secret') self.assertEqual(config.onedata_auth.type, 'ONEDATA') self.assertEqual(config.onedata_auth.get_credential('oneprovider_host'), 'test_oneprovider.host') self.assertEqual(config.onedata_auth.get_credential('token'), 'test_onedata_token') self.assertEqual(config.onedata_auth.get_credential('space'), 'test_onedata_space') self.assertEqual(config.s3_auth.type, 'S3')
def test_parse_config_invalid_onedata(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_INVALID_ONEDATA)}, clear=True): with self.assertRaises(SystemExit): StorageConfig() self.assertLogs('The storage authentication of \'ONEDATA\' is not well-defined.', level='ERROR')
def test_get_minio_auth(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_OK)}, clear=True): minio_auth = StorageConfig().get_auth_data_by_stg_type('MINIO') self.assertEqual(minio_auth.type, 'MINIO') self.assertEqual(minio_auth.get_credential('access_key'), 'test_minio_access') self.assertEqual(minio_auth.get_credential('secret_key'), 'test_minio_secret')
def create_response(self): return { "statusCode": 200, "headers": { "amz-lambda-request-id": self.lambda_instance.get_request_id(), "amz-log-group-name": self.lambda_instance.get_log_group_name(), "amz-log-stream-name": self.lambda_instance.get_log_stream_name() }, "body": StrUtils.bytes_to_base64str(self.body["udocker_output"]), "isBase64Encoded": True, }
def test_get_onedata_auth(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_OK)}, clear=True): onedata_auth = StorageConfig().get_auth_data_by_stg_type('ONEDATA') self.assertEqual(onedata_auth.type, 'ONEDATA') self.assertEqual(onedata_auth.get_credential('oneprovider_host'), 'test_oneprovider.host') self.assertEqual(onedata_auth.get_credential('token'), 'test_onedata_token') self.assertEqual(onedata_auth.get_credential('space'), 'test_onedata_space')
def _get_script_path(self): script_path = None if SysUtils.is_var_in_env('SCRIPT'): script_path = SysUtils.join_paths( SysUtils.get_env_var("TMP_INPUT_DIR"), self._SCRIPT_FILE_NAME) script_content = StrUtils.base64_to_str( SysUtils.get_env_var('SCRIPT')) FileUtils.create_file_with_content(script_path, script_content) get_logger().info("Script file created in '%s'", script_path) elif FileUtils.is_file(self._OSCAR_SCRIPT_PATH): script_path = self._OSCAR_SCRIPT_PATH return script_path
def create_error_response(self): exception_msg = traceback.format_exc() get_logger().error("Exception launched:\n %s", exception_msg) return { "statusCode": 500, "headers": { "amz-lambda-request-id": self.lambda_instance.get_request_id(), "amz-log-group-name": self.lambda_instance.get_log_group_name(), "amz-log-stream-name": self.lambda_instance.get_log_stream_name() }, "body": StrUtils.dict_to_base64str({"exception" : exception_msg}), "isBase64Encoded": True, }
def _get_input_auth_data(self, parsed_event): """Return the proper auth data from a storage_provider based on the event. This methods allows to filter ONEDATA provider when multiple inputs are defined.""" storage_type = parsed_event.get_type() if storage_type == 'ONEDATA': # Check input path and event object_key if hasattr(parsed_event, 'object_key'): # Get the onedata space from the event object_key event_space = parsed_event.object_key.strip('/').split( '/', maxsplit=1)[0] for input in self.input: provider_type = StrUtils.get_storage_type( input.get('storage_provider')) if provider_type == storage_type: provider_id = StrUtils.get_storage_id( input.get('storage_provider')) if self.onedata_auth[provider_id].get_credential( 'space') == event_space: return self._get_auth_data(storage_type, provider_id) raise StorageAuthError(auth_type='ONEDATA') else: return self._get_auth_data(storage_type)
def _parse_exec_script_and_commands(self): # Check for script in function event if 'script' in self.raw_event: self.script_path = f"{self.input_folder}/script.sh" script_content = StrUtils.base64_to_str(self.raw_event['script']) FileUtils.create_file_with_content(self.script_path, script_content) # Container invoked with arguments elif 'cmd_args' in self.raw_event: # Add args self.cmd_args = json.loads(self.raw_event['cmd_args']) # Script to be executed every time (if defined) elif ConfigUtils.read_cfg_var('init_script') is not '': # Add init script self.init_script_path = f"{self.input_folder}/init_script.sh" FileUtils.cp_file(ConfigUtils.read_cfg_var('init_script'), self.init_script_path)
def _parse_exec_script_and_commands(self): # Check for script in function event if 'script' in self.raw_event: self.script_path = f"{self.input_folder}/script.sh" script_content = StrUtils.base64_to_str(self.raw_event['script']) FileUtils.create_file_with_content(self.script_path, script_content) # Container invoked with arguments elif 'cmd_args' in self.raw_event: # Add args self.cmd_args = json.loads(self.raw_event['cmd_args']) # Script to be executed every time (if defined) elif SysUtils.is_var_in_env('INIT_SCRIPT_PATH'): # Add init script self.init_script_path = f"{self.input_folder}/init_script.sh" FileUtils.cp_file(SysUtils.get_env_var("INIT_SCRIPT_PATH"), self.init_script_path)
def test_upload_output(self, mock_minio, mock_s3, mock_get_files): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_OK)}, clear=True): files = [ '/tmp/test/file1.txt', '/tmp/test/file1.jpg', '/tmp/test/result-file.txt', '/tmp/test/result-file.out', '/tmp/test/file2.txt', '/tmp/test/file2.out' ] mock_get_files.return_value = files StorageConfig().upload_output('/tmp/test') self.assertEqual(mock_minio.call_count, 1) self.assertEqual(mock_minio.call_args, call('/tmp/test/result-file.txt', 'result-file.txt', 'bucket')) self.assertEqual(mock_s3.call_count, 6) for i, f in enumerate(files): self.assertEqual(mock_s3.call_args_list[i], call(f, f.split('/')[3], 'bucket/folder'))
def test_base64_to_str(self): self.assertEqual(StrUtils.base64_to_str("dGVzdGluZwplbmNvZGUu"), "testing\nencode.")
def test_dict_to_base64str(self): self.assertEqual(StrUtils.dict_to_base64str({ "k1": "v1", "k2": "v2" }), "eyJrMSI6ICJ2MSIsICJrMiI6ICJ2MiJ9")
def test_bytes_to_base64str(self): self.assertEqual(StrUtils.bytes_to_base64str(b'testing\nencode.'), "dGVzdGluZwplbmNvZGUu")
def test_get_storage_id(self): self.assertEqual(StrUtils.get_storage_id('bad.good'), 'good') self.assertEqual(StrUtils.get_storage_id('bad.good.with.dots'), 'good.with.dots')
def test_get_s3_auth(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE_OK)}, clear=True): s3_auth = StorageConfig().get_auth_data_by_stg_type('S3') self.assertEqual(s3_auth.type, 'S3')
def test_get_storage_type(self): self.assertEqual(StrUtils.get_storage_type('good.bad.asdf'), 'GOOD')
def test_utf8_to_base64_string(self): self.assertEqual(StrUtils.utf8_to_base64_string("testing\nencode."), "dGVzdGluZwplbmNvZGUu")
def test_read_cfg_var_config_encoded(self): with mock.patch.dict('os.environ', {'FUNCTION_CONFIG': StrUtils.utf8_to_base64_string(CONFIG_FILE)}, clear=True): self.assertEqual(ConfigUtils.read_cfg_var('name'), 'test-func')