file-based: update transfer mode validations

airbytehq · Feb 11, 2025 · d4f2b1d · d4f2b1d
1 parent 750935c
commit d4f2b1d
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/airbyte_cdk/sources/file_based/config/validate_config_transfer_modes.py b/airbyte_cdk/sources/file_based/config/validate_config_transfer_modes.py
@@ -2,18 +2,23 @@
 # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
 #
 
-from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import AbstractFileBasedSpec
+from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import (
+    AbstractFileBasedSpec,
+    DeliverRawFiles,
+)
+from airbyte_cdk.sources.specs.transfer_modes import DeliverPermissions
 
 DELIVERY_TYPE_KEY = "delivery_type"
 DELIVERY_TYPE_PERMISSION_TRANSFER_MODE_VALUE = "use_permissions_transfer"
+DELIVERY_TYPE_FILES_TRANSFER_MODE_VALUE = "use_file_transfer"
 PRESERVE_DIRECTORY_STRUCTURE_KEY = "preserve_directory_structure"
 INCLUDE_IDENTITIES_STREAM_KEY = "include_identities_stream"
 
 
 def use_file_transfer(parsed_config: AbstractFileBasedSpec) -> bool:
     return (
         hasattr(parsed_config.delivery_method, DELIVERY_TYPE_KEY)
-        and parsed_config.delivery_method.delivery_type == "use_file_transfer"
+        and parsed_config.delivery_method.delivery_type == DELIVERY_TYPE_FILES_TRANSFER_MODE_VALUE
     )
 
 
@@ -33,7 +38,7 @@ def preserve_directory_structure(parsed_config: AbstractFileBasedSpec) -> bool:
     if (
         use_file_transfer(parsed_config)
         and hasattr(parsed_config.delivery_method, PRESERVE_DIRECTORY_STRUCTURE_KEY)
-        and parsed_config.delivery_method.preserve_directory_structure is not None
+        and isinstance(parsed_config.delivery_method, DeliverRawFiles)
     ):
         return parsed_config.delivery_method.preserve_directory_structure
     return True
@@ -51,7 +56,7 @@ def include_identities_stream(parsed_config: AbstractFileBasedSpec) -> bool:
     if (
         use_permissions_transfer(parsed_config)
         and hasattr(parsed_config.delivery_method, INCLUDE_IDENTITIES_STREAM_KEY)
-        and parsed_config.delivery_method.include_identities_stream is not None
+        and isinstance(parsed_config.delivery_method, DeliverPermissions)
     ):
         return parsed_config.delivery_method.include_identities_stream
     return False