airbyte-source-azure-blob-storage 0.4.0__tar.gz → 0.4.2__tar.gz

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-source-azure-blob-storage
-Version: 0.4.0
+Version: 0.4.2
 Summary: Source implementation for Azure Blob Storage.
 Home-page: https://airbyte.com
 License: MIT

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "0.4.0"
+version = "0.4.2"
 name = "airbyte-source-azure-blob-storage"
 description = "Source implementation for Azure Blob Storage."
 authors = [
@@ -41,6 +41,7 @@ source-azure-blob-storage = "source_azure_blob_storage.run:run"
 
 [tool.poetry.group.dev.dependencies]
 docker = "^7.0.0"
+freezegun = "^1.4.0"
 pytest-mock = "^3.6.1"
 requests-mock = "^1.9.3"
 pandas = "2.2.1"

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/source_azure_blob_storage/__init__.py

@@ -3,8 +3,8 @@
 #
 
 
-from .config import Config
 from .source import SourceAzureBlobStorage
+from .spec import SourceAzureBlobStorageSpec
 from .stream_reader import SourceAzureBlobStorageStreamReader
 
-__all__ = ["SourceAzureBlobStorage", "SourceAzureBlobStorageStreamReader", "Config"]
+__all__ = ["SourceAzureBlobStorage", "SourceAzureBlobStorageStreamReader", "SourceAzureBlobStorageSpec"]

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/source_azure_blob_storage/config_migrations.py

@@ -4,6 +4,7 @@
 
 
 import logging
+from abc import ABC, abstractmethod
 from typing import Any, List, Mapping
 
 from airbyte_cdk.config_observation import create_connector_config_control_message
@@ -13,22 +14,16 @@ from airbyte_cdk.sources import Source
 logger = logging.getLogger("airbyte_logger")
 
 
-class MigrateCredentials:
-    """
-    This class stands for migrating the config azure_blob_storage_account_key inside object `credentials`
-    """
-
+class MigrateConfig(ABC):
     @classmethod
+    @abstractmethod
     def should_migrate(cls, config: Mapping[str, Any]) -> bool:
-        return "credentials" not in config
+        ...
 
     @classmethod
-    def set_azure_blob_storage_account_key(cls, config: Mapping[str, Any]) -> Mapping[str, Any]:
-        config["credentials"] = {
-            "auth_type": "storage_account_key",
-            "azure_blob_storage_account_key": config.pop("azure_blob_storage_account_key"),
-        }
-        return config
+    @abstractmethod
+    def migrate_config(cls, config: Mapping[str, Any]) -> Mapping[str, Any]:
+        ...
 
     @classmethod
     def modify_and_save(cls, config_path: str, source: Source, config: Mapping[str, Any]) -> Mapping[str, Any]:
@@ -43,7 +38,7 @@ class MigrateCredentials:
         Returns:
         - Mapping[str, Any]: The updated configuration.
         """
-        migrated_config = cls.set_azure_blob_storage_account_key(config)
+        migrated_config = cls.migrate_config(config)
         source.write_config(migrated_config, config_path)
         return migrated_config
 
@@ -75,3 +70,50 @@ class MigrateCredentials:
             config = source.read_config(config_path)
             if cls.should_migrate(config):
                 cls.emit_control_message(cls.modify_and_save(config_path, source, config))
+
+
+class MigrateLegacyConfig(MigrateConfig):
+    """
+    Class that takes in Azure Blob Storage source configs in the legacy format and transforms them into
+    configs that can be used by the new Azure Blob Storage source built with the file-based CDK.
+    """
+
+    @classmethod
+    def should_migrate(cls, config: Mapping[str, Any]) -> bool:
+        return "streams" not in config
+
+    @classmethod
+    def migrate_config(cls, legacy_config: Mapping[str, Any]) -> Mapping[str, Any]:
+        azure_blob_storage_blobs_prefix = legacy_config.get("azure_blob_storage_blobs_prefix", "")
+        return {
+            "azure_blob_storage_endpoint": legacy_config.get("azure_blob_storage_endpoint", None),
+            "azure_blob_storage_account_name": legacy_config["azure_blob_storage_account_name"],
+            "azure_blob_storage_account_key": legacy_config["azure_blob_storage_account_key"],
+            "azure_blob_storage_container_name": legacy_config["azure_blob_storage_container_name"],
+            "streams": [
+                {
+                    "name": legacy_config["azure_blob_storage_container_name"],
+                    "legacy_prefix": azure_blob_storage_blobs_prefix,
+                    "validation_policy": "Emit Record",
+                    "format": {"filetype": "jsonl"},
+                }
+            ],
+        }
+
+
+class MigrateCredentials(MigrateConfig):
+    """
+    This class stands for migrating the config azure_blob_storage_account_key inside object `credentials`
+    """
+
+    @classmethod
+    def should_migrate(cls, config: Mapping[str, Any]) -> bool:
+        return "credentials" not in config
+
+    @classmethod
+    def migrate_config(cls, config: Mapping[str, Any]) -> Mapping[str, Any]:
+        config["credentials"] = {
+            "auth_type": "storage_account_key",
+            "azure_blob_storage_account_key": config.pop("azure_blob_storage_account_key"),
+        }
+        return config

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/source_azure_blob_storage/run.py

@@ -9,8 +9,8 @@ from datetime import datetime
 from airbyte_cdk.entrypoint import AirbyteEntrypoint, launch
 from airbyte_cdk.models import AirbyteErrorTraceMessage, AirbyteMessage, AirbyteTraceMessage, TraceType, Type
 from airbyte_cdk.sources.file_based.stream.cursor import DefaultFileBasedCursor
-from source_azure_blob_storage import Config, SourceAzureBlobStorage, SourceAzureBlobStorageStreamReader
-from source_azure_blob_storage.config_migrations import MigrateCredentials
+from source_azure_blob_storage import SourceAzureBlobStorage, SourceAzureBlobStorageSpec, SourceAzureBlobStorageStreamReader
+from source_azure_blob_storage.config_migrations import MigrateCredentials, MigrateLegacyConfig
 
 
 def run():
@@ -21,12 +21,13 @@ def run():
     try:
         source = SourceAzureBlobStorage(
             SourceAzureBlobStorageStreamReader(),
-            Config,
+            SourceAzureBlobStorageSpec,
             SourceAzureBlobStorage.read_catalog(catalog_path) if catalog_path else None,
             SourceAzureBlobStorage.read_config(config_path) if catalog_path else None,
             SourceAzureBlobStorage.read_state(state_path) if catalog_path else None,
             cursor_cls=DefaultFileBasedCursor,
         )
+        MigrateLegacyConfig.migrate(sys.argv[1:], source)
         MigrateCredentials.migrate(sys.argv[1:], source)
     except Exception:
         print(

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/source_azure_blob_storage/source.py

@@ -2,35 +2,14 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from typing import Any, Mapping
+from typing import Any
 
-from airbyte_cdk.config_observation import emit_configuration_as_airbyte_control_message
 from airbyte_cdk.sources.declarative.models import OAuthConfigSpecification
 from airbyte_cdk.sources.file_based.file_based_source import FileBasedSource
 from airbyte_protocol.models import AdvancedAuth, ConnectorSpecification
 
-from .legacy_config_transformer import LegacyConfigTransformer
-
 
 class SourceAzureBlobStorage(FileBasedSource):
-    @classmethod
-    def read_config(cls, config_path: str) -> Mapping[str, Any]:
-        """
-        Used to override the default read_config so that when the new file-based Azure Blob Storage connector processes a config
-        in the legacy format, it can be transformed into the new config. This happens in entrypoint before we
-        validate the config against the new spec.
-        """
-        config = FileBasedSource.read_config(config_path)
-        if not cls._is_v1_config(config):
-            converted_config = LegacyConfigTransformer.convert(config)
-            emit_configuration_as_airbyte_control_message(converted_config)
-            return converted_config
-        return config
-
-    @staticmethod
-    def _is_v1_config(config: Mapping[str, Any]) -> bool:
-        return "streams" in config
-
     def spec(self, *args: Any, **kwargs: Any) -> ConnectorSpecification:
         """
         Returns the specification describing what fields can be configured by a user when setting up a file-based source.

airbyte_source_azure_blob_storage-0.4.0/source_azure_blob_storage/config.py → airbyte_source_azure_blob_storage-0.4.2/source_azure_blob_storage/spec.py

@@ -49,7 +49,7 @@ class StorageAccountKey(BaseModel):
     )
 
 
