From d2777afe56eec2f85a8825147e9cf6379dae10e0 Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Fri, 9 Dec 2022 22:09:23 +0530 Subject: [PATCH 1/6] ADLS gen2 hook ADLS gen2 hook --- .../providers/microsoft/azure/hooks/adls.py | 221 ++++++++++++++++++ .../providers/microsoft/azure/provider.yaml | 12 + .../connections/adls.rst | 66 ++++++ generated/provider_dependencies.json | 1 + 4 files changed, 300 insertions(+) create mode 100644 airflow/providers/microsoft/azure/hooks/adls.py create mode 100644 docs/apache-airflow-providers-microsoft-azure/connections/adls.rst diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py new file mode 100644 index 0000000000000..d9a2111754e72 --- /dev/null +++ b/airflow/providers/microsoft/azure/hooks/adls.py @@ -0,0 +1,221 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import Any + +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError +from azure.identity import ClientSecretCredential +from azure.storage.filedatalake import ( + DataLakeDirectoryClient, + DataLakeFileClient, + DataLakeServiceClient, + FileSystemClient, +) + +from airflow.hooks.base import BaseHook + + +class AzureDataLakeStorageV2(BaseHook): + + conn_name_attr = "adls_v2_conn_id" + default_conn_name = "adls_v2_default" + conn_type = "adls_v2" + hook_name = "Azure Date Lake Storage" + + @staticmethod + def get_connection_form_widgets() -> dict[str, Any]: + """Returns connection widgets to add to connection form""" + from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget + from flask_babel import lazy_gettext + from wtforms import PasswordField, StringField + + return { + "extra__adls_v2__connection_string": PasswordField( + lazy_gettext("Blob Storage Connection String (optional)"), widget=BS3PasswordFieldWidget() + ), + "extra__adls_v2__tenant_id": StringField( + lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() + ), + } + + @staticmethod + def get_ui_field_behaviour() -> dict[str, Any]: + """Returns custom field behaviour""" + return { + "hidden_fields": ["schema", "port"], + "relabeling": { + "login": "Blob Storage Login (optional)", + "password": "Blob Storage Key (optional)", + "host": "Account Name (Active Directory Auth)", + }, + "placeholders": { + "login": "account name", + "password": "secret", + "host": "account url", + "extra__adls_v2__connection_string": "connection string auth", + "extra__adls_v2__tenant_id": "tenant", + }, + } + + def __init__(self, adls_v2_conn_id: str = default_conn_name, public_read: bool = False) -> None: + super().__init__() + self.conn_id = adls_v2_conn_id + self.public_read = public_read + self.service_client = self.get_conn() + + def get_conn(self) -> DataLakeServiceClient: + """Return the DataLakeServiceClient object.""" + conn = self.get_connection(self.conn_id) + extra = conn.extra_dejson or {} + + connection_string = extra.pop( + "connection_string", extra.pop("extra__adls_v2__connection_string", None) + ) + if connection_string: + # connection_string auth takes priority + return DataLakeServiceClient.from_connection_string(connection_string, **extra) + + tenant = extra.pop("tenant_id", extra.pop("extra__adls_v2__tenant_id", None)) + if tenant: + # use Active Directory auth + app_id = conn.login + app_secret = conn.password + token_credential = ClientSecretCredential(tenant, app_id, app_secret) + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra + ) + credential = conn.password + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", + credential=credential, + **extra, + ) + + def create_file_system(self, file_system_name: str) -> None: + """ + A container acts as a file system for your files. Creates a new file system under + the specified account. + + If the file system with the same name already exists, a ResourceExistsError will + be raised. This method returns a client with which to interact with the newly + created file system. + """ + try: + file_system_client = self.service_client.create_file_system(file_system=file_system_name) + print("file_system_client ", file_system_client.file_system_name) + self.log.info("Created file system: %s", file_system_client.file_system_name) + except ResourceExistsError: + self.log.info("Attempted to create file system %r but it already exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) + raise + + def get_file_system(self, file_system_name: str) -> FileSystemClient: + try: + file_system_client = self.service_client.get_file_system_client(file_system=file_system_name) + return file_system_client + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) + raise + + def create_directory(self, file_system_name: str, directory_name: str) -> None: + try: + file_system = self.get_file_system(file_system_name) + file_system.create_directory(directory_name) + except Exception as e: + self.log.info(e) + raise + + def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: + try: + file_system_client = self.get_file_system(file_system_name) + directory_client = file_system_client.get_directory_client(directory_name) + return directory_client + except ResourceNotFoundError: + self.log.info( + "Directory %s doesn't exists in the file system %s", directory_name, file_system_name + ) + except Exception as e: + self.log.info(e) + raise + + def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: + try: + file_system = self.get_file_system(file_system_name) + file_client = file_system.create_file(file_name) + return file_client + except Exception as e: + self.log.info(e) + raise + + def upload_file(self, + file_system_name: str, + file_name: str, + file_path: str, + overwrite: bool = False) -> None: + file_client = self.create_file(file_system_name, file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite) + + def upload_file_to_directory( + self, + file_system_name: str, + directory_name: str, + file_name: str, + file_path: str, + overwrite: bool = False, + ) -> None: + directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) + file_client = directory_client.create_file(file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite) + + def list_directory(self, file_system_name: str, directory_name: str) -> list[str]: + file_system_client = self.get_file_system(file_system_name=file_system_name) + paths = file_system_client.get_paths(directory_name) + directory_lists = [] + for path in paths: + directory_lists.append(path.name) + return directory_lists + + def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: + file_system = self.service_client.list_file_systems( + name_starts_with=prefix, include_metadata=include_metadata + ) + file_system_list = [] + for fs in file_system: + file_system_list.append(fs.name) + return file_system_list + + def delete_file_system(self, file_system_name: str) -> None: + """Delete file system""" + try: + self.service_client.delete_file_system(file_system_name) + self.log.info("Deleted file system: %s", file_system_name) + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) + raise + + def delete_directory(self, file_system_name: str, directory_name: str) -> None: + directory_client = self.get_directory_client(file_system_name, directory_name) + directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 4127791f52667..f756af7f6b2c0 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -66,6 +66,7 @@ dependencies: - azure-servicebus>=7.6.1 - azure-synapse-spark - adal>=1.2.7 + - azure-storage-file-datalake>=12.9.1 integrations: - integration-name: Microsoft Azure Batch @@ -123,6 +124,12 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-microsoft-azure/operators/azure_synapse.rst tags: [azure] + - integration-name: Microsoft Azure Data Lake Storage gen2 + how-to-guide: + - /docs/apache-airflow-providers-microsoft-azure/operators/adls.rst + external-doc-url: https://azure.microsoft.com/en-us/services/storage/data-lake-storage/ + logo: /integration-logos/azure/Data Lake Storage.svg + tags: [azure] operators: - integration-name: Microsoft Azure Data Lake Storage @@ -197,6 +204,9 @@ hooks: - integration-name: Microsoft Azure Service Bus python-modules: - airflow.providers.microsoft.azure.hooks.asb + - integration-name: Microsoft Azure Data Lake Storage Gen2 + python-modules: + - airflow.providers.microsoft.azure.hooks.adls - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.hooks.synapse @@ -249,6 +259,8 @@ connection-types: connection-type: azure_container_registry - hook-class-name: airflow.providers.microsoft.azure.hooks.asb.BaseAzureServiceBusHook connection-type: azure_service_bus + - hook-class-name: airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageV2 + connection-type: adls_v2 - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst new file mode 100644 index 0000000000000..c531e67685405 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst @@ -0,0 +1,66 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + + +.. _howto/connection:adls_v2: + +Microsoft Azure Data Lake Storage Connection Gen2 +================================================= + +The Microsoft Azure Data Lake Storage connection type enables the Azure Data Storage Integrations. + +Authenticating to Azure Data Lake Storage +----------------------------------------- + +There are four ways to connect to Azure Data Lake Storage using Airflow. + +1. Use `token credentials + `_ + i.e. add specific credentials (client_id, secret, tenant) and subscription id to the Airflow connection. +2. Use a `Connection String + `_ + i.e. add connection string to ``connection_string`` in the Airflow connection. + +Only one authorization method can be used at a time. If you need to manage multiple credentials or keys then you should +configure multiple connections. + +Default Connection IDs +---------------------- + +All hooks and operators related to Microsoft Azure Data Lake Storage use ``adls_v2_default`` by default. + +Configuring the Connection +-------------------------- + +Login (optional) + Specify the login used for azure storage. + +Password (optional) + Specify the password used for azure storage. For use with + Active Directory (token credential) and shared key authentication. + +Host (optional) + Specify the account url for anonymous public read, Active Directory, shared access key authentication. + +Extra (optional) + Specify the extra parameters (as json dictionary) that can be used in Azure connection. + The following parameters are all optional: + + * ``tenant_id``: Specify the tenant to use. Needed for Active Directory (token) authentication. + * ``connection_string``: Connection string for use with connection string authentication. + diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 4c5ba768f3b05..4985983a99b35 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -452,6 +452,7 @@ "azure-servicebus>=7.6.1", "azure-storage-blob>=12.14.0", "azure-storage-common>=2.1.0", + "azure-storage-file-datalake>=12.9.1", "azure-storage-file>=2.1.0", "azure-synapse-spark" ], From 2e9846272cf7fbc8613064ed862ba2f8ca7f5588 Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Tue, 13 Dec 2022 12:01:06 +0530 Subject: [PATCH 2/6] Add Test case --- .../providers/microsoft/azure/hooks/adls.py | 93 +++++-------------- .../providers/microsoft/azure/provider.yaml | 2 - .../microsoft/azure/hooks/test_adls.py | 77 +++++++++++++++ 3 files changed, 101 insertions(+), 71 deletions(-) create mode 100644 tests/providers/microsoft/azure/hooks/test_adls.py diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py index d9a2111754e72..ea25923c70431 100644 --- a/airflow/providers/microsoft/azure/hooks/adls.py +++ b/airflow/providers/microsoft/azure/hooks/adls.py @@ -19,6 +19,7 @@ from typing import Any +from airflow.providers.microsoft.azure.hooks.wasb import WasbHook from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError from azure.identity import ClientSecretCredential from azure.storage.filedatalake import ( @@ -28,52 +29,10 @@ FileSystemClient, ) -from airflow.hooks.base import BaseHook - - -class AzureDataLakeStorageV2(BaseHook): - - conn_name_attr = "adls_v2_conn_id" - default_conn_name = "adls_v2_default" - conn_type = "adls_v2" - hook_name = "Azure Date Lake Storage" - - @staticmethod - def get_connection_form_widgets() -> dict[str, Any]: - """Returns connection widgets to add to connection form""" - from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget - from flask_babel import lazy_gettext - from wtforms import PasswordField, StringField - - return { - "extra__adls_v2__connection_string": PasswordField( - lazy_gettext("Blob Storage Connection String (optional)"), widget=BS3PasswordFieldWidget() - ), - "extra__adls_v2__tenant_id": StringField( - lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() - ), - } - - @staticmethod - def get_ui_field_behaviour() -> dict[str, Any]: - """Returns custom field behaviour""" - return { - "hidden_fields": ["schema", "port"], - "relabeling": { - "login": "Blob Storage Login (optional)", - "password": "Blob Storage Key (optional)", - "host": "Account Name (Active Directory Auth)", - }, - "placeholders": { - "login": "account name", - "password": "secret", - "host": "account url", - "extra__adls_v2__connection_string": "connection string auth", - "extra__adls_v2__tenant_id": "tenant", - }, - } - - def __init__(self, adls_v2_conn_id: str = default_conn_name, public_read: bool = False) -> None: + +class AzureDataLakeStorageGen2(WasbHook): + + def __init__(self, adls_v2_conn_id: str, public_read: bool = False) -> None: super().__init__() self.conn_id = adls_v2_conn_id self.public_read = public_read @@ -118,7 +77,6 @@ def create_file_system(self, file_system_name: str) -> None: """ try: file_system_client = self.service_client.create_file_system(file_system=file_system_name) - print("file_system_client ", file_system_client.file_system_name) self.log.info("Created file system: %s", file_system_client.file_system_name) except ResourceExistsError: self.log.info("Attempted to create file system %r but it already exists.", file_system_name) @@ -136,18 +94,13 @@ def get_file_system(self, file_system_name: str) -> FileSystemClient: self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) raise - def create_directory(self, file_system_name: str, directory_name: str) -> None: - try: - file_system = self.get_file_system(file_system_name) - file_system.create_directory(directory_name) - except Exception as e: - self.log.info(e) - raise + def create_directory(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: + result = self.get_file_system(file_system_name).create_directory(directory_name) + return result def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: try: - file_system_client = self.get_file_system(file_system_name) - directory_client = file_system_client.get_directory_client(directory_name) + directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) return directory_client except ResourceNotFoundError: self.log.info( @@ -158,13 +111,8 @@ def get_directory_client(self, file_system_name: str, directory_name: str) -> Da raise def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: - try: - file_system = self.get_file_system(file_system_name) - file_client = file_system.create_file(file_name) - return file_client - except Exception as e: - self.log.info(e) - raise + file_client = self.get_file_system(file_system_name).create_file(file_name) + return file_client def upload_file(self, file_system_name: str, @@ -182,21 +130,27 @@ def upload_file_to_directory( file_name: str, file_path: str, overwrite: bool = False, + **kwargs: any, ) -> None: + """ + Create a new file and return the file client to be interacted with and then + upload data to a file + """ directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) - file_client = directory_client.create_file(file_name) + file_client = directory_client.create_file(file_name, kwargs=kwargs) with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite) + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) - def list_directory(self, file_system_name: str, directory_name: str) -> list[str]: - file_system_client = self.get_file_system(file_system_name=file_system_name) - paths = file_system_client.get_paths(directory_name) + def list_files_directory(self, file_system_name: str, directory_name: str) -> list[str]: + """Get the list of files or directories under the specified file system""" + paths = self.get_file_system(file_system_name=file_system_name).get_paths(directory_name) directory_lists = [] for path in paths: directory_lists.append(path.name) return directory_lists def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: + """Get the list the file systems under the specified account.""" file_system = self.service_client.list_file_systems( name_starts_with=prefix, include_metadata=include_metadata ) @@ -206,7 +160,7 @@ def list_file_system(self, prefix: str | None = None, include_metadata: bool = F return file_system_list def delete_file_system(self, file_system_name: str) -> None: - """Delete file system""" + """Deletes the file system""" try: self.service_client.delete_file_system(file_system_name) self.log.info("Deleted file system: %s", file_system_name) @@ -217,5 +171,6 @@ def delete_file_system(self, file_system_name: str) -> None: raise def delete_directory(self, file_system_name: str, directory_name: str) -> None: + """Deletes specified directory in file system""" directory_client = self.get_directory_client(file_system_name, directory_name) directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index f756af7f6b2c0..d20b4e071c3e3 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -259,8 +259,6 @@ connection-types: connection-type: azure_container_registry - hook-class-name: airflow.providers.microsoft.azure.hooks.asb.BaseAzureServiceBusHook connection-type: azure_service_bus - - hook-class-name: airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageV2 - connection-type: adls_v2 - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse diff --git a/tests/providers/microsoft/azure/hooks/test_adls.py b/tests/providers/microsoft/azure/hooks/test_adls.py new file mode 100644 index 0000000000000..38faa2cd9889a --- /dev/null +++ b/tests/providers/microsoft/azure/hooks/test_adls.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import json +from unittest import mock + +import pytest +from airflow.models import Connection +from airflow.providers.microsoft.azure.hooks.adls import AzureDataLakeStorageGen2 + + +class TestAzureDataLakeStorageGen2: + def setup_class(self) -> None: + self.conn_id: str = "azure_data_lake_storage_v2" + self.file_system_name = "test_file_system" + self.directory_name = "test_directory" + self.file_name = "test_file_name" + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_file_system(self, mock_conn): + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.create_file_system("test_file_system") + expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.FileSystemClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_get_file_system(self, mock_conn, mock_file_system): + mock_conn.return_value.get_file_system_client.return_value = mock_file_system + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.get_file_system(self.file_system_name) + assert result == mock_file_system + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.create_directory.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.create_directory(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.get_directory_client(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeFileClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_file(self,mock_conn, mock_get_file_system, mock_file_client): + mock_get_file_system.return_value.create_file.return_value = mock_file_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.create_file(self.file_system_name, self.file_name) + assert result == mock_file_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_delete_file_system(self, mock_conn): + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.delete_file_system(self.file_system_name) + expected_calls = [mock.call().delete_file_system(self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_delete_directory(self, mock_conn, mock_directory_client): + mock_conn.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.delete_directory(self.file_system_name, self.directory_name) + expected_calls = [mock.call().get_file_system_client( + self.file_system_name).get_directory_client(self.directory_name).delete_directory()] + mock_conn.assert_has_calls(expected_calls) + + From fc15cc6d015ba9749198668451d4f7ea12de954f Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Tue, 13 Dec 2022 15:23:53 +0530 Subject: [PATCH 3/6] Moved hook to existing data lake hook file Fix mypy issues Fix mypy issues Fix mypy issues Fix mypy issues Add adls v2 connection Fix doc Fix doc Moved hook to existing data lake hook file Moved hook test case Moved hook test case Fix doc --- .../providers/microsoft/azure/hooks/adls.py | 176 ---------- .../microsoft/azure/hooks/data_lake.py | 311 +++++++++++++++++- .../providers/microsoft/azure/provider.yaml | 12 +- .../connections/{adls.rst => adls_v2.rst} | 26 +- .../microsoft/azure/hooks/test_adls.py | 77 ----- .../azure/hooks/test_azure_data_lake.py | 96 +++++- 6 files changed, 418 insertions(+), 280 deletions(-) delete mode 100644 airflow/providers/microsoft/azure/hooks/adls.py rename docs/apache-airflow-providers-microsoft-azure/connections/{adls.rst => adls_v2.rst} (71%) delete mode 100644 tests/providers/microsoft/azure/hooks/test_adls.py diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py deleted file mode 100644 index ea25923c70431..0000000000000 --- a/airflow/providers/microsoft/azure/hooks/adls.py +++ /dev/null @@ -1,176 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -from __future__ import annotations - -from typing import Any - -from airflow.providers.microsoft.azure.hooks.wasb import WasbHook -from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError -from azure.identity import ClientSecretCredential -from azure.storage.filedatalake import ( - DataLakeDirectoryClient, - DataLakeFileClient, - DataLakeServiceClient, - FileSystemClient, -) - - -class AzureDataLakeStorageGen2(WasbHook): - - def __init__(self, adls_v2_conn_id: str, public_read: bool = False) -> None: - super().__init__() - self.conn_id = adls_v2_conn_id - self.public_read = public_read - self.service_client = self.get_conn() - - def get_conn(self) -> DataLakeServiceClient: - """Return the DataLakeServiceClient object.""" - conn = self.get_connection(self.conn_id) - extra = conn.extra_dejson or {} - - connection_string = extra.pop( - "connection_string", extra.pop("extra__adls_v2__connection_string", None) - ) - if connection_string: - # connection_string auth takes priority - return DataLakeServiceClient.from_connection_string(connection_string, **extra) - - tenant = extra.pop("tenant_id", extra.pop("extra__adls_v2__tenant_id", None)) - if tenant: - # use Active Directory auth - app_id = conn.login - app_secret = conn.password - token_credential = ClientSecretCredential(tenant, app_id, app_secret) - return DataLakeServiceClient( - account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra - ) - credential = conn.password - return DataLakeServiceClient( - account_url=f"https://{conn.login}.dfs.core.windows.net", - credential=credential, - **extra, - ) - - def create_file_system(self, file_system_name: str) -> None: - """ - A container acts as a file system for your files. Creates a new file system under - the specified account. - - If the file system with the same name already exists, a ResourceExistsError will - be raised. This method returns a client with which to interact with the newly - created file system. - """ - try: - file_system_client = self.service_client.create_file_system(file_system=file_system_name) - self.log.info("Created file system: %s", file_system_client.file_system_name) - except ResourceExistsError: - self.log.info("Attempted to create file system %r but it already exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) - raise - - def get_file_system(self, file_system_name: str) -> FileSystemClient: - try: - file_system_client = self.service_client.get_file_system_client(file_system=file_system_name) - return file_system_client - except ResourceNotFoundError: - self.log.info("file system %r doesn't exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) - raise - - def create_directory(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: - result = self.get_file_system(file_system_name).create_directory(directory_name) - return result - - def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: - try: - directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) - return directory_client - except ResourceNotFoundError: - self.log.info( - "Directory %s doesn't exists in the file system %s", directory_name, file_system_name - ) - except Exception as e: - self.log.info(e) - raise - - def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: - file_client = self.get_file_system(file_system_name).create_file(file_name) - return file_client - - def upload_file(self, - file_system_name: str, - file_name: str, - file_path: str, - overwrite: bool = False) -> None: - file_client = self.create_file(file_system_name, file_name) - with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite) - - def upload_file_to_directory( - self, - file_system_name: str, - directory_name: str, - file_name: str, - file_path: str, - overwrite: bool = False, - **kwargs: any, - ) -> None: - """ - Create a new file and return the file client to be interacted with and then - upload data to a file - """ - directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) - file_client = directory_client.create_file(file_name, kwargs=kwargs) - with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) - - def list_files_directory(self, file_system_name: str, directory_name: str) -> list[str]: - """Get the list of files or directories under the specified file system""" - paths = self.get_file_system(file_system_name=file_system_name).get_paths(directory_name) - directory_lists = [] - for path in paths: - directory_lists.append(path.name) - return directory_lists - - def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: - """Get the list the file systems under the specified account.""" - file_system = self.service_client.list_file_systems( - name_starts_with=prefix, include_metadata=include_metadata - ) - file_system_list = [] - for fs in file_system: - file_system_list.append(fs.name) - return file_system_list - - def delete_file_system(self, file_system_name: str) -> None: - """Deletes the file system""" - try: - self.service_client.delete_file_system(file_system_name) - self.log.info("Deleted file system: %s", file_system_name) - except ResourceNotFoundError: - self.log.info("file system %r doesn't exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) - raise - - def delete_directory(self, file_system_name: str, directory_name: str) -> None: - """Deletes specified directory in file system""" - directory_client = self.get_directory_client(file_system_name, directory_name) - directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/hooks/data_lake.py b/airflow/providers/microsoft/azure/hooks/data_lake.py index 4a9bc98f925fb..ad5e9818a958a 100644 --- a/airflow/providers/microsoft/azure/hooks/data_lake.py +++ b/airflow/providers/microsoft/azure/hooks/data_lake.py @@ -15,19 +15,21 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -""" -This module contains integration with Azure Data Lake. - -AzureDataLakeHook communicates via a REST API compatible with WebHDFS. Make sure that a -Airflow connection of type `azure_data_lake` exists. Authorization can be done by supplying a -login (=Client ID), password (=Client Secret) and extra fields tenant (Tenant) and account_name (Account Name) -(see connection `azure_data_lake_default` for an example). -""" from __future__ import annotations from typing import Any +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError from azure.datalake.store import core, lib, multithread +from azure.identity import ClientSecretCredential +from azure.storage.filedatalake import ( + DataLakeDirectoryClient, + DataLakeFileClient, + DataLakeServiceClient, + DirectoryProperties, + FileSystemClient, + FileSystemProperties, +) from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook @@ -36,6 +38,13 @@ class AzureDataLakeHook(BaseHook): """ + This module contains integration with Azure Data Lake. + + AzureDataLakeHook communicates via a REST API compatible with WebHDFS. Make sure that a + Airflow connection of type `azure_data_lake` exists. Authorization can be done by supplying a + login (=Client ID), password (=Client Secret) and extra fields tenant (Tenant) and account_name + (Account Name)(see connection `azure_data_lake_default` for an example). + Interacts with Azure Data Lake. Client ID and client secret should be in user and password parameters. @@ -230,3 +239,289 @@ def remove(self, path: str, recursive: bool = False, ignore_not_found: bool = Tr self.log.info("File %s not found", path) else: raise AirflowException(f"File {path} not found") + + +class AzureDataLakeStorageV2Hook(BaseHook): + """ + This Hook interacts with ADLS gen2 storage account it mainly helps to create and manage + directories and files in storage accounts that have a hierarchical namespace. Using Adls_v2 connection + details create DataLakeServiceClient object + + Due to Wasb is marked as legacy and and retirement of the (ADLS1) it would be nice to + implement ADLS gen2 hook for interacting with the storage account. + + .. seealso:: + https://learn.microsoft.com/en-us/azure/storage/blobs/data-lake-storage-directory-file-acl-python + + :param adls_conn_id: Reference to the :ref:`adls connection `. + :param public_read: Whether an anonymous public read access should be used. default is False + """ + + conn_name_attr = "adls_conn_id" + default_conn_name = "adls_default" + conn_type = "adls" + hook_name = "Azure Date Lake Storage V2" + + @staticmethod + def get_connection_form_widgets() -> dict[str, Any]: + """Returns connection widgets to add to connection form""" + from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget + from flask_babel import lazy_gettext + from wtforms import PasswordField, StringField + + return { + "connection_string": PasswordField( + lazy_gettext("ADLS gen2 Connection String (optional)"), widget=BS3PasswordFieldWidget() + ), + "tenant_id": StringField( + lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() + ), + } + + @staticmethod + def get_ui_field_behaviour() -> dict[str, Any]: + """Returns custom field behaviour""" + return { + "hidden_fields": ["schema", "port"], + "relabeling": { + "login": "ADLS gen2 Storage Login (optional)", + "password": "ADLS gen2 Storage Key (optional)", + "host": "Account Name (Active Directory Auth)", + }, + "placeholders": { + "extra": "additional options for use with FileService and AzureFileVolume", + "login": "account name", + "password": "secret", + "host": "account url", + "connection_string": "connection string auth", + "tenant_id": "tenant", + }, + } + + def __init__(self, adls_conn_id: str, public_read: bool = False) -> None: + super().__init__() + self.conn_id = adls_conn_id + self.public_read = public_read + self.service_client = self.get_conn() + + def get_conn(self) -> DataLakeServiceClient: # type: ignore[override] + """Return the DataLakeServiceClient object.""" + conn = self.get_connection(self.conn_id) + extra = conn.extra_dejson or {} + + connection_string = self._get_field(extra, "connection_string") + if connection_string: + # connection_string auth takes priority + return DataLakeServiceClient.from_connection_string(connection_string, **extra) + + tenant = self._get_field(extra, "tenant_id") + if tenant: + # use Active Directory auth + app_id = conn.login + app_secret = conn.password + token_credential = ClientSecretCredential(tenant, app_id, app_secret) + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra + ) + credential = conn.password + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", + credential=credential, + **extra, + ) + + def _get_field(self, extra_dict, field_name): + prefix = "extra__adls__" + if field_name.startswith("extra__"): + raise ValueError( + f"Got prefixed name {field_name}; please remove the '{prefix}' prefix " + f"when using this method." + ) + if field_name in extra_dict: + return extra_dict[field_name] or None + return extra_dict.get(f"{prefix}{field_name}") or None + + def create_file_system(self, file_system_name: str) -> None: + """ + A container acts as a file system for your files. Creates a new file system under + the specified account. + + If the file system with the same name already exists, a ResourceExistsError will + be raised. This method returns a client with which to interact with the newly + created file system. + """ + try: + file_system_client = self.service_client.create_file_system(file_system=file_system_name) + self.log.info("Created file system: %s", file_system_client.file_system_name) + except ResourceExistsError: + self.log.info("Attempted to create file system %r but it already exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) + raise + + def get_file_system(self, file_system: FileSystemProperties | str) -> FileSystemClient: + """ + Get a client to interact with the specified file system + + :param file_system: This can either be the name of the file system + or an instance of FileSystemProperties. + """ + try: + file_system_client = self.service_client.get_file_system_client(file_system=file_system) + return file_system_client + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system) + raise + except Exception as e: + self.log.info("Error while attempting to get file system %r: %s", file_system, e) + raise + + def create_directory( + self, file_system_name: FileSystemProperties | str, directory_name: str, **kwargs + ) -> DataLakeDirectoryClient: + """ + Create a directory under the specified file system. + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory which needs to be created in the file system. + """ + result = self.get_file_system(file_system_name).create_directory(directory_name, kwargs) + return result + + def get_directory_client( + self, + file_system_name: FileSystemProperties | str, + directory_name: DirectoryProperties | str, + ) -> DataLakeDirectoryClient: + """ + Get the specific directory under the specified file system. + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory or instance of DirectoryProperties which needs to be + retrieved from the file system. + """ + try: + directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) + return directory_client + except ResourceNotFoundError: + self.log.info( + "Directory %s doesn't exists in the file system %s", directory_name, file_system_name + ) + raise + except Exception as e: + self.log.info(e) + raise + + def create_file(self, file_system_name: FileSystemProperties | str, file_name: str) -> DataLakeFileClient: + """ + Creates a file under the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param file_name: Name of the file which needs to be created in the file system. + """ + file_client = self.get_file_system(file_system_name).create_file(file_name) + return file_client + + def upload_file( + self, + file_system_name: FileSystemProperties | str, + file_name: str, + file_path: str, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """ + Create a file with data in the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param file_name: Name of the file to be created with name. + :param file_path: Path to the file to load. + :param overwrite: Boolean flag to overwrite an existing file or not. + """ + file_client = self.create_file(file_system_name, file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) + + def upload_file_to_directory( + self, + file_system_name: str, + directory_name: str, + file_name: str, + file_path: str, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """ + Create a new file and return the file client to be interacted with and then + upload data to a file + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + :param file_name: Name of the file to be created with name. + :param file_path: Path to the file to load. + :param overwrite: Boolean flag to overwrite an existing file or not. + """ + directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) + file_client = directory_client.create_file(file_name, kwargs=kwargs) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) + + def list_files_directory( + self, file_system_name: FileSystemProperties | str, directory_name: str + ) -> list[str]: + """ + Get the list of files or directories under the specified file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + """ + paths = self.get_file_system(file_system=file_system_name).get_paths(directory_name) + directory_lists = [] + for path in paths: + directory_lists.append(path.name) + return directory_lists + + def list_file_system( + self, prefix: str | None = None, include_metadata: bool = False, **kwargs: Any + ) -> list[str]: + """ + Get the list the file systems under the specified account. + + :param prefix: + Filters the results to return only file systems whose names + begin with the specified prefix. + :param include_metadata: Specifies that file system metadata be returned in the response. + The default value is `False`. + """ + file_system = self.service_client.list_file_systems( + name_starts_with=prefix, include_metadata=include_metadata + ) + file_system_list = [] + for fs in file_system: + file_system_list.append(fs.name) + return file_system_list + + def delete_file_system(self, file_system_name: FileSystemProperties | str) -> None: + """ + Deletes the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + """ + try: + self.service_client.delete_file_system(file_system_name) + self.log.info("Deleted file system: %s", file_system_name) + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) + raise + + def delete_directory(self, file_system_name: FileSystemProperties | str, directory_name: str) -> None: + """ + Deletes specified directory in file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + """ + directory_client = self.get_directory_client(file_system_name, directory_name) + directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index d20b4e071c3e3..41f890fd6f9e0 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -124,10 +124,8 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-microsoft-azure/operators/azure_synapse.rst tags: [azure] - - integration-name: Microsoft Azure Data Lake Storage gen2 - how-to-guide: - - /docs/apache-airflow-providers-microsoft-azure/operators/adls.rst - external-doc-url: https://azure.microsoft.com/en-us/services/storage/data-lake-storage/ + - integration-name: Microsoft Azure Data Lake Storage Client Gen2 + external-doc-url: https://azure.microsoft.com/en-us/products/storage/data-lake-storage/ logo: /integration-logos/azure/Data Lake Storage.svg tags: [azure] @@ -204,9 +202,9 @@ hooks: - integration-name: Microsoft Azure Service Bus python-modules: - airflow.providers.microsoft.azure.hooks.asb - - integration-name: Microsoft Azure Data Lake Storage Gen2 + - integration-name: Microsoft Azure Data Lake Storage Client Gen2 python-modules: - - airflow.providers.microsoft.azure.hooks.adls + - airflow.providers.microsoft.azure.hooks.data_lake - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.hooks.synapse @@ -261,6 +259,8 @@ connection-types: connection-type: azure_service_bus - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse + - hook-class-name: airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook + connection-type: adls secrets-backends: - airflow.providers.microsoft.azure.secrets.key_vault.AzureKeyVaultBackend diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst similarity index 71% rename from docs/apache-airflow-providers-microsoft-azure/connections/adls.rst rename to docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst index c531e67685405..9ec4015679def 100644 --- a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst @@ -15,19 +15,17 @@ specific language governing permissions and limitations under the License. +.. _howto/connection:adls: +Microsoft Azure Data Lake Storage Gen2 Connection +================================================== -.. _howto/connection:adls_v2: +The Microsoft Azure Data Lake Storage Gen2 connection type enables the ADLS gen2 Integrations. -Microsoft Azure Data Lake Storage Connection Gen2 -================================================= +Authenticating to Azure Data Lake Storage Gen2 +---------------------------------------------- -The Microsoft Azure Data Lake Storage connection type enables the Azure Data Storage Integrations. - -Authenticating to Azure Data Lake Storage ------------------------------------------ - -There are four ways to connect to Azure Data Lake Storage using Airflow. +Currently, there are two ways to connect to Azure Data Lake Storage Gen2 using Airflow. 1. Use `token credentials `_ @@ -42,16 +40,16 @@ configure multiple connections. Default Connection IDs ---------------------- -All hooks and operators related to Microsoft Azure Data Lake Storage use ``adls_v2_default`` by default. +All hooks and operators related to Microsoft Azure Blob Storage use ``azure_data_lake_default`` by default. Configuring the Connection -------------------------- Login (optional) - Specify the login used for azure storage. + Specify the login used for azure blob storage. For use with Shared Key Credential and SAS Token authentication. Password (optional) - Specify the password used for azure storage. For use with + Specify the password used for azure blob storage. For use with Active Directory (token credential) and shared key authentication. Host (optional) @@ -64,3 +62,7 @@ Extra (optional) * ``tenant_id``: Specify the tenant to use. Needed for Active Directory (token) authentication. * ``connection_string``: Connection string for use with connection string authentication. +When specifying the connection in environment variable you should specify +it using URI syntax. + +Note that all components of the URI should be URL-encoded. diff --git a/tests/providers/microsoft/azure/hooks/test_adls.py b/tests/providers/microsoft/azure/hooks/test_adls.py deleted file mode 100644 index 38faa2cd9889a..0000000000000 --- a/tests/providers/microsoft/azure/hooks/test_adls.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -import json -from unittest import mock - -import pytest -from airflow.models import Connection -from airflow.providers.microsoft.azure.hooks.adls import AzureDataLakeStorageGen2 - - -class TestAzureDataLakeStorageGen2: - def setup_class(self) -> None: - self.conn_id: str = "azure_data_lake_storage_v2" - self.file_system_name = "test_file_system" - self.directory_name = "test_directory" - self.file_name = "test_file_name" - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_file_system(self, mock_conn): - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.create_file_system("test_file_system") - expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] - mock_conn.assert_has_calls(expected_calls) - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.FileSystemClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_get_file_system(self, mock_conn, mock_file_system): - mock_conn.return_value.get_file_system_client.return_value = mock_file_system - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.get_file_system(self.file_system_name) - assert result == mock_file_system - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): - mock_get_file_system.return_value.create_directory.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.create_directory(self.file_system_name, self.directory_name) - assert result == mock_directory_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): - mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.get_directory_client(self.file_system_name, self.directory_name) - assert result == mock_directory_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeFileClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_file(self,mock_conn, mock_get_file_system, mock_file_client): - mock_get_file_system.return_value.create_file.return_value = mock_file_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.create_file(self.file_system_name, self.file_name) - assert result == mock_file_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_delete_file_system(self, mock_conn): - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.delete_file_system(self.file_system_name) - expected_calls = [mock.call().delete_file_system(self.file_system_name)] - mock_conn.assert_has_calls(expected_calls) - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_delete_directory(self, mock_conn, mock_directory_client): - mock_conn.return_value.get_directory_client.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.delete_directory(self.file_system_name, self.directory_name) - expected_calls = [mock.call().get_file_system_client( - self.file_system_name).get_directory_client(self.directory_name).delete_directory()] - mock_conn.assert_has_calls(expected_calls) - - diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py index f2b2d5a65439f..e8d378ab5c12a 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py @@ -21,7 +21,7 @@ from unittest import mock from airflow.models import Connection -from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook +from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook, AzureDataLakeStorageV2Hook from airflow.utils import db from tests.test_utils.providers import get_provider_min_airflow_version @@ -151,3 +151,97 @@ def test_get_ui_field_behaviour_placeholders(self): "You must now remove `_ensure_prefixes` from azure utils." " The functionality is now taken care of by providers manager." ) + + +class TestAzureDataLakeStorageV2Hook: + def setup_class(self) -> None: + self.conn_id: str = "adls_conn_id" + self.file_system_name = "test_file_system" + self.directory_name = "test_directory" + self.file_name = "test_file_name" + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.create_file_system("test_file_system") + expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.FileSystemClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_get_file_system(self, mock_conn, mock_file_system): + mock_conn.return_value.get_file_system_client.return_value = mock_file_system + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.get_file_system(self.file_system_name) + assert result == mock_file_system + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.create_directory.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.create_directory(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.get_directory_client(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeFileClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_file(self, mock_conn, mock_get_file_system, mock_file_client): + mock_get_file_system.return_value.create_file.return_value = mock_file_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.create_file(self.file_system_name, self.file_name) + assert result == mock_file_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_delete_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.delete_file_system(self.file_system_name) + expected_calls = [mock.call().delete_file_system(self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_delete_directory(self, mock_conn, mock_directory_client): + mock_conn.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.delete_directory(self.file_system_name, self.directory_name) + expected_calls = [ + mock.call() + .get_file_system_client(self.file_system_name) + .get_directory_client(self.directory_name) + .delete_directory() + ] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_list_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.list_file_system(prefix="prefix") + mock_conn.return_value.list_file_systems.assert_called_once_with( + name_starts_with="prefix", include_metadata=False + ) + + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_list_files_directory(self, mock_conn, mock_get_file_system): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.list_files_directory(self.file_system_name, self.directory_name) + mock_get_file_system.return_value.get_paths.assert_called_once_with(self.directory_name) From 5725a227ea28ae72de9de268497b8e0418e6f5be Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Fri, 9 Dec 2022 22:09:23 +0530 Subject: [PATCH 4/6] ADLS gen2 hook ADLS gen2 hook --- .../providers/microsoft/azure/hooks/adls.py | 221 ++++++++++++++++++ .../providers/microsoft/azure/provider.yaml | 12 + .../connections/adls.rst | 66 ++++++ generated/provider_dependencies.json | 1 + 4 files changed, 300 insertions(+) create mode 100644 airflow/providers/microsoft/azure/hooks/adls.py create mode 100644 docs/apache-airflow-providers-microsoft-azure/connections/adls.rst diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py new file mode 100644 index 0000000000000..d9a2111754e72 --- /dev/null +++ b/airflow/providers/microsoft/azure/hooks/adls.py @@ -0,0 +1,221 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import Any + +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError +from azure.identity import ClientSecretCredential +from azure.storage.filedatalake import ( + DataLakeDirectoryClient, + DataLakeFileClient, + DataLakeServiceClient, + FileSystemClient, +) + +from airflow.hooks.base import BaseHook + + +class AzureDataLakeStorageV2(BaseHook): + + conn_name_attr = "adls_v2_conn_id" + default_conn_name = "adls_v2_default" + conn_type = "adls_v2" + hook_name = "Azure Date Lake Storage" + + @staticmethod + def get_connection_form_widgets() -> dict[str, Any]: + """Returns connection widgets to add to connection form""" + from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget + from flask_babel import lazy_gettext + from wtforms import PasswordField, StringField + + return { + "extra__adls_v2__connection_string": PasswordField( + lazy_gettext("Blob Storage Connection String (optional)"), widget=BS3PasswordFieldWidget() + ), + "extra__adls_v2__tenant_id": StringField( + lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() + ), + } + + @staticmethod + def get_ui_field_behaviour() -> dict[str, Any]: + """Returns custom field behaviour""" + return { + "hidden_fields": ["schema", "port"], + "relabeling": { + "login": "Blob Storage Login (optional)", + "password": "Blob Storage Key (optional)", + "host": "Account Name (Active Directory Auth)", + }, + "placeholders": { + "login": "account name", + "password": "secret", + "host": "account url", + "extra__adls_v2__connection_string": "connection string auth", + "extra__adls_v2__tenant_id": "tenant", + }, + } + + def __init__(self, adls_v2_conn_id: str = default_conn_name, public_read: bool = False) -> None: + super().__init__() + self.conn_id = adls_v2_conn_id + self.public_read = public_read + self.service_client = self.get_conn() + + def get_conn(self) -> DataLakeServiceClient: + """Return the DataLakeServiceClient object.""" + conn = self.get_connection(self.conn_id) + extra = conn.extra_dejson or {} + + connection_string = extra.pop( + "connection_string", extra.pop("extra__adls_v2__connection_string", None) + ) + if connection_string: + # connection_string auth takes priority + return DataLakeServiceClient.from_connection_string(connection_string, **extra) + + tenant = extra.pop("tenant_id", extra.pop("extra__adls_v2__tenant_id", None)) + if tenant: + # use Active Directory auth + app_id = conn.login + app_secret = conn.password + token_credential = ClientSecretCredential(tenant, app_id, app_secret) + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra + ) + credential = conn.password + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", + credential=credential, + **extra, + ) + + def create_file_system(self, file_system_name: str) -> None: + """ + A container acts as a file system for your files. Creates a new file system under + the specified account. + + If the file system with the same name already exists, a ResourceExistsError will + be raised. This method returns a client with which to interact with the newly + created file system. + """ + try: + file_system_client = self.service_client.create_file_system(file_system=file_system_name) + print("file_system_client ", file_system_client.file_system_name) + self.log.info("Created file system: %s", file_system_client.file_system_name) + except ResourceExistsError: + self.log.info("Attempted to create file system %r but it already exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) + raise + + def get_file_system(self, file_system_name: str) -> FileSystemClient: + try: + file_system_client = self.service_client.get_file_system_client(file_system=file_system_name) + return file_system_client + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) + raise + + def create_directory(self, file_system_name: str, directory_name: str) -> None: + try: + file_system = self.get_file_system(file_system_name) + file_system.create_directory(directory_name) + except Exception as e: + self.log.info(e) + raise + + def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: + try: + file_system_client = self.get_file_system(file_system_name) + directory_client = file_system_client.get_directory_client(directory_name) + return directory_client + except ResourceNotFoundError: + self.log.info( + "Directory %s doesn't exists in the file system %s", directory_name, file_system_name + ) + except Exception as e: + self.log.info(e) + raise + + def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: + try: + file_system = self.get_file_system(file_system_name) + file_client = file_system.create_file(file_name) + return file_client + except Exception as e: + self.log.info(e) + raise + + def upload_file(self, + file_system_name: str, + file_name: str, + file_path: str, + overwrite: bool = False) -> None: + file_client = self.create_file(file_system_name, file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite) + + def upload_file_to_directory( + self, + file_system_name: str, + directory_name: str, + file_name: str, + file_path: str, + overwrite: bool = False, + ) -> None: + directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) + file_client = directory_client.create_file(file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite) + + def list_directory(self, file_system_name: str, directory_name: str) -> list[str]: + file_system_client = self.get_file_system(file_system_name=file_system_name) + paths = file_system_client.get_paths(directory_name) + directory_lists = [] + for path in paths: + directory_lists.append(path.name) + return directory_lists + + def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: + file_system = self.service_client.list_file_systems( + name_starts_with=prefix, include_metadata=include_metadata + ) + file_system_list = [] + for fs in file_system: + file_system_list.append(fs.name) + return file_system_list + + def delete_file_system(self, file_system_name: str) -> None: + """Delete file system""" + try: + self.service_client.delete_file_system(file_system_name) + self.log.info("Deleted file system: %s", file_system_name) + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) + raise + + def delete_directory(self, file_system_name: str, directory_name: str) -> None: + directory_client = self.get_directory_client(file_system_name, directory_name) + directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index 4127791f52667..f756af7f6b2c0 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -66,6 +66,7 @@ dependencies: - azure-servicebus>=7.6.1 - azure-synapse-spark - adal>=1.2.7 + - azure-storage-file-datalake>=12.9.1 integrations: - integration-name: Microsoft Azure Batch @@ -123,6 +124,12 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-microsoft-azure/operators/azure_synapse.rst tags: [azure] + - integration-name: Microsoft Azure Data Lake Storage gen2 + how-to-guide: + - /docs/apache-airflow-providers-microsoft-azure/operators/adls.rst + external-doc-url: https://azure.microsoft.com/en-us/services/storage/data-lake-storage/ + logo: /integration-logos/azure/Data Lake Storage.svg + tags: [azure] operators: - integration-name: Microsoft Azure Data Lake Storage @@ -197,6 +204,9 @@ hooks: - integration-name: Microsoft Azure Service Bus python-modules: - airflow.providers.microsoft.azure.hooks.asb + - integration-name: Microsoft Azure Data Lake Storage Gen2 + python-modules: + - airflow.providers.microsoft.azure.hooks.adls - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.hooks.synapse @@ -249,6 +259,8 @@ connection-types: connection-type: azure_container_registry - hook-class-name: airflow.providers.microsoft.azure.hooks.asb.BaseAzureServiceBusHook connection-type: azure_service_bus + - hook-class-name: airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageV2 + connection-type: adls_v2 - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst new file mode 100644 index 0000000000000..c531e67685405 --- /dev/null +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst @@ -0,0 +1,66 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + + + +.. _howto/connection:adls_v2: + +Microsoft Azure Data Lake Storage Connection Gen2 +================================================= + +The Microsoft Azure Data Lake Storage connection type enables the Azure Data Storage Integrations. + +Authenticating to Azure Data Lake Storage +----------------------------------------- + +There are four ways to connect to Azure Data Lake Storage using Airflow. + +1. Use `token credentials + `_ + i.e. add specific credentials (client_id, secret, tenant) and subscription id to the Airflow connection. +2. Use a `Connection String + `_ + i.e. add connection string to ``connection_string`` in the Airflow connection. + +Only one authorization method can be used at a time. If you need to manage multiple credentials or keys then you should +configure multiple connections. + +Default Connection IDs +---------------------- + +All hooks and operators related to Microsoft Azure Data Lake Storage use ``adls_v2_default`` by default. + +Configuring the Connection +-------------------------- + +Login (optional) + Specify the login used for azure storage. + +Password (optional) + Specify the password used for azure storage. For use with + Active Directory (token credential) and shared key authentication. + +Host (optional) + Specify the account url for anonymous public read, Active Directory, shared access key authentication. + +Extra (optional) + Specify the extra parameters (as json dictionary) that can be used in Azure connection. + The following parameters are all optional: + + * ``tenant_id``: Specify the tenant to use. Needed for Active Directory (token) authentication. + * ``connection_string``: Connection string for use with connection string authentication. + diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 4c5ba768f3b05..4985983a99b35 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -452,6 +452,7 @@ "azure-servicebus>=7.6.1", "azure-storage-blob>=12.14.0", "azure-storage-common>=2.1.0", + "azure-storage-file-datalake>=12.9.1", "azure-storage-file>=2.1.0", "azure-synapse-spark" ], From bb6f8fbba9e3a233d5917a1638fe88fc9b56a486 Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Tue, 13 Dec 2022 12:01:06 +0530 Subject: [PATCH 5/6] Add Test case --- .../providers/microsoft/azure/hooks/adls.py | 93 +++++-------------- .../providers/microsoft/azure/provider.yaml | 2 - .../microsoft/azure/hooks/test_adls.py | 77 +++++++++++++++ 3 files changed, 101 insertions(+), 71 deletions(-) create mode 100644 tests/providers/microsoft/azure/hooks/test_adls.py diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py index d9a2111754e72..ea25923c70431 100644 --- a/airflow/providers/microsoft/azure/hooks/adls.py +++ b/airflow/providers/microsoft/azure/hooks/adls.py @@ -19,6 +19,7 @@ from typing import Any +from airflow.providers.microsoft.azure.hooks.wasb import WasbHook from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError from azure.identity import ClientSecretCredential from azure.storage.filedatalake import ( @@ -28,52 +29,10 @@ FileSystemClient, ) -from airflow.hooks.base import BaseHook - - -class AzureDataLakeStorageV2(BaseHook): - - conn_name_attr = "adls_v2_conn_id" - default_conn_name = "adls_v2_default" - conn_type = "adls_v2" - hook_name = "Azure Date Lake Storage" - - @staticmethod - def get_connection_form_widgets() -> dict[str, Any]: - """Returns connection widgets to add to connection form""" - from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget - from flask_babel import lazy_gettext - from wtforms import PasswordField, StringField - - return { - "extra__adls_v2__connection_string": PasswordField( - lazy_gettext("Blob Storage Connection String (optional)"), widget=BS3PasswordFieldWidget() - ), - "extra__adls_v2__tenant_id": StringField( - lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() - ), - } - - @staticmethod - def get_ui_field_behaviour() -> dict[str, Any]: - """Returns custom field behaviour""" - return { - "hidden_fields": ["schema", "port"], - "relabeling": { - "login": "Blob Storage Login (optional)", - "password": "Blob Storage Key (optional)", - "host": "Account Name (Active Directory Auth)", - }, - "placeholders": { - "login": "account name", - "password": "secret", - "host": "account url", - "extra__adls_v2__connection_string": "connection string auth", - "extra__adls_v2__tenant_id": "tenant", - }, - } - - def __init__(self, adls_v2_conn_id: str = default_conn_name, public_read: bool = False) -> None: + +class AzureDataLakeStorageGen2(WasbHook): + + def __init__(self, adls_v2_conn_id: str, public_read: bool = False) -> None: super().__init__() self.conn_id = adls_v2_conn_id self.public_read = public_read @@ -118,7 +77,6 @@ def create_file_system(self, file_system_name: str) -> None: """ try: file_system_client = self.service_client.create_file_system(file_system=file_system_name) - print("file_system_client ", file_system_client.file_system_name) self.log.info("Created file system: %s", file_system_client.file_system_name) except ResourceExistsError: self.log.info("Attempted to create file system %r but it already exists.", file_system_name) @@ -136,18 +94,13 @@ def get_file_system(self, file_system_name: str) -> FileSystemClient: self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) raise - def create_directory(self, file_system_name: str, directory_name: str) -> None: - try: - file_system = self.get_file_system(file_system_name) - file_system.create_directory(directory_name) - except Exception as e: - self.log.info(e) - raise + def create_directory(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: + result = self.get_file_system(file_system_name).create_directory(directory_name) + return result def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: try: - file_system_client = self.get_file_system(file_system_name) - directory_client = file_system_client.get_directory_client(directory_name) + directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) return directory_client except ResourceNotFoundError: self.log.info( @@ -158,13 +111,8 @@ def get_directory_client(self, file_system_name: str, directory_name: str) -> Da raise def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: - try: - file_system = self.get_file_system(file_system_name) - file_client = file_system.create_file(file_name) - return file_client - except Exception as e: - self.log.info(e) - raise + file_client = self.get_file_system(file_system_name).create_file(file_name) + return file_client def upload_file(self, file_system_name: str, @@ -182,21 +130,27 @@ def upload_file_to_directory( file_name: str, file_path: str, overwrite: bool = False, + **kwargs: any, ) -> None: + """ + Create a new file and return the file client to be interacted with and then + upload data to a file + """ directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) - file_client = directory_client.create_file(file_name) + file_client = directory_client.create_file(file_name, kwargs=kwargs) with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite) + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) - def list_directory(self, file_system_name: str, directory_name: str) -> list[str]: - file_system_client = self.get_file_system(file_system_name=file_system_name) - paths = file_system_client.get_paths(directory_name) + def list_files_directory(self, file_system_name: str, directory_name: str) -> list[str]: + """Get the list of files or directories under the specified file system""" + paths = self.get_file_system(file_system_name=file_system_name).get_paths(directory_name) directory_lists = [] for path in paths: directory_lists.append(path.name) return directory_lists def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: + """Get the list the file systems under the specified account.""" file_system = self.service_client.list_file_systems( name_starts_with=prefix, include_metadata=include_metadata ) @@ -206,7 +160,7 @@ def list_file_system(self, prefix: str | None = None, include_metadata: bool = F return file_system_list def delete_file_system(self, file_system_name: str) -> None: - """Delete file system""" + """Deletes the file system""" try: self.service_client.delete_file_system(file_system_name) self.log.info("Deleted file system: %s", file_system_name) @@ -217,5 +171,6 @@ def delete_file_system(self, file_system_name: str) -> None: raise def delete_directory(self, file_system_name: str, directory_name: str) -> None: + """Deletes specified directory in file system""" directory_client = self.get_directory_client(file_system_name, directory_name) directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index f756af7f6b2c0..d20b4e071c3e3 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -259,8 +259,6 @@ connection-types: connection-type: azure_container_registry - hook-class-name: airflow.providers.microsoft.azure.hooks.asb.BaseAzureServiceBusHook connection-type: azure_service_bus - - hook-class-name: airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageV2 - connection-type: adls_v2 - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse diff --git a/tests/providers/microsoft/azure/hooks/test_adls.py b/tests/providers/microsoft/azure/hooks/test_adls.py new file mode 100644 index 0000000000000..38faa2cd9889a --- /dev/null +++ b/tests/providers/microsoft/azure/hooks/test_adls.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import json +from unittest import mock + +import pytest +from airflow.models import Connection +from airflow.providers.microsoft.azure.hooks.adls import AzureDataLakeStorageGen2 + + +class TestAzureDataLakeStorageGen2: + def setup_class(self) -> None: + self.conn_id: str = "azure_data_lake_storage_v2" + self.file_system_name = "test_file_system" + self.directory_name = "test_directory" + self.file_name = "test_file_name" + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_file_system(self, mock_conn): + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.create_file_system("test_file_system") + expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.FileSystemClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_get_file_system(self, mock_conn, mock_file_system): + mock_conn.return_value.get_file_system_client.return_value = mock_file_system + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.get_file_system(self.file_system_name) + assert result == mock_file_system + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.create_directory.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.create_directory(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.get_directory_client(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeFileClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_create_file(self,mock_conn, mock_get_file_system, mock_file_client): + mock_get_file_system.return_value.create_file.return_value = mock_file_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + result = hook.create_file(self.file_system_name, self.file_name) + assert result == mock_file_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_delete_file_system(self, mock_conn): + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.delete_file_system(self.file_system_name) + expected_calls = [mock.call().delete_file_system(self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") + def test_delete_directory(self, mock_conn, mock_directory_client): + mock_conn.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) + hook.delete_directory(self.file_system_name, self.directory_name) + expected_calls = [mock.call().get_file_system_client( + self.file_system_name).get_directory_client(self.directory_name).delete_directory()] + mock_conn.assert_has_calls(expected_calls) + + From e4fe4c44e492ed1a064a8389f89bb626743bd0e8 Mon Sep 17 00:00:00 2001 From: bharanidharan14 Date: Tue, 13 Dec 2022 15:23:53 +0530 Subject: [PATCH 6/6] Moved hook to existing data lake hook file Fix mypy issues Fix mypy issues Fix mypy issues Fix mypy issues Add adls v2 connection Fix doc Fix doc Moved hook to existing data lake hook file Moved hook test case Moved hook test case Fix doc --- .../providers/microsoft/azure/hooks/adls.py | 176 ---------- .../microsoft/azure/hooks/data_lake.py | 311 +++++++++++++++++- .../providers/microsoft/azure/provider.yaml | 12 +- .../connections/{adls.rst => adls_v2.rst} | 26 +- .../microsoft/azure/hooks/test_adls.py | 77 ----- .../azure/hooks/test_azure_data_lake.py | 96 +++++- 6 files changed, 418 insertions(+), 280 deletions(-) delete mode 100644 airflow/providers/microsoft/azure/hooks/adls.py rename docs/apache-airflow-providers-microsoft-azure/connections/{adls.rst => adls_v2.rst} (71%) delete mode 100644 tests/providers/microsoft/azure/hooks/test_adls.py diff --git a/airflow/providers/microsoft/azure/hooks/adls.py b/airflow/providers/microsoft/azure/hooks/adls.py deleted file mode 100644 index ea25923c70431..0000000000000 --- a/airflow/providers/microsoft/azure/hooks/adls.py +++ /dev/null @@ -1,176 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -from __future__ import annotations - -from typing import Any - -from airflow.providers.microsoft.azure.hooks.wasb import WasbHook -from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError -from azure.identity import ClientSecretCredential -from azure.storage.filedatalake import ( - DataLakeDirectoryClient, - DataLakeFileClient, - DataLakeServiceClient, - FileSystemClient, -) - - -class AzureDataLakeStorageGen2(WasbHook): - - def __init__(self, adls_v2_conn_id: str, public_read: bool = False) -> None: - super().__init__() - self.conn_id = adls_v2_conn_id - self.public_read = public_read - self.service_client = self.get_conn() - - def get_conn(self) -> DataLakeServiceClient: - """Return the DataLakeServiceClient object.""" - conn = self.get_connection(self.conn_id) - extra = conn.extra_dejson or {} - - connection_string = extra.pop( - "connection_string", extra.pop("extra__adls_v2__connection_string", None) - ) - if connection_string: - # connection_string auth takes priority - return DataLakeServiceClient.from_connection_string(connection_string, **extra) - - tenant = extra.pop("tenant_id", extra.pop("extra__adls_v2__tenant_id", None)) - if tenant: - # use Active Directory auth - app_id = conn.login - app_secret = conn.password - token_credential = ClientSecretCredential(tenant, app_id, app_secret) - return DataLakeServiceClient( - account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra - ) - credential = conn.password - return DataLakeServiceClient( - account_url=f"https://{conn.login}.dfs.core.windows.net", - credential=credential, - **extra, - ) - - def create_file_system(self, file_system_name: str) -> None: - """ - A container acts as a file system for your files. Creates a new file system under - the specified account. - - If the file system with the same name already exists, a ResourceExistsError will - be raised. This method returns a client with which to interact with the newly - created file system. - """ - try: - file_system_client = self.service_client.create_file_system(file_system=file_system_name) - self.log.info("Created file system: %s", file_system_client.file_system_name) - except ResourceExistsError: - self.log.info("Attempted to create file system %r but it already exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) - raise - - def get_file_system(self, file_system_name: str) -> FileSystemClient: - try: - file_system_client = self.service_client.get_file_system_client(file_system=file_system_name) - return file_system_client - except ResourceNotFoundError: - self.log.info("file system %r doesn't exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to get file system %r: %s", file_system_name, e) - raise - - def create_directory(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: - result = self.get_file_system(file_system_name).create_directory(directory_name) - return result - - def get_directory_client(self, file_system_name: str, directory_name: str) -> DataLakeDirectoryClient: - try: - directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) - return directory_client - except ResourceNotFoundError: - self.log.info( - "Directory %s doesn't exists in the file system %s", directory_name, file_system_name - ) - except Exception as e: - self.log.info(e) - raise - - def create_file(self, file_system_name: str, file_name: str) -> DataLakeFileClient: - file_client = self.get_file_system(file_system_name).create_file(file_name) - return file_client - - def upload_file(self, - file_system_name: str, - file_name: str, - file_path: str, - overwrite: bool = False) -> None: - file_client = self.create_file(file_system_name, file_name) - with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite) - - def upload_file_to_directory( - self, - file_system_name: str, - directory_name: str, - file_name: str, - file_path: str, - overwrite: bool = False, - **kwargs: any, - ) -> None: - """ - Create a new file and return the file client to be interacted with and then - upload data to a file - """ - directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) - file_client = directory_client.create_file(file_name, kwargs=kwargs) - with open(file_path, "rb") as data: - file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) - - def list_files_directory(self, file_system_name: str, directory_name: str) -> list[str]: - """Get the list of files or directories under the specified file system""" - paths = self.get_file_system(file_system_name=file_system_name).get_paths(directory_name) - directory_lists = [] - for path in paths: - directory_lists.append(path.name) - return directory_lists - - def list_file_system(self, prefix: str | None = None, include_metadata: bool = False) -> list[str]: - """Get the list the file systems under the specified account.""" - file_system = self.service_client.list_file_systems( - name_starts_with=prefix, include_metadata=include_metadata - ) - file_system_list = [] - for fs in file_system: - file_system_list.append(fs.name) - return file_system_list - - def delete_file_system(self, file_system_name: str) -> None: - """Deletes the file system""" - try: - self.service_client.delete_file_system(file_system_name) - self.log.info("Deleted file system: %s", file_system_name) - except ResourceNotFoundError: - self.log.info("file system %r doesn't exists.", file_system_name) - except Exception as e: - self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) - raise - - def delete_directory(self, file_system_name: str, directory_name: str) -> None: - """Deletes specified directory in file system""" - directory_client = self.get_directory_client(file_system_name, directory_name) - directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/hooks/data_lake.py b/airflow/providers/microsoft/azure/hooks/data_lake.py index 4a9bc98f925fb..ad5e9818a958a 100644 --- a/airflow/providers/microsoft/azure/hooks/data_lake.py +++ b/airflow/providers/microsoft/azure/hooks/data_lake.py @@ -15,19 +15,21 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -""" -This module contains integration with Azure Data Lake. - -AzureDataLakeHook communicates via a REST API compatible with WebHDFS. Make sure that a -Airflow connection of type `azure_data_lake` exists. Authorization can be done by supplying a -login (=Client ID), password (=Client Secret) and extra fields tenant (Tenant) and account_name (Account Name) -(see connection `azure_data_lake_default` for an example). -""" from __future__ import annotations from typing import Any +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError from azure.datalake.store import core, lib, multithread +from azure.identity import ClientSecretCredential +from azure.storage.filedatalake import ( + DataLakeDirectoryClient, + DataLakeFileClient, + DataLakeServiceClient, + DirectoryProperties, + FileSystemClient, + FileSystemProperties, +) from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook @@ -36,6 +38,13 @@ class AzureDataLakeHook(BaseHook): """ + This module contains integration with Azure Data Lake. + + AzureDataLakeHook communicates via a REST API compatible with WebHDFS. Make sure that a + Airflow connection of type `azure_data_lake` exists. Authorization can be done by supplying a + login (=Client ID), password (=Client Secret) and extra fields tenant (Tenant) and account_name + (Account Name)(see connection `azure_data_lake_default` for an example). + Interacts with Azure Data Lake. Client ID and client secret should be in user and password parameters. @@ -230,3 +239,289 @@ def remove(self, path: str, recursive: bool = False, ignore_not_found: bool = Tr self.log.info("File %s not found", path) else: raise AirflowException(f"File {path} not found") + + +class AzureDataLakeStorageV2Hook(BaseHook): + """ + This Hook interacts with ADLS gen2 storage account it mainly helps to create and manage + directories and files in storage accounts that have a hierarchical namespace. Using Adls_v2 connection + details create DataLakeServiceClient object + + Due to Wasb is marked as legacy and and retirement of the (ADLS1) it would be nice to + implement ADLS gen2 hook for interacting with the storage account. + + .. seealso:: + https://learn.microsoft.com/en-us/azure/storage/blobs/data-lake-storage-directory-file-acl-python + + :param adls_conn_id: Reference to the :ref:`adls connection `. + :param public_read: Whether an anonymous public read access should be used. default is False + """ + + conn_name_attr = "adls_conn_id" + default_conn_name = "adls_default" + conn_type = "adls" + hook_name = "Azure Date Lake Storage V2" + + @staticmethod + def get_connection_form_widgets() -> dict[str, Any]: + """Returns connection widgets to add to connection form""" + from flask_appbuilder.fieldwidgets import BS3PasswordFieldWidget, BS3TextFieldWidget + from flask_babel import lazy_gettext + from wtforms import PasswordField, StringField + + return { + "connection_string": PasswordField( + lazy_gettext("ADLS gen2 Connection String (optional)"), widget=BS3PasswordFieldWidget() + ), + "tenant_id": StringField( + lazy_gettext("Tenant Id (Active Directory Auth)"), widget=BS3TextFieldWidget() + ), + } + + @staticmethod + def get_ui_field_behaviour() -> dict[str, Any]: + """Returns custom field behaviour""" + return { + "hidden_fields": ["schema", "port"], + "relabeling": { + "login": "ADLS gen2 Storage Login (optional)", + "password": "ADLS gen2 Storage Key (optional)", + "host": "Account Name (Active Directory Auth)", + }, + "placeholders": { + "extra": "additional options for use with FileService and AzureFileVolume", + "login": "account name", + "password": "secret", + "host": "account url", + "connection_string": "connection string auth", + "tenant_id": "tenant", + }, + } + + def __init__(self, adls_conn_id: str, public_read: bool = False) -> None: + super().__init__() + self.conn_id = adls_conn_id + self.public_read = public_read + self.service_client = self.get_conn() + + def get_conn(self) -> DataLakeServiceClient: # type: ignore[override] + """Return the DataLakeServiceClient object.""" + conn = self.get_connection(self.conn_id) + extra = conn.extra_dejson or {} + + connection_string = self._get_field(extra, "connection_string") + if connection_string: + # connection_string auth takes priority + return DataLakeServiceClient.from_connection_string(connection_string, **extra) + + tenant = self._get_field(extra, "tenant_id") + if tenant: + # use Active Directory auth + app_id = conn.login + app_secret = conn.password + token_credential = ClientSecretCredential(tenant, app_id, app_secret) + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", credential=token_credential, **extra + ) + credential = conn.password + return DataLakeServiceClient( + account_url=f"https://{conn.login}.dfs.core.windows.net", + credential=credential, + **extra, + ) + + def _get_field(self, extra_dict, field_name): + prefix = "extra__adls__" + if field_name.startswith("extra__"): + raise ValueError( + f"Got prefixed name {field_name}; please remove the '{prefix}' prefix " + f"when using this method." + ) + if field_name in extra_dict: + return extra_dict[field_name] or None + return extra_dict.get(f"{prefix}{field_name}") or None + + def create_file_system(self, file_system_name: str) -> None: + """ + A container acts as a file system for your files. Creates a new file system under + the specified account. + + If the file system with the same name already exists, a ResourceExistsError will + be raised. This method returns a client with which to interact with the newly + created file system. + """ + try: + file_system_client = self.service_client.create_file_system(file_system=file_system_name) + self.log.info("Created file system: %s", file_system_client.file_system_name) + except ResourceExistsError: + self.log.info("Attempted to create file system %r but it already exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to create file system %r: %s", file_system_name, e) + raise + + def get_file_system(self, file_system: FileSystemProperties | str) -> FileSystemClient: + """ + Get a client to interact with the specified file system + + :param file_system: This can either be the name of the file system + or an instance of FileSystemProperties. + """ + try: + file_system_client = self.service_client.get_file_system_client(file_system=file_system) + return file_system_client + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system) + raise + except Exception as e: + self.log.info("Error while attempting to get file system %r: %s", file_system, e) + raise + + def create_directory( + self, file_system_name: FileSystemProperties | str, directory_name: str, **kwargs + ) -> DataLakeDirectoryClient: + """ + Create a directory under the specified file system. + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory which needs to be created in the file system. + """ + result = self.get_file_system(file_system_name).create_directory(directory_name, kwargs) + return result + + def get_directory_client( + self, + file_system_name: FileSystemProperties | str, + directory_name: DirectoryProperties | str, + ) -> DataLakeDirectoryClient: + """ + Get the specific directory under the specified file system. + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory or instance of DirectoryProperties which needs to be + retrieved from the file system. + """ + try: + directory_client = self.get_file_system(file_system_name).get_directory_client(directory_name) + return directory_client + except ResourceNotFoundError: + self.log.info( + "Directory %s doesn't exists in the file system %s", directory_name, file_system_name + ) + raise + except Exception as e: + self.log.info(e) + raise + + def create_file(self, file_system_name: FileSystemProperties | str, file_name: str) -> DataLakeFileClient: + """ + Creates a file under the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param file_name: Name of the file which needs to be created in the file system. + """ + file_client = self.get_file_system(file_system_name).create_file(file_name) + return file_client + + def upload_file( + self, + file_system_name: FileSystemProperties | str, + file_name: str, + file_path: str, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """ + Create a file with data in the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param file_name: Name of the file to be created with name. + :param file_path: Path to the file to load. + :param overwrite: Boolean flag to overwrite an existing file or not. + """ + file_client = self.create_file(file_system_name, file_name) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) + + def upload_file_to_directory( + self, + file_system_name: str, + directory_name: str, + file_name: str, + file_path: str, + overwrite: bool = False, + **kwargs: Any, + ) -> None: + """ + Create a new file and return the file client to be interacted with and then + upload data to a file + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + :param file_name: Name of the file to be created with name. + :param file_path: Path to the file to load. + :param overwrite: Boolean flag to overwrite an existing file or not. + """ + directory_client = self.get_directory_client(file_system_name, directory_name=directory_name) + file_client = directory_client.create_file(file_name, kwargs=kwargs) + with open(file_path, "rb") as data: + file_client.upload_data(data, overwrite=overwrite, kwargs=kwargs) + + def list_files_directory( + self, file_system_name: FileSystemProperties | str, directory_name: str + ) -> list[str]: + """ + Get the list of files or directories under the specified file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + """ + paths = self.get_file_system(file_system=file_system_name).get_paths(directory_name) + directory_lists = [] + for path in paths: + directory_lists.append(path.name) + return directory_lists + + def list_file_system( + self, prefix: str | None = None, include_metadata: bool = False, **kwargs: Any + ) -> list[str]: + """ + Get the list the file systems under the specified account. + + :param prefix: + Filters the results to return only file systems whose names + begin with the specified prefix. + :param include_metadata: Specifies that file system metadata be returned in the response. + The default value is `False`. + """ + file_system = self.service_client.list_file_systems( + name_starts_with=prefix, include_metadata=include_metadata + ) + file_system_list = [] + for fs in file_system: + file_system_list.append(fs.name) + return file_system_list + + def delete_file_system(self, file_system_name: FileSystemProperties | str) -> None: + """ + Deletes the file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + """ + try: + self.service_client.delete_file_system(file_system_name) + self.log.info("Deleted file system: %s", file_system_name) + except ResourceNotFoundError: + self.log.info("file system %r doesn't exists.", file_system_name) + except Exception as e: + self.log.info("Error while attempting to deleting file system %r: %s", file_system_name, e) + raise + + def delete_directory(self, file_system_name: FileSystemProperties | str, directory_name: str) -> None: + """ + Deletes specified directory in file system + + :param file_system_name: Name of the file system or instance of FileSystemProperties. + :param directory_name: Name of the directory. + """ + directory_client = self.get_directory_client(file_system_name, directory_name) + directory_client.delete_directory() diff --git a/airflow/providers/microsoft/azure/provider.yaml b/airflow/providers/microsoft/azure/provider.yaml index d20b4e071c3e3..41f890fd6f9e0 100644 --- a/airflow/providers/microsoft/azure/provider.yaml +++ b/airflow/providers/microsoft/azure/provider.yaml @@ -124,10 +124,8 @@ integrations: how-to-guide: - /docs/apache-airflow-providers-microsoft-azure/operators/azure_synapse.rst tags: [azure] - - integration-name: Microsoft Azure Data Lake Storage gen2 - how-to-guide: - - /docs/apache-airflow-providers-microsoft-azure/operators/adls.rst - external-doc-url: https://azure.microsoft.com/en-us/services/storage/data-lake-storage/ + - integration-name: Microsoft Azure Data Lake Storage Client Gen2 + external-doc-url: https://azure.microsoft.com/en-us/products/storage/data-lake-storage/ logo: /integration-logos/azure/Data Lake Storage.svg tags: [azure] @@ -204,9 +202,9 @@ hooks: - integration-name: Microsoft Azure Service Bus python-modules: - airflow.providers.microsoft.azure.hooks.asb - - integration-name: Microsoft Azure Data Lake Storage Gen2 + - integration-name: Microsoft Azure Data Lake Storage Client Gen2 python-modules: - - airflow.providers.microsoft.azure.hooks.adls + - airflow.providers.microsoft.azure.hooks.data_lake - integration-name: Microsoft Azure Synapse python-modules: - airflow.providers.microsoft.azure.hooks.synapse @@ -261,6 +259,8 @@ connection-types: connection-type: azure_service_bus - hook-class-name: airflow.providers.microsoft.azure.hooks.synapse.AzureSynapseHook connection-type: azure_synapse + - hook-class-name: airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook + connection-type: adls secrets-backends: - airflow.providers.microsoft.azure.secrets.key_vault.AzureKeyVaultBackend diff --git a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst b/docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst similarity index 71% rename from docs/apache-airflow-providers-microsoft-azure/connections/adls.rst rename to docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst index c531e67685405..9ec4015679def 100644 --- a/docs/apache-airflow-providers-microsoft-azure/connections/adls.rst +++ b/docs/apache-airflow-providers-microsoft-azure/connections/adls_v2.rst @@ -15,19 +15,17 @@ specific language governing permissions and limitations under the License. +.. _howto/connection:adls: +Microsoft Azure Data Lake Storage Gen2 Connection +================================================== -.. _howto/connection:adls_v2: +The Microsoft Azure Data Lake Storage Gen2 connection type enables the ADLS gen2 Integrations. -Microsoft Azure Data Lake Storage Connection Gen2 -================================================= +Authenticating to Azure Data Lake Storage Gen2 +---------------------------------------------- -The Microsoft Azure Data Lake Storage connection type enables the Azure Data Storage Integrations. - -Authenticating to Azure Data Lake Storage ------------------------------------------ - -There are four ways to connect to Azure Data Lake Storage using Airflow. +Currently, there are two ways to connect to Azure Data Lake Storage Gen2 using Airflow. 1. Use `token credentials `_ @@ -42,16 +40,16 @@ configure multiple connections. Default Connection IDs ---------------------- -All hooks and operators related to Microsoft Azure Data Lake Storage use ``adls_v2_default`` by default. +All hooks and operators related to Microsoft Azure Blob Storage use ``azure_data_lake_default`` by default. Configuring the Connection -------------------------- Login (optional) - Specify the login used for azure storage. + Specify the login used for azure blob storage. For use with Shared Key Credential and SAS Token authentication. Password (optional) - Specify the password used for azure storage. For use with + Specify the password used for azure blob storage. For use with Active Directory (token credential) and shared key authentication. Host (optional) @@ -64,3 +62,7 @@ Extra (optional) * ``tenant_id``: Specify the tenant to use. Needed for Active Directory (token) authentication. * ``connection_string``: Connection string for use with connection string authentication. +When specifying the connection in environment variable you should specify +it using URI syntax. + +Note that all components of the URI should be URL-encoded. diff --git a/tests/providers/microsoft/azure/hooks/test_adls.py b/tests/providers/microsoft/azure/hooks/test_adls.py deleted file mode 100644 index 38faa2cd9889a..0000000000000 --- a/tests/providers/microsoft/azure/hooks/test_adls.py +++ /dev/null @@ -1,77 +0,0 @@ -from __future__ import annotations - -import json -from unittest import mock - -import pytest -from airflow.models import Connection -from airflow.providers.microsoft.azure.hooks.adls import AzureDataLakeStorageGen2 - - -class TestAzureDataLakeStorageGen2: - def setup_class(self) -> None: - self.conn_id: str = "azure_data_lake_storage_v2" - self.file_system_name = "test_file_system" - self.directory_name = "test_directory" - self.file_name = "test_file_name" - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_file_system(self, mock_conn): - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.create_file_system("test_file_system") - expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] - mock_conn.assert_has_calls(expected_calls) - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.FileSystemClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_get_file_system(self, mock_conn, mock_file_system): - mock_conn.return_value.get_file_system_client.return_value = mock_file_system - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.get_file_system(self.file_system_name) - assert result == mock_file_system - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): - mock_get_file_system.return_value.create_directory.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.create_directory(self.file_system_name, self.directory_name) - assert result == mock_directory_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): - mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.get_directory_client(self.file_system_name, self.directory_name) - assert result == mock_directory_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeFileClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_file_system") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_create_file(self,mock_conn, mock_get_file_system, mock_file_client): - mock_get_file_system.return_value.create_file.return_value = mock_file_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - result = hook.create_file(self.file_system_name, self.file_name) - assert result == mock_file_client - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_delete_file_system(self, mock_conn): - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.delete_file_system(self.file_system_name) - expected_calls = [mock.call().delete_file_system(self.file_system_name)] - mock_conn.assert_has_calls(expected_calls) - - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.DataLakeDirectoryClient") - @mock.patch("airflow.providers.microsoft.azure.hooks.adls.AzureDataLakeStorageGen2.get_conn") - def test_delete_directory(self, mock_conn, mock_directory_client): - mock_conn.return_value.get_directory_client.return_value = mock_directory_client - hook = AzureDataLakeStorageGen2(adls_v2_conn_id=self.conn_id) - hook.delete_directory(self.file_system_name, self.directory_name) - expected_calls = [mock.call().get_file_system_client( - self.file_system_name).get_directory_client(self.directory_name).delete_directory()] - mock_conn.assert_has_calls(expected_calls) - - diff --git a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py index f2b2d5a65439f..e8d378ab5c12a 100644 --- a/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py +++ b/tests/providers/microsoft/azure/hooks/test_azure_data_lake.py @@ -21,7 +21,7 @@ from unittest import mock from airflow.models import Connection -from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook +from airflow.providers.microsoft.azure.hooks.data_lake import AzureDataLakeHook, AzureDataLakeStorageV2Hook from airflow.utils import db from tests.test_utils.providers import get_provider_min_airflow_version @@ -151,3 +151,97 @@ def test_get_ui_field_behaviour_placeholders(self): "You must now remove `_ensure_prefixes` from azure utils." " The functionality is now taken care of by providers manager." ) + + +class TestAzureDataLakeStorageV2Hook: + def setup_class(self) -> None: + self.conn_id: str = "adls_conn_id" + self.file_system_name = "test_file_system" + self.directory_name = "test_directory" + self.file_name = "test_file_name" + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.create_file_system("test_file_system") + expected_calls = [mock.call().create_file_system(file_system=self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.FileSystemClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_get_file_system(self, mock_conn, mock_file_system): + mock_conn.return_value.get_file_system_client.return_value = mock_file_system + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.get_file_system(self.file_system_name) + assert result == mock_file_system + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.create_directory.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.create_directory(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_get_directory(self, mock_conn, mock_get_file_system, mock_directory_client): + mock_get_file_system.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.get_directory_client(self.file_system_name, self.directory_name) + assert result == mock_directory_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeFileClient") + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_create_file(self, mock_conn, mock_get_file_system, mock_file_client): + mock_get_file_system.return_value.create_file.return_value = mock_file_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + result = hook.create_file(self.file_system_name, self.file_name) + assert result == mock_file_client + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_delete_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.delete_file_system(self.file_system_name) + expected_calls = [mock.call().delete_file_system(self.file_system_name)] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.DataLakeDirectoryClient") + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_delete_directory(self, mock_conn, mock_directory_client): + mock_conn.return_value.get_directory_client.return_value = mock_directory_client + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.delete_directory(self.file_system_name, self.directory_name) + expected_calls = [ + mock.call() + .get_file_system_client(self.file_system_name) + .get_directory_client(self.directory_name) + .delete_directory() + ] + mock_conn.assert_has_calls(expected_calls) + + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_list_file_system(self, mock_conn): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.list_file_system(prefix="prefix") + mock_conn.return_value.list_file_systems.assert_called_once_with( + name_starts_with="prefix", include_metadata=False + ) + + @mock.patch( + "airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_file_system" + ) + @mock.patch("airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook.get_conn") + def test_list_files_directory(self, mock_conn, mock_get_file_system): + hook = AzureDataLakeStorageV2Hook(adls_conn_id=self.conn_id) + hook.list_files_directory(self.file_system_name, self.directory_name) + mock_get_file_system.return_value.get_paths.assert_called_once_with(self.directory_name)