Use correct paths in GcsBucket.put/get_directory (#14432)

Co-authored-by: nate nowack <thrast36@gmail.com>
PrefectHQ · Jul 5, 2024 · eb367af · eb367af
1 parent 200cf21
commit eb367af
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 9 deletions.
diff --git a/src/integrations/prefect-gcp/prefect_gcp/cloud_storage.py b/src/integrations/prefect-gcp/prefect_gcp/cloud_storage.py
@@ -677,7 +677,8 @@ async def get_directory(
             if blob_path[-1] == "/":
                 # object is a folder and will be created if it contains any objects
                 continue
-            local_file_path = os.path.join(local_path, blob_path)
+            relative_blob_path = os.path.relpath(blob_path, from_path)
+            local_file_path = os.path.join(local_path, relative_blob_path)
             os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
 
             with disable_run_logger():
@@ -739,7 +740,13 @@ async def put_directory(
                     PurePosixPath(to_path, local_file_path.relative_to(local_path))
                 )
                 local_file_content = local_file_path.read_bytes()
-                await self.write_path(remote_file_path, content=local_file_content)
+                with disable_run_logger():
+                    await cloud_storage_upload_blob_from_string.fn(
+                        data=local_file_content,
+                        bucket=self.bucket,
+                        blob=remote_file_path,
+                        gcp_credentials=self.gcp_credentials,
+                    )
                 uploaded_file_count += 1
 
         return uploaded_file_count

diff --git a/src/integrations/prefect-gcp/tests/test_cloud_storage.py b/src/integrations/prefect-gcp/tests/test_cloud_storage.py
@@ -157,6 +157,8 @@ def test_get_directory(self, gcs_bucket, tmp_path, from_path, local_path):
         prefix = os.path.join(gcs_bucket.bucket_folder, from_path or "")
         if local_path is None:
             local_path = os.path.abspath(".")
+        else:
+            local_path = os.path.abspath(os.path.expanduser(local_path))
 
         if from_path is None:
             from_path = gcs_bucket.bucket_folder
@@ -165,17 +167,28 @@ def test_get_directory(self, gcs_bucket, tmp_path, from_path, local_path):
         for file_path in actual:
             assert os.path.isfile(file_path)
 
-        # blob.txt, base_folder/nested_blob.txt, base_folder/sub_folder/nested_blob.txt
+        # check that returned file paths are correct
         if not prefix:
-            assert len(actual) == 4
-        # base_folder/sub_folder/nested_blob.txt
+            expected = [
+                os.path.join(local_path, "blob.txt"),
+                os.path.join(local_path, "base_folder/nested_blob.txt"),
+                os.path.join(local_path, "base_folder/sub_folder/nested_blob.txt"),
+                os.path.join(local_path, "dotted.folder/nested_blob.txt"),
+            ]
         elif prefix == "base_folder/sub_folder":
-            assert len(actual) == 1
-        # base_folder/nested_blob.txt, base_folder/sub_folder/nested_blob.txt
+            # base_folder/sub_folder/nested_blob.txt
+            expected = [
+                os.path.join(local_path, "nested_blob.txt"),
+            ]
         elif prefix == "base_folder" or prefix == "base_folder/":
-            assert len(actual) == 2
+            # base_folder/nested_blob.txt, base_folder/sub_folder/nested_blob.txt
+            expected = [
+                os.path.join(local_path, "nested_blob.txt"),
+                os.path.join(local_path, "sub_folder/nested_blob.txt"),
+            ]
         else:
-            assert len(actual) == 0
+            expected = []
+        assert actual == expected
 
     @pytest.mark.parametrize("to_path", [None, "to_path"])
     @pytest.mark.parametrize("ignore", [True, False])