Add support for S3 dag bundle

Author: ismailsimsek

airflow/dag_processing/bundles/s3.py

@@ -116,25 +116,26 @@ def _delete_stale_local_files(self, current_s3_objects: List[Path]):
 
     def _download_s3_object_if_changed(self, s3_bucket, s3_object, local_target_path: Path):
         should_download = False
+        download_msg = ""
         if not local_target_path.exists():
             should_download = True
-            self.log.debug(f"Local file {local_target_path} does not exist. Downloading.")
+            download_msg = f"Local file {local_target_path} does not exist."
         else:
             local_stats = local_target_path.stat()
 
             if s3_object.size != local_stats.st_size:
                 should_download = True
-                self.log.debug(f"S3 object size ({s3_object.size}) and local file size ({local_stats.st_size}) differ. Downloading.")
+                download_msg = f"S3 object size ({s3_object.size}) and local file size ({local_stats.st_size}) differ."
 
             s3_last_modified = s3_object.last_modified
             local_last_modified = datetime.fromtimestamp(local_stats.st_mtime, tz=timezone.utc)
             if s3_last_modified.replace(microsecond=0) != local_last_modified.replace(microsecond=0):
                 should_download = True
-                self.log.debug(f"S3 object last modified ({s3_last_modified}) and local file last modified ({local_last_modified}) differ. Downloading.")
+                download_msg = f"S3 object last modified ({s3_last_modified}) and local file last modified ({local_last_modified}) differ."
 
         if should_download:
             s3_bucket.download_file(s3_object.key, local_target_path)
-            self.log.debug(f"Downloaded {s3_object.key} to {local_target_path.as_posix()}")
+            self.log.debug(f"{download_msg} Downloaded {s3_object.key} to {local_target_path.as_posix()}")
         else:
             self.log.debug(
                 f"Local file {local_target_path.as_posix()} is up-to-date with S3 object {s3_object.key}. Skipping download."

tests/dag_processing/bundles/test_s3.py

@@ -205,19 +205,29 @@ def test_refresh(self, s3_bucket, s3_client, caplog):
             level="DEBUG",
             regex=rf"Downloaded.*{S3_BUCKET_PREFIX}.*dag_03.py.*/s3/{bundle.name}/dag_03.py",
         )
+        assert bundle.bare_repo_path.joinpath("dag_03.py").read_text() == "test data"
         bundle.bare_repo_path.joinpath("dag_should_be_deleted.py").write_text("test dag")
         bundle.bare_repo_path.joinpath("dag_should_be_deleted_folder").mkdir(exist_ok=True)
+        s3_client.put_object(
+            Bucket=s3_bucket.name, Key=S3_BUCKET_PREFIX + "/dag_03.py", Body="test data-changed".encode("utf-8")
+        )
         bundle.refresh()
         assert caplog.text.count("Downloading DAGs from s3") == 4
         self.assert_log_matches_regex(
             caplog=caplog,
             level="DEBUG",
-            regex=rf"Deleted stale empty directory.*dag_should_be_deleted_folder",
+            regex=rf"S3 object size.*and local file size.*differ.*Downloaded.*dag_03.py.*",
+        )
+        assert bundle.bare_repo_path.joinpath("dag_03.py").read_text() == "test data-changed"
+        self.assert_log_matches_regex(
+            caplog=caplog,
+            level="DEBUG",
+            regex=rf"Deleted stale empty directory.*dag_should_be_deleted_folder.*",
         )
         self.assert_log_matches_regex(
             caplog=caplog,
             level="DEBUG",
-            regex=rf"Deleted stale local file.*dag_should_be_deleted.py",
+            regex=rf"Deleted stale local file.*dag_should_be_deleted.py.*",
         )
 
     def assert_log_matches_regex(self, caplog, level, regex):
@@ -229,4 +239,4 @@ def assert_log_matches_regex(self, caplog, level, regex):
                 break  # Stop searching once a match is found
         assert (
             matched
-        ), f"No log message at level {level} matching regex '{regex}' found. Logged messages:\n{caplog.text}"
+        ), f"No log message at level {level} matching regex '{regex}' found."