airflow-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From bo...@apache.org
Subject incubator-airflow git commit: [AIRFLOW-1797] S3Hook.load_string didn't work on Python3
Date Thu, 09 Nov 2017 19:47:39 GMT
Repository: incubator-airflow
Updated Branches:
  refs/heads/v1-9-test 8dea4a64e -> d592f891e


[AIRFLOW-1797] S3Hook.load_string didn't work on Python3

With the switch to Boto3 we now need the content
to be bytes, not a
string. On Python2 there is no difference, but for
Python3 this matters.

And since there were no real tests covering the
S3Hook I've added some
basic ones.

Closes #2771 from ashb/AIRFLOW-1797

(cherry picked from commit 28411b1e7eddb3338a329db3e52ee09de3676784)
Signed-off-by: Bolke de Bruin <bolke@xs4all.nl>


Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo
Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/d592f891
Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/d592f891
Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/d592f891

Branch: refs/heads/v1-9-test
Commit: d592f891e58650472c8fba89bace3cce54a7972b
Parents: 8dea4a6
Author: Ash Berlin-Taylor <ash_github@firemirror.com>
Authored: Thu Nov 9 20:47:08 2017 +0100
Committer: Bolke de Bruin <bolke@xs4all.nl>
Committed: Thu Nov 9 20:47:34 2017 +0100

----------------------------------------------------------------------
 airflow/hooks/S3_hook.py    |  7 ++--
 tests/core.py               | 20 -----------
 tests/hooks/test_s3_hook.py | 74 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 78 insertions(+), 23 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/d592f891/airflow/hooks/S3_hook.py
----------------------------------------------------------------------
diff --git a/airflow/hooks/S3_hook.py b/airflow/hooks/S3_hook.py
index f8052ca..b16566f 100644
--- a/airflow/hooks/S3_hook.py
+++ b/airflow/hooks/S3_hook.py
@@ -15,7 +15,7 @@
 from airflow.exceptions import AirflowException
 from airflow.contrib.hooks.aws_hook import AwsHook
 
-from six import StringIO
+from six import BytesIO
 from urllib.parse import urlparse
 import re
 import fnmatch
@@ -217,7 +217,8 @@ class S3Hook(AwsHook):
                     key, 
                     bucket_name=None,
                     replace=False,
-                    encrypt=False):
+                    encrypt=False,
+                    encoding='utf-8'):
         """
         Loads a string to S3
 
@@ -247,7 +248,7 @@ class S3Hook(AwsHook):
         if encrypt:
             extra_args['ServerSideEncryption'] = "AES256"
         
-        filelike_buffer = StringIO(string_data)
+        filelike_buffer = BytesIO(string_data.encode(encoding))
         
         client = self.get_conn()
         client.upload_fileobj(filelike_buffer, bucket_name, key, ExtraArgs=extra_args)

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/d592f891/tests/core.py
----------------------------------------------------------------------
diff --git a/tests/core.py b/tests/core.py
index c08d63d..0bd0c87 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -2386,26 +2386,6 @@ class HttpHookTest(unittest.TestCase):
         self.assertEqual(hook.base_url, 'https://localhost')
 
 
-try:
-    from airflow.hooks.S3_hook import S3Hook
-except ImportError:
-    S3Hook = None
-
-
-@unittest.skipIf(S3Hook is None,
-                 "Skipping test because S3Hook is not installed")
-class S3HookTest(unittest.TestCase):
-    def setUp(self):
-        configuration.load_test_config()
-        self.s3_test_url = "s3://test/this/is/not/a-real-key.txt"
-
-    def test_parse_s3_url(self):
-        parsed = S3Hook.parse_s3_url(self.s3_test_url)
-        self.assertEqual(parsed,
-                         ("test", "this/is/not/a-real-key.txt"),
-                         "Incorrect parsing of the s3 url")
-
-
 send_email_test = mock.Mock()
 
 

http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/d592f891/tests/hooks/test_s3_hook.py
----------------------------------------------------------------------
diff --git a/tests/hooks/test_s3_hook.py b/tests/hooks/test_s3_hook.py
new file mode 100644
index 0000000..48c9fde
--- /dev/null
+++ b/tests/hooks/test_s3_hook.py
@@ -0,0 +1,74 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import unittest
+
+from airflow import configuration
+
+try:
+    from airflow.hooks.S3_hook import S3Hook
+except ImportError:
+    S3Hook = None
+
+
+try:
+    import boto3
+    from moto import mock_s3
+except ImportError:
+    mock_s3 = None
+
+
+@unittest.skipIf(S3Hook is None,
+                 "Skipping test because S3Hook is not available")
+@unittest.skipIf(mock_s3 is None,
+                 "Skipping test because moto.mock_s3 is not available")
+class TestS3Hook(unittest.TestCase):
+    def setUp(self):
+        configuration.load_test_config()
+        self.s3_test_url = "s3://test/this/is/not/a-real-key.txt"
+
+    def test_parse_s3_url(self):
+        parsed = S3Hook.parse_s3_url(self.s3_test_url)
+        self.assertEqual(parsed,
+                         ("test", "this/is/not/a-real-key.txt"),
+                         "Incorrect parsing of the s3 url")
+
+    @mock_s3
+    def test_load_string(self):
+        hook = S3Hook(aws_conn_id=None)
+        conn = hook.get_conn()
+        # We need to create the bucket since this is all in Moto's 'virtual'
+        # AWS account
+        conn.create_bucket(Bucket="mybucket")
+
+        hook.load_string(u"Contént", "my_key", "mybucket")
+        body = boto3.resource('s3').Object('mybucket', 'my_key').get()['Body'].read()
+
+        self.assertEqual(body, b'Cont\xC3\xA9nt')
+
+    @mock_s3
+    def test_read_key(self):
+        hook = S3Hook(aws_conn_id=None)
+        conn = hook.get_conn()
+        # We need to create the bucket since this is all in Moto's 'virtual'
+        # AWS account
+        conn.create_bucket(Bucket='mybucket')
+        conn.put_object(Bucket='mybucket', Key='my_key', Body=b'Cont\xC3\xA9nt')
+
+        self.assertEqual(hook.read_key('my_key', 'mybucket'), u'Contént')
+
+
+if __name__ == '__main__':
+    unittest.main()


Mime
View raw message