Repository: incubator-airflow Updated Branches: refs/heads/master 26b747981 -> 7ece95686
[AIRFLOW-1593] expose load_string in WasbHook Closes #2596 from NielsZeilemaker/AIRFLOW-1593 Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/7ece9568 Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/7ece9568 Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/7ece9568 Branch: refs/heads/master Commit: 7ece9568684da50edb2f65aad8a5e98ba6b802e0 Parents: 26b7479 Author: Niels Zeilemaker <[email protected]> Authored: Wed Sep 13 10:15:01 2017 +0200 Committer: Bolke de Bruin <[email protected]> Committed: Wed Sep 13 10:15:01 2017 +0200 ---------------------------------------------------------------------- airflow/contrib/hooks/wasb_hook.py | 18 ++++++++++++++++++ tests/contrib/hooks/test_wasb_hook.py | 14 ++++++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7ece9568/airflow/contrib/hooks/wasb_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/wasb_hook.py b/airflow/contrib/hooks/wasb_hook.py index 89eaa5b..cc9922b 100644 --- a/airflow/contrib/hooks/wasb_hook.py +++ b/airflow/contrib/hooks/wasb_hook.py @@ -92,3 +92,21 @@ class WasbHook(BaseHook): # Reorder the argument order from airflow.hooks.S3_hook.load_file. self.connection.create_blob_from_path(container_name, blob_name, file_path, **kwargs) + + def load_string(self, string_data, container_name, blob_name, **kwargs): + """ + Upload a string to Azure Blob Storage. + + :param string_data: String to load. + :type string_data: str + :param container_name: Name of the container. + :type container_name: str + :param blob_name: Name of the blob. + :type blob_name: str + :param kwargs: Optional keyword arguments that + `BlockBlobService.create_blob_from_text()` takes. + :type kwargs: object + """ + # Reorder the argument order from airflow.hooks.S3_hook.load_string. + self.connection.create_blob_from_text(container_name, blob_name, + string_data, **kwargs) http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/7ece9568/tests/contrib/hooks/test_wasb_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_wasb_hook.py b/tests/contrib/hooks/test_wasb_hook.py index aa92937..df9cefe 100644 --- a/tests/contrib/hooks/test_wasb_hook.py +++ b/tests/contrib/hooks/test_wasb_hook.py @@ -87,7 +87,7 @@ class TestWasbHook(unittest.TestCase): self.assertTrue(hook.check_for_prefix('container', 'prefix', timeout=3)) mock_instance.list_blobs.assert_called_once_with( - 'container', 'prefix', timeout=3 + 'container', 'prefix', num_results=1, timeout=3 ) @mock.patch('airflow.contrib.hooks.wasb_hook.BlockBlobService', @@ -100,7 +100,7 @@ class TestWasbHook(unittest.TestCase): @mock.patch('airflow.contrib.hooks.wasb_hook.BlockBlobService', autospec=True) - def test_check_for_prefix(self, mock_service): + def test_load_file(self, mock_service): mock_instance = mock_service.return_value hook = WasbHook(wasb_conn_id='wasb_test_sas_token') hook.load_file('path', 'container', 'blob', max_connections=1) @@ -108,6 +108,16 @@ class TestWasbHook(unittest.TestCase): 'container', 'blob', 'path', max_connections=1 ) + @mock.patch('airflow.contrib.hooks.wasb_hook.BlockBlobService', + autospec=True) + def test_load_string(self, mock_service): + mock_instance = mock_service.return_value + hook = WasbHook(wasb_conn_id='wasb_test_sas_token') + hook.load_string('big string', 'container', 'blob', max_connections=1) + mock_instance.create_blob_from_text.assert_called_once_with( + 'container', 'blob', 'big string', max_connections=1 + ) + if __name__ == '__main__': unittest.main()