-class Config(AbstractFileBasedSpec):
+class SourceAzureBlobStorageSpec(AbstractFileBasedSpec):
     """
     NOTE: When this Spec is changed, legacy_config_transformer.py must also be modified to uptake the changes
     because it is responsible for converting legacy Azure Blob Storage v0 configs into v1 configs using the File-Based CDK.
@@ -59,17 +59,17 @@ class Config(AbstractFileBasedSpec):
     def documentation_url(cls) -> AnyUrl:
         return AnyUrl("https://docs.airbyte.com/integrations/sources/azure-blob-storage", scheme="https")
 
-    azure_blob_storage_account_name: str = Field(
-        title="Azure Blob Storage account name",
-        description="The account's name of the Azure Blob Storage.",
-        examples=["airbyte5storage"],
-        order=2,
-    )
     credentials: Union[Oauth2, StorageAccountKey] = Field(
         title="Authentication",
         description="Credentials for connecting to the Azure Blob Storage",
         discriminator="auth_type",
         type="object",
+        order=2,
+    )
+    azure_blob_storage_account_name: str = Field(
+        title="Azure Blob Storage account name",
+        description="The account's name of the Azure Blob Storage.",
+        examples=["airbyte5storage"],
         order=3,
     )
     azure_blob_storage_container_name: str = Field(

{airbyte_source_azure_blob_storage-0.4.0 → airbyte_source_azure_blob_storage-0.4.2}/source_azure_blob_storage/stream_reader.py

@@ -8,11 +8,14 @@ import pytz
 from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader, FileReadMode
 from airbyte_cdk.sources.file_based.remote_file import RemoteFile
 from airbyte_cdk.sources.streams.http.requests_native_auth import Oauth2Authenticator
+from airbyte_cdk.utils import AirbyteTracedException
+from airbyte_protocol.models import FailureType
 from azure.core.credentials import AccessToken
+from azure.core.exceptions import ResourceNotFoundError
 from azure.storage.blob import BlobServiceClient, ContainerClient
 from smart_open import open
 
-from .config import Config
+from .spec import SourceAzureBlobStorageSpec
 
 
 class AzureOauth2Authenticator(Oauth2Authenticator):
@@ -35,11 +38,11 @@ class SourceAzureBlobStorageStreamReader(AbstractFileBasedStreamReader):
         self._config = None
 
     @property
-    def config(self) -> Config:
+    def config(self) -> SourceAzureBlobStorageSpec:
         return self._config
 
     @config.setter
-    def config(self, value: Config) -> None:
+    def config(self, value: SourceAzureBlobStorageSpec) -> None:
         self._config = value
 
     @property
@@ -80,11 +83,13 @@ class SourceAzureBlobStorageStreamReader(AbstractFileBasedStreamReader):
     ) -> Iterable[RemoteFile]:
         prefixes = [prefix] if prefix else self.get_prefixes_from_globs(globs)
         prefixes = prefixes or [None]
-        for prefix in prefixes:
-            for blob in self.azure_container_client.list_blobs(name_starts_with=prefix):
-                remote_file = RemoteFile(uri=blob.name, last_modified=blob.last_modified.astimezone(pytz.utc).replace(tzinfo=None))
-                if not globs or self.file_matches_globs(remote_file, globs):
-                    yield remote_file
+        try:
+            for prefix in prefixes:
+                for blob in self.azure_container_client.list_blobs(name_starts_with=prefix):
+                    remote_file = RemoteFile(uri=blob.name, last_modified=blob.last_modified.astimezone(pytz.utc).replace(tzinfo=None))
+                    yield from self.filter_files_by_globs_and_start_date([remote_file], globs)
+        except ResourceNotFoundError as e:
+            raise AirbyteTracedException(failure_type=FailureType.config_error, internal_message=e.message, message=e.reason or e.message)
 
     def open_file(self, file: RemoteFile, mode: FileReadMode, encoding: Optional[str], logger: logging.Logger) -> IOBase:
         try:

airbyte_source_azure_blob_storage-0.4.0/source_azure_blob_storage/legacy_config_transformer.py

@@ -1,31 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-from typing import Any, Mapping, MutableMapping
-
-
-class LegacyConfigTransformer:
-    """
-    Class that takes in Azure Blob Storage source configs in the legacy format and transforms them into
-    configs that can be used by the new Azure Blob Storage source built with the file-based CDK.
-    """
-
-    @classmethod
-    def convert(cls, legacy_config: Mapping) -> MutableMapping[str, Any]:
-        azure_blob_storage_blobs_prefix = legacy_config.get("azure_blob_storage_blobs_prefix", "")
-
-        return {
-            "azure_blob_storage_endpoint": legacy_config.get("azure_blob_storage_endpoint", None),
-            "azure_blob_storage_account_name": legacy_config["azure_blob_storage_account_name"],
-            "azure_blob_storage_account_key": legacy_config["azure_blob_storage_account_key"],
-            "azure_blob_storage_container_name": legacy_config["azure_blob_storage_container_name"],
-            "streams": [
-                {
-                    "name": legacy_config["azure_blob_storage_container_name"],
-                    "legacy_prefix": azure_blob_storage_blobs_prefix,
-                    "validation_policy": "Emit Record",
-                    "format": {"filetype": "jsonl"},
-                }
-            ],
-        }