airbyte-source-azure-blob-storage 0.3.6__tar.gz → 0.4.1__tar.gz

{airbyte_source_azure_blob_storage-0.3.6 → airbyte_source_azure_blob_storage-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: airbyte-source-azure-blob-storage
-Version: 0.3.6
+Version: 0.4.1
 Summary: Source implementation for Azure Blob Storage.
 Home-page: https://airbyte.com
 License: MIT
@@ -32,6 +32,26 @@ For information about how to use this connector within Airbyte, see [the documen
 * Poetry (~=1.7) - installation instructions [here](https://python-poetry.org/docs/#installation)
 
 
+### Generate new oauth token
+
+Tenant id should be provided by user, reason:
+https://learn.microsoft.com/en-us/answers/questions/1531138/which-tenant-id-do-i-have-to-use-to-get-tokens-and
+
+1. GET https://login.microsoftonline.com/<tenant_id>/oauth2/v2.0/authorize
+    ?response_type=code
+    &client_id=<client_id>
+    &scope=offline_access https://storage.azure.com/.default
+    &redirect_uri=http://localhost:8000/auth_flow
+    &response_mode=query
+    &state=1234
+
+2. POST https://login.microsoftonline.com/<tenant_id>/oauth2/v2.0/token
+client_id:<client_id>
+code:<code obtained from previous request>
+redirect_uri:http://localhost:8000/auth_flow
+grant_type:authorization_code
+client_secret:<client_secret>
+
 ### Installing the connector
 From this connector directory, run:
 ```bash

{airbyte_source_azure_blob_storage-0.3.6 → airbyte_source_azure_blob_storage-0.4.1}/README.md

@@ -11,6 +11,26 @@ For information about how to use this connector within Airbyte, see [the documen
 * Poetry (~=1.7) - installation instructions [here](https://python-poetry.org/docs/#installation)
 
 
+### Generate new oauth token
+
+Tenant id should be provided by user, reason:
+https://learn.microsoft.com/en-us/answers/questions/1531138/which-tenant-id-do-i-have-to-use-to-get-tokens-and
+
+1. GET https://login.microsoftonline.com/<tenant_id>/oauth2/v2.0/authorize
+    ?response_type=code
+    &client_id=<client_id>
+    &scope=offline_access https://storage.azure.com/.default
+    &redirect_uri=http://localhost:8000/auth_flow
+    &response_mode=query
+    &state=1234
+
+2. POST https://login.microsoftonline.com/<tenant_id>/oauth2/v2.0/token
+client_id:<client_id>
+code:<code obtained from previous request>
+redirect_uri:http://localhost:8000/auth_flow
+grant_type:authorization_code
+client_secret:<client_secret>
+
 ### Installing the connector
 From this connector directory, run:
 ```bash

{airbyte_source_azure_blob_storage-0.3.6 → airbyte_source_azure_blob_storage-0.4.1}/pyproject.toml

@@ -5,7 +5,7 @@ requires = [
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-version = "0.3.6"
+version = "0.4.1"
 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.3.6 → airbyte_source_azure_blob_storage-0.4.1}/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.1/source_azure_blob_storage/config_migrations.py

@@ -0,0 +1,119 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, List, Mapping
+
+from airbyte_cdk.config_observation import create_connector_config_control_message
+from airbyte_cdk.entrypoint import AirbyteEntrypoint
+from airbyte_cdk.sources import Source
+
+logger = logging.getLogger("airbyte_logger")
+
+
+class MigrateConfig(ABC):
+    @classmethod
+    @abstractmethod
+    def should_migrate(cls, config: Mapping[str, Any]) -> bool:
+        ...
+
+    @classmethod
+    @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]:
+        """
+        Modifies the configuration and then saves it back to the source.
+
+        Args:
+        - config_path (str): The path where the configuration is stored.
+        - source (Source): The data source.
+        - config (Mapping[str, Any]): The current configuration.
+
+        Returns:
+        - Mapping[str, Any]: The updated configuration.
+        """
+        migrated_config = cls.migrate_config(config)
+        source.write_config(migrated_config, config_path)
+        return migrated_config
+
+    @classmethod
+    def emit_control_message(cls, migrated_config: Mapping[str, Any]) -> None:
+        """
+        Emits the control messages related to configuration migration.
+
+        Args:
+        - migrated_config (Mapping[str, Any]): The migrated configuration.
+        """
+        print(create_connector_config_control_message(migrated_config).json(exclude_unset=True))
+
+    @classmethod
+    def migrate(cls, args: List[str], source: Source) -> None:
+        """
+        Orchestrates the configuration migration process.
+
+        It first checks if the `--config` argument is provided, and if so,
+        determines whether migration is needed, and then performs the migration
+        if required.
+
+        Args:
+        - args (List[str]): List of command-line arguments.
+        - source (Source): The data source.
+        """
+        config_path = AirbyteEntrypoint(source).extract_config(args)
+        if config_path:
+            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.3.6 → airbyte_source_azure_blob_storage-0.4.1}/source_azure_blob_storage/run.py

@@ -9,7 +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 import SourceAzureBlobStorage, SourceAzureBlobStorageSpec, SourceAzureBlobStorageStreamReader
+from source_azure_blob_storage.config_migrations import MigrateCredentials, MigrateLegacyConfig
 
 
 def run():
@@ -20,12 +21,14 @@ 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(
             AirbyteMessage(

airbyte_source_azure_blob_storage-0.4.1/source_azure_blob_storage/source.py

@@ -0,0 +1,52 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Mapping
+
+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
+
+
+class SourceAzureBlobStorage(FileBasedSource):
+    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.
+        """
+
+        return ConnectorSpecification(
+            documentationUrl=self.spec_class.documentation_url(),
+            connectionSpecification=self.spec_class.schema(),
+            advanced_auth=AdvancedAuth(
+                auth_flow_type="oauth2.0",
+                predicate_key=["credentials", "auth_type"],
+                predicate_value="oauth2",
+                oauth_config_specification=OAuthConfigSpecification(
+                    complete_oauth_output_specification={
+                        "type": "object",
+                        "additionalProperties": False,
+                        "properties": {"refresh_token": {"type": "string", "path_in_connector_config": ["credentials", "refresh_token"]}},
+                    },
+                    complete_oauth_server_input_specification={
+                        "type": "object",
+                        "additionalProperties": False,
+                        "properties": {"client_id": {"type": "string"}, "client_secret": {"type": "string"}},
+                    },
+                    complete_oauth_server_output_specification={
+                        "type": "object",
+                        "additionalProperties": False,
+                        "properties": {
+                            "client_id": {"type": "string", "path_in_connector_config": ["credentials", "client_id"]},
+                            "client_secret": {"type": "string", "path_in_connector_config": ["credentials", "client_secret"]},
+                        },
+                    },
+                    oauth_user_input_from_connector_config_specification={
+                        "type": "object",
+                        "additionalProperties": False,
+                        "properties": {"tenant_id": {"type": "string", "path_in_connector_config": ["credentials", "tenant_id"]}},
+                    },
+                ),
+            ),
+        )

airbyte_source_azure_blob_storage-0.3.6/source_azure_blob_storage/config.py → airbyte_source_azure_blob_storage-0.4.1/source_azure_blob_storage/spec.py

@@ -2,14 +2,54 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Literal, Optional, Union
 
 import dpath.util
 from airbyte_cdk.sources.file_based.config.abstract_file_based_spec import AbstractFileBasedSpec
-from pydantic import AnyUrl, Field
+from airbyte_cdk.utils.oneof_option_config import OneOfOptionConfig
+from pydantic import AnyUrl, BaseModel, Field
 
 
-class Config(AbstractFileBasedSpec):
+class Oauth2(BaseModel):
+    class Config(OneOfOptionConfig):
+        title = "Authenticate via Oauth2"
+        discriminator = "auth_type"
+
+    auth_type: Literal["oauth2"] = Field("oauth2", const=True)
+    tenant_id: str = Field(title="Tenant ID", description="Tenant ID of the Microsoft Azure Application user", airbyte_secret=True)
+    client_id: str = Field(
+        title="Client ID",
+        description="Client ID of your Microsoft developer application",
+        airbyte_secret=True,
+    )
+    client_secret: str = Field(
+        title="Client Secret",
+        description="Client Secret of your Microsoft developer application",
+        airbyte_secret=True,
+    )
+    refresh_token: str = Field(
+        title="Refresh Token",
+        description="Refresh Token of your Microsoft developer application",
+        airbyte_secret=True,
+    )
+
+
+class StorageAccountKey(BaseModel):
+    class Config(OneOfOptionConfig):
+        title = "Authenticate via Storage Account Key"
+        discriminator = "auth_type"
+
+    auth_type: Literal["storage_account_key"] = Field("storage_account_key", const=True)
+    azure_blob_storage_account_key: str = Field(
+        title="Azure Blob Storage account key",
+        description="The Azure blob storage account key.",
+        airbyte_secret=True,
+        examples=["Z8ZkZpteggFx394vm+PJHnGTvdRncaYS+JhLKdj789YNmD+iyGTnG+PV+POiuYNhBg/ACS+LKjd%4FG3FHGN12Nd=="],
+        order=3,
+    )
+
+
+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.
@@ -25,11 +65,11 @@ class Config(AbstractFileBasedSpec):
         examples=["airbyte5storage"],
         order=2,
     )
-    azure_blob_storage_account_key: str = Field(
-        title="Azure Blob Storage account key",
-        description="The Azure blob storage account key.",
-        airbyte_secret=True,
-        examples=["Z8ZkZpteggFx394vm+PJHnGTvdRncaYS+JhLKdj789YNmD+iyGTnG+PV+POiuYNhBg/ACS+LKjd%4FG3FHGN12Nd=="],
+    credentials: Union[Oauth2, StorageAccountKey] = Field(
+        title="Authentication",
+        description="Credentials for connecting to the Azure Blob Storage",
+        discriminator="auth_type",
+        type="object",
         order=3,
     )
     azure_blob_storage_container_name: str = Field(

{airbyte_source_azure_blob_storage-0.3.6 → airbyte_source_azure_blob_storage-0.4.1}/source_azure_blob_storage/stream_reader.py

@@ -2,28 +2,44 @@
 
 import logging
 from io import IOBase
-from typing import Iterable, List, Optional
+from typing import Iterable, List, Optional, Union
 
 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 azure.core.credentials import AccessToken
 from azure.storage.blob import BlobServiceClient, ContainerClient
 from smart_open import open
 
-from .config import Config
+from .spec import SourceAzureBlobStorageSpec
+
+
+class AzureOauth2Authenticator(Oauth2Authenticator):
+    """
+    Authenticator for Azure Blob Storage SDK to align with azure.core.credentials.TokenCredential protocol
+    """
+
+    def get_token(self, *args, **kwargs) -> AccessToken:
+        """Parent class handles Oauth Refresh token logic.
+        `expires_on` is ignored and set to year 2222 to align with protocol.
+        """
+        return AccessToken(token=self.get_access_token(), expires_on=7952342400)
 
 
 class SourceAzureBlobStorageStreamReader(AbstractFileBasedStreamReader):
+    _credentials = None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         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
@@ -35,14 +51,26 @@ class SourceAzureBlobStorageStreamReader(AbstractFileBasedStreamReader):
     @property
     def azure_container_client(self):
         return ContainerClient(
-            self.account_url,
-            container_name=self.config.azure_blob_storage_container_name,
-            credential=self.config.azure_blob_storage_account_key,
+            self.account_url, container_name=self.config.azure_blob_storage_container_name, credential=self.azure_credentials
         )
 
     @property
     def azure_blob_service_client(self):
-        return BlobServiceClient(self.account_url, credential=self.config.azure_blob_storage_account_key)
+        return BlobServiceClient(self.account_url, credential=self._credentials)
+
+    @property
+    def azure_credentials(self) -> Union[str, AzureOauth2Authenticator]:
+        if not self._credentials:
+            if self.config.credentials.auth_type == "storage_account_key":
+                self._credentials = self.config.credentials.azure_blob_storage_account_key
+            else:
+                self._credentials = AzureOauth2Authenticator(
+                    token_refresh_endpoint=f"https://login.microsoftonline.com/{self.config.credentials.tenant_id}/oauth2/v2.0/token",
+                    client_id=self.config.credentials.client_id,
+                    client_secret=self.config.credentials.client_secret,
+                    refresh_token=self.config.credentials.refresh_token,
+                )
+        return self._credentials
 
     def get_matching_files(
         self,
@@ -55,8 +83,7 @@ class SourceAzureBlobStorageStreamReader(AbstractFileBasedStreamReader):
         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
+                yield from self.filter_files_by_globs_and_start_date([remote_file], globs)
 
     def open_file(self, file: RemoteFile, mode: FileReadMode, encoding: Optional[str], logger: logging.Logger) -> IOBase:
         try:

airbyte_source_azure_blob_storage-0.3.6/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"},
-                }
-            ],
-        }

airbyte_source_azure_blob_storage-0.3.6/source_azure_blob_storage/source.py

@@ -1,30 +0,0 @@
-#
-# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
-#
-
-from typing import Any, Mapping
-
-from airbyte_cdk.config_observation import emit_configuration_as_airbyte_control_message
-from airbyte_cdk.sources.file_based.file_based_source import FileBasedSource
-
-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