Quellcode durchsuchen

[raz_adls] Refactor rename operation section (#2539)

- Move it to AdlsClient so that the core abfs.py is not cluttered with Raz related code.
- Update the tests to check the header getting signed with SAS separately.
- Fix pylint long lines issue.
Harsh Gupta vor 4 Jahren
Ursprung
Commit
271c77c296

+ 11 - 1
desktop/core/src/desktop/lib/raz/clients.py

@@ -36,7 +36,9 @@ class S3RazClient():
       u'Host': u'prakashmowdev1.s3-us-west-2.amazonaws.com',
       u'X-Amz-Security-Token': u'IQoJb3JpZ2luX2Vj...C',
       u'X-Amz-Date': u'20210604T102022Z',
-      u'Authorization': u'AWS4-HMAC-SHA256 Credential=ASIAYO3P24NAOAYMMDNN/20210604/us-west-2/s3/aws4_request, SignedHeaders=host;user-agent;x-amz-content-sha256;x-amz-date;x-amz-security-token, Signature=d341a194c2998c64b6fc726b69d0c3c2b97d520265f80df7e1bc1ac59a21ef94',
+      u'Authorization': u'AWS4-HMAC-SHA256 Credential=ASIAYO3P24NAOAYMMDNN/20210604/us-west-2/s3/aws4_request, 
+                          SignedHeaders=host;user-agent;x-amz-content-sha256;x-amz-date;x-amz-security-token, 
+                          Signature=d341a194c2998c64b6fc726b69d0c3c2b97d520265f80df7e1bc1ac59a21ef94',
       u'User-Agent': u'user:csso_romain'
     }
     '''
@@ -63,4 +65,12 @@ class AdlsRazClient():
       service='adls',
     )
 
+    # We need to sign the header source path separately for rename operation
+    if headers.get('x-ms-rename-source'):
+      partition_path = path.partition('.dfs.core.windows.net')
+      source_path = partition_path[0] + partition_path[1] + headers.get('x-ms-rename-source')
+
+      sas_token = c.check_access(method=action, url=source_path, headers=None)
+      headers['x-ms-rename-source'] += '?' + sas_token.get('token')
+
     return c.check_access(method=action, url=path, headers=headers)

+ 41 - 2
desktop/core/src/desktop/lib/raz/clients_test.py

@@ -14,14 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
 import unittest
 
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_equal, assert_false, assert_true, assert_raises
 
 from desktop.conf import RAZ
-from desktop.lib.raz.clients import S3RazClient
+from desktop.lib.raz.clients import S3RazClient, AdlsRazClient
 
+if sys.version_info[0] > 2:
+  from unittest.mock import patch, Mock
+else:
+  from mock import patch, Mock
 
 class S3RazClientLiveTest(unittest.TestCase):
 
@@ -40,7 +45,8 @@ class S3RazClientLiveTest(unittest.TestCase):
 
 
   def test_check_acccess_s3_list_file(self):
-    # e.g. 'https://gethue-test.s3.amazonaws.com/data/query-hive-weblogs.csv?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL&Signature=3lhK%2BwtQ9Q2u5VDIqb4MEpoY3X4%3D&Expires=1617207304'
+    # e.g. 'https://gethue-test.s3.amazonaws.com/data/query-hive-weblogs.csv?AWSAccessKeyId=AKIA23E77ZX2HVY76YGL&'
+    # 'Signature=3lhK%2BwtQ9Q2u5VDIqb4MEpoY3X4%3D&Expires=1617207304'
 
     url = S3RazClient().get_url(bucket='gethue-test', path='/data/query-hive-weblogs.csv')
 
@@ -58,3 +64,36 @@ class S3RazClientLiveTest(unittest.TestCase):
 
 
   def test_check_acccess_s3_list_file_no_access(self): pass
+
+class AdlsRazClientTest(unittest.TestCase):
+
+  def setUp(self):
+    self.username = 'csso_hueuser'
+  
+  def test_check_rename_operation(self):
+    with patch('desktop.lib.raz.raz_client.RazToken.get_delegation_token') as raz_token:
+      with patch('desktop.lib.raz.raz_client.requests.post') as requests_post:
+        with patch('desktop.lib.raz.raz_client.uuid.uuid4') as uuid:
+          with patch('desktop.lib.raz.raz_client.RazClient.check_access') as check_access:
+
+            reset = RAZ.API_URL.set_for_testing('https://raz_url:8000')
+            check_access.return_value = {'token': 'some_random_sas_token'}
+
+            try:
+              sas_token = AdlsRazClient(
+                username=self.username
+              ).get_url(
+                action='PUT',
+                path='https://gethuestorage.dfs.core.windows.net/data/user/csso_hueuser/rename_destination_dir',
+                headers={'x-ms-version': '2019-12-12', 'x-ms-rename-source': '/data/user/csso_hueuser/rename_source_dir'})
+
+              check_access.assert_called_with(
+                headers={
+                  'x-ms-version': '2019-12-12', 
+                  'x-ms-rename-source': '/data/user/csso_hueuser/rename_source_dir?some_random_sas_token'
+                },
+                method='PUT',
+                url='https://gethuestorage.dfs.core.windows.net/data/user/csso_hueuser/rename_destination_dir'
+              )
+            finally:
+              reset()

+ 0 - 16
desktop/core/src/desktop/lib/raz/raz_client_test.py

@@ -254,22 +254,6 @@ class RazClientTest(unittest.TestCase):
     access_type = client.handle_adls_req_mapping(method, url_params)
     assert_equal(access_type, 'set-permission')
 
-    # Rename
-    method = 'PUT'
-    relative_path = '/user/csso_hueuser/old_dir' # First call to fetch SAS to sign header path
-    url_params = {}
-
-    access_type = client.handle_adls_req_mapping(method, url_params)
-    assert_equal(access_type, 'rename-source')
-
-    method = 'PUT'
-    relative_path = '/user/csso_hueuser/new_dir' 
-    headers = {'x-ms-rename-source': '/user/csso_hueuser/old_dir?some_sas_token'} # Second call having signed header path
-    url_params = {}
-
-    access_type = client.handle_adls_req_mapping(method, url_params)
-    assert_equal(access_type, 'rename-source')
-
 
   def test_handle_relative_path(self):
     client = RazClient(self.raz_url, self.raz_token, username=self.username, service="adls", service_name="cm_adls", cluster_name="cl1")

+ 0 - 8
desktop/libs/azure/src/azure/abfs/abfs.py

@@ -573,14 +573,6 @@ class ABFS(object):
     """
     headers = {'x-ms-rename-source': '/' + urllib_quote(Init_ABFS.strip_scheme(old))}
 
-    # Required to sign the header with SAS token for RAZ
-    if RAZ.IS_ENABLED.get():
-      raz_http_client = RazHttpClient(self._user, self._url, exc_class=WebHdfsException, logger=LOG)
-      url = raz_http_client._make_url(headers['x-ms-rename-source'], params=None)
-
-      sas_token = raz_http_client.get_sas_token('PUT', self._user, url)
-      headers['x-ms-rename-source'] += '?' + sas_token
-
     try:
       self._create_path(new, headers=headers, overwrite=True)
     except WebHdfsException as e: