diff --git a/dev/breeze/src/airflow_breeze/global_constants.py b/dev/breeze/src/airflow_breeze/global_constants.py index 060cb2dccf6b6..272787c06199c 100644 --- a/dev/breeze/src/airflow_breeze/global_constants.py +++ b/dev/breeze/src/airflow_breeze/global_constants.py @@ -106,8 +106,9 @@ OTEL_INTEGRATION = "otel" OPENLINEAGE_INTEGRATION = "openlineage" OPENSEARCH_INTEGRATION = "opensearch" +OZONE_INTEGRATION = "ozone" OTHER_CORE_INTEGRATIONS = [STATSD_INTEGRATION, KEYCLOAK_INTEGRATION] -OTHER_PROVIDERS_INTEGRATIONS = [OPENLINEAGE_INTEGRATION, OPENSEARCH_INTEGRATION] +OTHER_PROVIDERS_INTEGRATIONS = [OPENLINEAGE_INTEGRATION, OPENSEARCH_INTEGRATION, OZONE_INTEGRATION] ALLOWED_DEBIAN_VERSIONS = ["bookworm"] ALL_CORE_INTEGRATIONS = sorted( [ @@ -474,6 +475,8 @@ def get_default_platform_machine() -> str: GREMLIN_HOST_PORT = "8182" MSSQL_HOST_PORT = "21433" MYSQL_HOST_PORT = "23306" +OZONE_OM_HOST_PORT = "29862" +OZONE_SCM_HOST_PORT = "29860" POSTGRES_HOST_PORT = "25433" RABBITMQ_HOST_PORT = "25672" REDIS_HOST_PORT = "26379" diff --git a/dev/breeze/src/airflow_breeze/params/shell_params.py b/dev/breeze/src/airflow_breeze/params/shell_params.py index 4905fb200136f..205ef1356cc41 100644 --- a/dev/breeze/src/airflow_breeze/params/shell_params.py +++ b/dev/breeze/src/airflow_breeze/params/shell_params.py @@ -62,6 +62,8 @@ MOUNT_TESTS, MSSQL_HOST_PORT, MYSQL_HOST_PORT, + OZONE_OM_HOST_PORT, + OZONE_SCM_HOST_PORT, POSTGRES_BACKEND, POSTGRES_HOST_PORT, RABBITMQ_HOST_PORT, @@ -677,6 +679,8 @@ def env_variables_for_docker_commands(self) -> dict[str, str]: _set_var(_env, "MSSQL_HOST_PORT", None, MSSQL_HOST_PORT) _set_var(_env, "MYSQL_HOST_PORT", None, MYSQL_HOST_PORT) _set_var(_env, "MYSQL_VERSION", self.mysql_version) + _set_var(_env, "OZONE_OM_HOST_PORT", None, OZONE_OM_HOST_PORT) + _set_var(_env, "OZONE_SCM_HOST_PORT", None, OZONE_SCM_HOST_PORT) _set_var(_env, "MOUNT_SOURCES", self.mount_sources) _set_var(_env, "MOUNT_UI_DIST", self.mount_ui_dist) _set_var(_env, "NUM_RUNS", self.num_runs) diff --git a/generated/provider_metadata.json b/generated/provider_metadata.json index 0dd414cfe7c45..77bb774742ea2 100644 --- a/generated/provider_metadata.json +++ b/generated/provider_metadata.json @@ -673,6 +673,12 @@ "date_released": "2026-03-13T17:32:34Z" } }, + "arenadata.ozone": { + "1.1.0": { + "associated_airflow_version": "3.2.1", + "date_released": "2026-06-18T00:00:00Z" + } + }, "apache.beam": { "1.0.0": { "associated_airflow_version": "2.1.0", diff --git a/providers/arenadata/ozone/.gitignore b/providers/arenadata/ozone/.gitignore new file mode 100644 index 0000000000000..bff2d7629604d --- /dev/null +++ b/providers/arenadata/ozone/.gitignore @@ -0,0 +1 @@ +*.iml diff --git a/providers/arenadata/ozone/CHANGELOG.rst b/providers/arenadata/ozone/CHANGELOG.rst new file mode 100644 index 0000000000000..2c11828d0b186 --- /dev/null +++ b/providers/arenadata/ozone/CHANGELOG.rst @@ -0,0 +1,57 @@ +.. 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. + + +.. NOTE TO CONTRIBUTORS: + Please, only add notes to the Changelog just below the "Changelog" header when there are some breaking changes + and you want to add an explanation to the users on how they are supposed to deal with them. + The changelog is updated and maintained semi-automatically by release manager. + +``apache-airflow-providers-arenadata-ozone`` + + +Changelog +--------- + +1.1.0 +..... + +Ozone provider migration: + +* Migrated the provider to the Apache Airflow 3.2.1 provider package layout. + +1.0.1 +..... + +Ozone provider improvements: + +* Replaced module-level environment-derived constants in Ozone example DAGs with DAG ``params`` / + ``Param(...)`` definitions. +* Split ``OzoneKeySensor`` CLI timeout handling from the parent sensor timeout. The public + ``timeout`` parameter now controls one Ozone CLI existence check and is stored internally as + ``cli_timeout``. +* Added ``if_exists`` support for selected Ozone FS operations. +* Added ``ExistingTargetPolicy`` enum for existing-target behavior. +* Added ``OzoneFsHook.make_path`` as the preferred path creation API. ``OzoneFsHook.create_path`` + is kept as a deprecated compatibility wrapper; its ``fail_if_exists`` parameter is deprecated in + favor of ``if_exists``. +* Added Kerberos authentication via principal and password. + +1.0.0 +..... + +Initial release of the Apache Ozone provider. diff --git a/providers/arenadata/ozone/LICENSE b/providers/arenadata/ozone/LICENSE new file mode 100644 index 0000000000000..35ad4bd192864 --- /dev/null +++ b/providers/arenadata/ozone/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2026 Arenadata company + +Licensed 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. diff --git a/providers/arenadata/ozone/NOTICE b/providers/arenadata/ozone/NOTICE new file mode 100644 index 0000000000000..a51bd9390d030 --- /dev/null +++ b/providers/arenadata/ozone/NOTICE @@ -0,0 +1,5 @@ +Apache Airflow +Copyright 2016-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). diff --git a/providers/arenadata/ozone/README.rst b/providers/arenadata/ozone/README.rst new file mode 100644 index 0000000000000..dbee24cc7ffdf --- /dev/null +++ b/providers/arenadata/ozone/README.rst @@ -0,0 +1,85 @@ + +.. 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. + +.. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + +.. IF YOU WANT TO MODIFY TEMPLATE FOR THIS FILE, YOU SHOULD MODIFY THE TEMPLATE + ``PROVIDER_README_TEMPLATE.rst.jinja2`` IN the ``dev/breeze/src/airflow_breeze/templates`` DIRECTORY + +Package ``apache-airflow-providers-arenadata-ozone`` + +Release: ``1.1.0`` + + +`Apache Ozone `__ + +Provider package for interacting with Apache Ozone via: + +- Native CLI (``ozone``) for ``ofs://`` / ``o3fs://`` paths +- Native Ozone administration and filesystem workflows + + +Provider package +---------------- + +This is a provider package for ``arenadata.ozone`` provider. All classes for this provider package +are in ``airflow.providers.arenadata.ozone`` python package. + +You can find package information and changelog for the provider +in the `documentation `_. + +Installation +------------ + +You can install this package on top of an existing Airflow installation (see ``Requirements`` below +for the minimum Airflow version supported) via +``pip install apache-airflow-providers-arenadata-ozone`` + +The package supports the following python versions: 3.10,3.11,3.12,3.13,3.14 + +Requirements +------------ + +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=2.11.0`` +``apache-airflow-providers-common-compat`` ``>=1.12.0`` +========================================== ================== + +Cross provider package dependencies +----------------------------------- + +Those are dependencies that might be needed in order to use all the features of the package. +You need to install the specified providers in order to use them. + +You can install such cross-provider dependencies when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-arenadata-ozone[common.compat] + + +================================================================================================================== ================= +Dependent package Extra +================================================================================================================== ================= +`apache-airflow-providers-common-compat `_ ``common.compat`` +================================================================================================================== ================= + +The changelog for the provider package can be found in the +`changelog `_. diff --git a/providers/arenadata/ozone/docs/changelog.rst b/providers/arenadata/ozone/docs/changelog.rst new file mode 100644 index 0000000000000..1628c9344be70 --- /dev/null +++ b/providers/arenadata/ozone/docs/changelog.rst @@ -0,0 +1,25 @@ + + .. 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. + + .. NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE + OVERWRITTEN WHEN PREPARING PACKAGES. + + .. IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE + `PROVIDER_CHANGELOG_TEMPLATE.rst.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY + +.. include:: ../CHANGELOG.rst diff --git a/providers/arenadata/ozone/docs/commits.rst b/providers/arenadata/ozone/docs/commits.rst new file mode 100644 index 0000000000000..1e9802470054a --- /dev/null +++ b/providers/arenadata/ozone/docs/commits.rst @@ -0,0 +1,18 @@ + .. 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. + + .. THIS FILE IS UPDATED AUTOMATICALLY_AT_RELEASE_TIME diff --git a/providers/arenadata/ozone/docs/conf.py b/providers/arenadata/ozone/docs/conf.py new file mode 100644 index 0000000000000..0cff6439d47c2 --- /dev/null +++ b/providers/arenadata/ozone/docs/conf.py @@ -0,0 +1,27 @@ +# Disable Flake8 because of all the sphinx imports +# +# 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. +"""Configuration of Providers docs building.""" + +from __future__ import annotations + +import os + +os.environ["AIRFLOW_PACKAGE_NAME"] = "apache-airflow-providers-arenadata-ozone" + +from docs.provider_conf import * # noqa: F403 diff --git a/providers/arenadata/ozone/docs/connections.rst b/providers/arenadata/ozone/docs/connections.rst new file mode 100644 index 0000000000000..0d2d6a0751708 --- /dev/null +++ b/providers/arenadata/ozone/docs/connections.rst @@ -0,0 +1,338 @@ +.. 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:ozone: + +Apache Ozone Connections +======================== + +The ``ozone`` connection type is the main runtime contract for this provider. +It is used by hooks/operators/sensors/transfers for Native Ozone CLI execution +(``ozone sh`` / ``ozone fs``) and for security/runtime settings. + +Provider runtime reads this connection via: +``airflow/providers/arenadata/ozone/utils/connection_schema.py`` +(``OzoneConnSnapshot.from_connection(...)``). + +The worker runtime must have: + +* Java runtime +* Ozone CLI binary available in ``PATH`` +* client config directory/files reachable from the worker runtime + +Default connection IDs +---------------------- + +* ``ozone_default`` is used by ``OzoneCliHook`` and ``OzoneFsHook`` by default. +* ``ozone_admin_default`` is used by ``OzoneAdminHook`` by default. + +Connection fields +----------------- + +Host (required in Ozone CLI scope) + Ozone Manager hostname for CLI endpoint wiring. + +Port (required in Ozone CLI scope) + Ozone Manager port for CLI endpoint wiring. + +Extra + JSON dictionary with canonical provider keys. Built-in provider code reads + only fields defined in ``OzoneConnSnapshot``. + + Additional custom keys are allowed and available to DAG developers through + ``hook.connection_snapshot.raw_extra``. + +Connection extra keys (canonical) +--------------------------------- + +Ozone SSL/TLS scope: + +* ``ozone_security_enabled`` +* ``ozone_om_https_port`` +* ``ozone_ssl_keystore_location`` +* ``ozone_ssl_keystore_password`` (supports ``secret://...``) +* ``ozone_ssl_keystore_type`` +* ``ozone_ssl_truststore_location`` +* ``ozone_ssl_truststore_password`` (supports ``secret://...``) +* ``ozone_ssl_truststore_type`` + +Ozone Kerberos scope: + +* ``hadoop_security_authentication`` (set ``kerberos`` to enable) +* ``kerberos_principal`` +* ``kerberos_keytab`` (supports ``secret://...``) +* ``kerberos_password`` (supports ``secret://...``) +* ``krb5_conf`` (optional) +* ``ozone_conf_dir`` (recommended for all modes) + +Connection extra by scenario +---------------------------- + +Use one of the following templates depending on your runtime mode. + +1) Plain Ozone (no SSL, no Kerberos) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: json + + { + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +2) Ozone with SSL/TLS +^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: json + + { + "ozone_security_enabled": "true", + "ozone_om_https_port": "9879", + "ozone_ssl_keystore_location": "/opt/airflow/ssl/ozone-keystore.jks", + "ozone_ssl_keystore_password": "secret://vault/ozone/keystore_password", + "ozone_ssl_keystore_type": "JKS", + "ozone_ssl_truststore_location": "/opt/airflow/ssl/ozone-truststore.jks", + "ozone_ssl_truststore_password": "secret://vault/ozone/truststore_password", + "ozone_ssl_truststore_type": "JKS", + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +3) Ozone with SSL/TLS and Kerberos using keytab +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: json + + { + "ozone_security_enabled": "true", + "ozone_om_https_port": "9879", + "ozone_ssl_keystore_location": "/opt/airflow/ssl/ozone-keystore.jks", + "ozone_ssl_keystore_password": "secret://vault/ozone/keystore_password", + "ozone_ssl_keystore_type": "JKS", + "ozone_ssl_truststore_location": "/opt/airflow/ssl/ozone-truststore.jks", + "ozone_ssl_truststore_password": "secret://vault/ozone/truststore_password", + "ozone_ssl_truststore_type": "JKS", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "testuser@EXAMPLE.COM", + "kerberos_keytab": "/opt/airflow/keytabs/testuser.keytab", + "krb5_conf": "/opt/airflow/kerberos-config/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +4) Ozone with SSL/TLS and Kerberos using password +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: json + + { + "ozone_security_enabled": "true", + "ozone_om_https_port": "9879", + "ozone_ssl_keystore_location": "/opt/airflow/ssl/ozone-keystore.jks", + "ozone_ssl_keystore_password": "secret://vault/ozone/keystore_password", + "ozone_ssl_keystore_type": "JKS", + "ozone_ssl_truststore_location": "/opt/airflow/ssl/ozone-truststore.jks", + "ozone_ssl_truststore_password": "secret://vault/ozone/truststore_password", + "ozone_ssl_truststore_type": "JKS", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "testuser@EXAMPLE.COM", + "kerberos_password": "secret://vault/ozone/kerberos_password", + "krb5_conf": "/opt/airflow/kerberos-config/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +Field-by-field explanation for SSL + Kerberos extra +---------------------------------------------------- + +``ozone_security_enabled`` + Enables Ozone security mode in provider runtime. Use ``"true"`` for SSL and + SSL+Kerberos scenarios. + +``ozone_om_https_port`` + HTTPS port used by Ozone OM in your cluster (for example ``9879`` in local + demo setups). Must match actual OM SSL config. + +``ozone_ssl_keystore_location`` / ``ozone_ssl_truststore_location`` + Absolute paths inside the Airflow worker runtime to keystore and truststore + files used by Ozone CLI/JVM SSL layer. + +``ozone_ssl_keystore_password`` / ``ozone_ssl_truststore_password`` + Passwords for the stores above. Prefer ``secret://...`` values so passwords + are resolved from Airflow secrets backend instead of hardcoding plain text. + +``ozone_ssl_keystore_type`` / ``ozone_ssl_truststore_type`` + Store format (typically ``JKS``). Must match how files were created. + +``hadoop_security_authentication`` + Kerberos switch for provider runtime. Set exactly to ``"kerberos"`` to + enable Kerberos env wiring and ``kinit`` flow. + +``kerberos_principal`` + Principal used by tasks before CLI calls (for example + ``testuser@EXAMPLE.COM``). + +``kerberos_keytab`` + Absolute path to a keytab file inside the Airflow worker runtime. + The file must exist and be readable by the task process. When both + ``kerberos_keytab`` and ``kerberos_password`` are provided, the provider + keeps the existing keytab-based flow. + +``kerberos_password`` + Password for ``kerberos_principal``. Prefer ``secret://...`` values so the + password is resolved from Airflow secrets backend. The provider uses the + password only for ``kinit`` stdin and does not export it to Ozone CLI + subprocess environment. + +``krb5_conf`` + Optional explicit path to ``krb5.conf``. Recommended in containerized runs + to avoid relying on container-image defaults. + +``ozone_conf_dir`` + Directory with Ozone/Hadoop client configs (usually ``ozone-site.xml`` and + ``core-site.xml``). This is one of the most important fields: provider uses + it to set ``OZONE_CONF_DIR`` and (when needed) ``HADOOP_CONF_DIR``. + +``host`` and ``port`` still matter even when ``ozone_conf_dir`` points to a +valid runtime config directory: the provider uses them to prepare CLI endpoint +wiring for task execution. + +Common validation checklist for DAG developers +---------------------------------------------- + +* Paths in extra must be valid inside the worker runtime, not on your laptop. +* SSL store files and keytab files, when configured, must be mounted/readable + for Airflow task user. +* Kerberos requires ``kerberos_principal`` plus either ``kerberos_keytab`` or + ``kerberos_password``. +* ``ozone_conf_dir`` must contain config files pointing to real OM/SCM endpoints. + +Why ``ozone_conf_dir`` is important +----------------------------------- + +``ozone_conf_dir`` points to the directory inside the worker runtime where +Ozone/Hadoop client configs are available (usually ``ozone-site.xml`` and +``core-site.xml``). + +The provider uses this value to prepare CLI runtime environment +(``OZONE_CONF_DIR`` and, when needed, ``HADOOP_CONF_DIR``), so ``ozone sh`` / +``ozone fs`` can resolve OM/SCM addresses and security policy. + +If ``ozone_conf_dir`` is missing or points to the wrong path, commands may fail +even with correct connection host/port (for example: unresolved endpoints, +auth/security initialization failures, or generic Ozone CLI connection errors). + +HDFS DistCp transfer scope (when ``hdfs_conn_id`` is provided) +--------------------------------------------------------------- + +``HdfsToOzoneOperator`` reads HDFS-specific security keys from the connection referenced by +``hdfs_conn_id`` and applies them only to DistCp subprocess environment. + +HDFS SSL keys: + +* ``hdfs_ssl_enabled`` +* ``dfs_encrypt_data_transfer`` +* ``hdfs_ssl_keystore_location`` +* ``hdfs_ssl_keystore_password`` (supports ``secret://...``) +* ``hdfs_ssl_truststore_location`` +* ``hdfs_ssl_truststore_password`` (supports ``secret://...``) + +HDFS Kerberos keys: + +* ``hdfs_kerberos_enabled`` +* ``hdfs_kerberos_principal`` +* ``hdfs_kerberos_keytab`` (supports ``secret://...``) +* ``hdfs_kerberos_password`` (supports ``secret://...``) +* optional ``krb5_conf`` for explicit Kerberos config path +* optional ``hdfs_distcp_mapreduce_local`` to run DistCp through local MapReduce + when no YARN/MapReduce cluster config is available on the worker +* optional ``hdfs_distcp_renewer_principal`` to pass a Kerberos renewer principal + to DistCp delegation token setup + +For HDFS DistCp Kerberos, provide ``hdfs_kerberos_principal`` plus either +``hdfs_kerberos_keytab`` or ``hdfs_kerberos_password``. When both credentials +are configured, the provider keeps the existing keytab-based flow. Passwords +are used only for ``kinit`` stdin and are not exported to the DistCp subprocess +environment. + +Example HDFS Kerberos extra for ``HdfsToOzoneOperator`` using password: + +.. code-block:: json + + { + "hdfs_kerberos_enabled": "true", + "hdfs_kerberos_principal": "hdfs@EXAMPLE.COM", + "hdfs_kerberos_password": "secret://vault/hdfs/kerberos_password", + "hdfs_distcp_mapreduce_local": "true", + "hdfs_distcp_renewer_principal": "hdfs@EXAMPLE.COM", + "krb5_conf": "/opt/airflow/kerberos-config/krb5.conf" + } + +Recommended runtime notes +------------------------- + +* Set ``ozone_conf_dir`` explicitly in connection extra for plain/SSL/Kerberos modes. +* Keep all security-sensitive values in Airflow Connection/Secrets backend, + not in DAG source code. +* Do not rely on key aliases: only canonical key names are supported. + +Recommended Kerberos extra (explicit paths): + +.. code-block:: json + + { + "ozone_security_enabled": "true", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "testuser@EXAMPLE.COM", + "kerberos_keytab": "/opt/airflow/keytabs/testuser.keytab", + "krb5_conf": "/opt/airflow/kerberos-config/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +Recommended Kerberos extra (password from secrets backend): + +.. code-block:: json + + { + "ozone_security_enabled": "true", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "testuser@EXAMPLE.COM", + "kerberos_password": "secret://vault/ozone/kerberos_password", + "krb5_conf": "/opt/airflow/kerberos-config/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf" + } + +Upload/download size limit +-------------------------- + +The provider supports a connection-level limit for payload/file transfers: + +* ``max_content_size_bytes`` (default: ``1073741824`` = 1 GiB) + +This value is read from ``OzoneConnSnapshot`` and is used by: + +* ``OzoneUploadContentOperator`` +* ``OzoneUploadFileOperator`` +* ``OzoneDownloadFileOperator`` (via ``OzoneFsHook.get_key_property(...)[\"data_size\"]``) + +Each of these operators also accepts an optional ``max_content_size_bytes`` parameter +to override the connection-level value for that task. + +Testing connections in Airflow UI +--------------------------------- + +The ``ozone`` connection type implements ``test_connection()``: + +* ``ozone``: runs a minimal CLI command against Ozone Manager + and returns Airflow-compatible status/message. diff --git a/providers/arenadata/ozone/docs/example-dags.rst b/providers/arenadata/ozone/docs/example-dags.rst new file mode 100644 index 0000000000000..d35df38c611b0 --- /dev/null +++ b/providers/arenadata/ozone/docs/example-dags.rst @@ -0,0 +1,309 @@ +.. 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. + +Apache Ozone Example DAGs +========================= + +This provider ships example DAGs that demonstrate native Ozone workflows. +Each example DAG now supports two layers of configuration: + +* environment variables (used only as defaults during DAG parsing); +* DAG ``params`` / Trigger UI / ``dag_run.conf`` overrides at run time. + +In practice this means: + +* local dev scripts and CI can continue using ``OZONE_EXAMPLE_*`` environment variables; +* testers can trigger the DAG manually from Airflow UI and override values without editing source code + or changing worker environment. + +Where example DAGs are located +------------------------------ + +``airflow/providers/arenadata/ozone/example_dags/`` + +Main examples: + +* ``example_ozone_usage.py`` +* ``example_ozone_copy_from_hdfs.py`` +* ``example_ozone_data_lifecycle.py`` +* ``example_ozone_multi_tenant_management.py`` + +What each example DAG does +-------------------------- + +``example_ozone_usage.py`` (plain / SSL / Kerberos / SSL+Kerberos): + +* Creates a volume and bucket with admin operators. +* Creates an ``ofs://`` directory path in Ozone FS. +* Uploads inline text content to a file in that directory. +* Waits for the file with ``OzoneKeySensor``. +* Overwrites the demo upload target, so the same DAG can be triggered + repeatedly without manually deleting the previous file. +* Works in all Ozone security modes; pick mode via + ``admin_conn_id`` / ``OZONE_EXAMPLE_USAGE_ADMIN_CONN_ID`` and corresponding + ``Connection Extra``. +* Best for first smoke-check of provider installation and CLI connectivity. + +``example_ozone_copy_from_hdfs.py`` (HDFS -> Ozone copy): + +* Waits for a trigger marker file in Ozone landing zone. +* Creates the trigger marker file after a short delay for standalone demos. +* Applies storage quota to the target bucket. +* Migrates data from HDFS to Ozone with ``HdfsToOzoneOperator`` (DistCp). +* Useful as a template for ingestion pipelines where Ozone is destination storage. + +``example_ozone_data_lifecycle.py`` (archive + backup + cleanup): + +* Ensures landing/archive volumes and buckets exist. +* Lists files in landing zone with ``OzoneListOperator``. +* Simulates processing stage (``TaskGroup`` with a processing task). +* Moves processed files into date-partitioned archive path. +* Creates a snapshot backup and then cleans original landing files. +* Demonstrates full lifecycle automation: discover -> process -> archive -> backup -> cleanup. + +``example_ozone_multi_tenant_management.py`` (provisioning): + +* Creates a dedicated project volume. +* Sets volume quota. +* Creates standard landing/processed buckets with bucket quotas. +* Creates standard ``data`` directories in each bucket. +* Useful for platform teams that provision isolated tenant/project storage. + +How configuration is applied +---------------------------- + +Example DAG configuration works in this order: + +1. built-in fallback in DAG code; +2. ``OZONE_EXAMPLE_*`` environment variable default (if present); +3. manual override from Trigger UI / ``dag_run.conf`` for a specific DAG run. + +The environment-variable layer exists mainly for local bootstrap scripts and CI. +The recommended interactive workflow for testers is: + +1. open the DAG in Airflow UI; +2. click ``Trigger DAG``; +3. change values in the generated params form; +4. or paste JSON into ``Trigger DAG with config``. + +The underlying implementation uses DAG-level ``params`` and templated operator +fields, so values from ``dag_run.conf`` are rendered into operator arguments +before task execution. + +Common variable +--------------- + +Used by all examples: + +* ``OZONE_EXAMPLE_OM_HOST`` + +Common runtime param: + +* ``om_host`` + +Per-example variables +--------------------- + +Basic usage (``example_ozone_usage.py``): + +Environment defaults: + +* ``OZONE_EXAMPLE_USAGE_ADMIN_CONN_ID`` +* ``OZONE_EXAMPLE_USAGE_VOLUME`` +* ``OZONE_EXAMPLE_USAGE_BUCKET`` +* ``OZONE_EXAMPLE_USAGE_DIR`` +* ``OZONE_EXAMPLE_USAGE_FILE`` +* ``OZONE_EXAMPLE_USAGE_VOLUME_QUOTA`` +* ``OZONE_EXAMPLE_USAGE_BUCKET_QUOTA`` + +Trigger/UI params: + +* ``admin_conn_id`` +* ``volume`` +* ``bucket`` +* ``directory`` +* ``file_name`` +* ``volume_quota`` +* ``bucket_quota`` + +To switch mode for this single usage DAG, point +``OZONE_EXAMPLE_USAGE_ADMIN_CONN_ID`` to a connection configured for: + +* plain Ozone; +* Ozone + SSL; +* Ozone + Kerberos; +* Ozone + SSL + Kerberos. + +For Kerberos mode, required connection-extra keys include +``hadoop_security_authentication=kerberos``, ``kerberos_principal``, +either ``kerberos_keytab`` or ``kerberos_password``, optional ``krb5_conf``, +and ``ozone_conf_dir``. + +Example trigger config: + +.. code-block:: json + + { + "om_host": "adho", + "admin_conn_id": "ozone_admin_default", + "volume": "vol1", + "bucket": "bucket-native", + "directory": "data_dir", + "file_name": "file.txt", + "volume_quota": "10GB", + "bucket_quota": "10GB" + } + +Ozone copy from HDFS (``example_ozone_copy_from_hdfs.py``): + +This DAG models an HDFS -> Ozone copy flow. It waits for ``trigger_file`` with +``OzoneKeySensor`` and, for standalone demos, starts a parallel helper branch +that waits 10 seconds and creates the same trigger file with +``OzoneUploadContentOperator`` using ``if_exists="ignore"``. In a real pipeline +this helper branch can be removed and the trigger file can be produced by an +upstream ingestion job. + +Environment defaults: + +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_CONN_ID`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_HDFS_CONN_ID`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_VOLUME`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_BUCKET`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_TRIGGER_FILE`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_QUOTA`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_SOURCE_PATH`` +* ``OZONE_EXAMPLE_COPY_FROM_HDFS_DEST_SUBPATH`` + +Trigger/UI params: + +* ``pipeline_conn_id`` +* ``hdfs_conn_id`` +* ``volume`` +* ``bucket`` +* ``trigger_file`` +* ``quota`` +* ``source_path`` +* ``dest_subpath`` + +Example trigger config: + +.. code-block:: json + + { + "om_host": "adho", + "pipeline_conn_id": "ozone_admin_default", + "hdfs_conn_id": "hdfs_admin_default", + "volume": "vol1", + "bucket": "bucket1", + "trigger_file": "trigger.lck", + "quota": "5GB", + "source_path": "hdfs:///user/data/legacy/", + "dest_subpath": "migrated/" + } + +Data lifecycle (``example_ozone_data_lifecycle.py``): + +Environment defaults: + +* ``OZONE_EXAMPLE_LIFECYCLE_CONN_ID`` +* ``OZONE_EXAMPLE_LIFECYCLE_LANDING_VOLUME`` +* ``OZONE_EXAMPLE_LIFECYCLE_LANDING_BUCKET`` +* ``OZONE_EXAMPLE_LIFECYCLE_ARCHIVE_VOLUME`` +* ``OZONE_EXAMPLE_LIFECYCLE_ARCHIVE_BUCKET`` +* ``OZONE_EXAMPLE_LIFECYCLE_SNAPSHOT_PREFIX`` + +Trigger/UI params: + +* ``lifecycle_conn_id`` +* ``landing_volume`` +* ``landing_bucket`` +* ``archive_volume`` +* ``archive_bucket`` +* ``snapshot_prefix`` + +Example trigger config: + +.. code-block:: json + + { + "om_host": "adho", + "lifecycle_conn_id": "ozone_admin_default", + "landing_volume": "landing", + "landing_bucket": "raw", + "archive_volume": "archive", + "archive_bucket": "processed", + "snapshot_prefix": "snap" + } + +Multi-tenant management (``example_ozone_multi_tenant_management.py``): + +Environment defaults: + +* ``OZONE_EXAMPLE_MULTI_TENANT_CONN_ID`` +* ``OZONE_EXAMPLE_MULTI_TENANT_PROJECT_VOLUME`` +* ``OZONE_EXAMPLE_MULTI_TENANT_PROJECT_QUOTA`` +* ``OZONE_EXAMPLE_MULTI_TENANT_LANDING_BUCKET`` +* ``OZONE_EXAMPLE_MULTI_TENANT_PROCESSED_BUCKET`` +* ``OZONE_EXAMPLE_MULTI_TENANT_BUCKET_QUOTA`` + +Trigger/UI params: + +* ``multi_tenant_conn_id`` +* ``project_volume`` +* ``project_quota`` +* ``landing_bucket`` +* ``processed_bucket`` +* ``bucket_quota`` + +Example trigger config: + +.. code-block:: json + + { + "om_host": "adho", + "multi_tenant_conn_id": "ozone_admin_default", + "project_volume": "project-alpha", + "project_quota": "10GB", + "landing_bucket": "landing", + "processed_bucket": "processed", + "bucket_quota": "1GB" + } + +DAG developer notes +------------------- + +* Example DAG params and ``OZONE_EXAMPLE_*`` variables are for demonstration only + and are not part of provider runtime API. +* Production DAGs should keep operational parameters in task args and Airflow Connections. +* Connection parsing for runtime behavior is defined by + ``airflow/providers/arenadata/ozone/utils/connection_schema.py``. +* Custom DAG extensions can read non-provider extra keys through + ``hook.connection_snapshot.raw_extra``. + +Provider runtime tuning (for example runs) +------------------------------------------ + +Example DAGs use the same provider runtime policy as regular tasks. +Retry/timeout behavior is configured in provider code and per-task arguments: + +* Hook defaults: ``RETRY_ATTEMPTS``, ``FAST_TIMEOUT_SECONDS``, ``SLOW_TIMEOUT_SECONDS`` + (see ``airflow/providers/arenadata/ozone/hooks/ozone.py``). +* Operators/transfers can override ``retry_attempts`` and ``timeout`` per task + where needed. +* ``OzoneKeySensor.timeout`` configures one Ozone CLI existence check and is + stored internally as ``cli_timeout``. It is not forwarded to + ``BaseSensorOperator.timeout``. diff --git a/providers/arenadata/ozone/docs/hooks.rst b/providers/arenadata/ozone/docs/hooks.rst new file mode 100644 index 0000000000000..dc122ee770fcb --- /dev/null +++ b/providers/arenadata/ozone/docs/hooks.rst @@ -0,0 +1,115 @@ +.. 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. + +Apache Ozone Hooks +================== + +The provider keeps most Ozone runtime logic in hooks. Operators and sensors are +thin by design and delegate CLI execution, connection parsing, security wiring, +and retries to this layer. + +OzoneCliHook +------------ + +``OzoneCliHook`` is the base hook for native Ozone CLI execution. + +Main responsibilities: + +* read Airflow connection data into ``OzoneConnSnapshot``; +* prepare Ozone CLI environment variables; +* apply SSL and Kerberos settings from connection extra; +* run commands through the common CLI runner with retries and timeouts; +* implement ``test_connection()`` for the ``ozone`` connection type. + +This hook uses ``ozone_default`` by default. + +OzoneFsHook +----------- + +``OzoneFsHook`` extends ``OzoneCliHook`` for filesystem workflows based on +``ozone fs``. + +Typical API areas: + +* create and delete paths; +* check path existence; +* list keys and wildcard matches; +* upload and download files; +* move and copy data inside Ozone; +* inspect key metadata and selected key properties. + +Write helpers use the ``ExistingTargetPolicy`` enum for already existing targets. +DAG code may still pass the string values directly: + +* ``ExistingTargetPolicy.ERROR`` / ``"error"``: fail fast with a clear Airflow + exception; +* ``ExistingTargetPolicy.IGNORE`` / ``"ignore"``: treat the existing target as + success; +* ``ExistingTargetPolicy.OVERWRITE`` / ``"overwrite"``: only for file uploads, + where the provider deletes the existing key first and then writes through + ``ozone sh key put``. + +Use ``make_path(..., if_exists=...)`` for new path creation code. The older +``create_path(..., fail_if_exists=...)`` method remains available only for +backwards compatibility, logs a deprecation warning, and delegates to +``make_path`` internally. + +This is the main hook behind filesystem operators and sensors. + +OzoneAdminHook +-------------- + +``OzoneAdminHook`` extends ``OzoneCliHook`` for administrative workflows based +on ``ozone sh``. + +Typical API areas: + +* create and delete volumes; +* create and delete buckets; +* inspect volume and bucket metadata; +* set and clear quotas; +* list volumes and buckets. + +Volume and bucket creation defaults to idempotent ``if_exists="ignore"`` and +can be switched to ``if_exists="error"`` when a DAG should fail if the resource +already exists. + +This hook uses ``ozone_admin_default`` by default. + +OzoneAdminExtraHook +------------------- + +``OzoneAdminExtraHook`` extends ``OzoneAdminHook`` with advanced administrative +operations. + +Typical API areas: + +* snapshot management; +* bucket link creation; +* bucket replication configuration; +* tenant management; +* SCM container reports. + +When to use hooks directly +-------------------------- + +Using operators and sensors is preferred for standard DAG tasks. +Use hooks directly when you need: + +* custom Python branching around Ozone results; +* direct access to structured metadata in task code; +* provider extension code that should reuse connection parsing and CLI helpers. diff --git a/providers/arenadata/ozone/docs/index.rst b/providers/arenadata/ozone/docs/index.rst new file mode 100644 index 0000000000000..e8860e37120c6 --- /dev/null +++ b/providers/arenadata/ozone/docs/index.rst @@ -0,0 +1,235 @@ + +.. 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. + +``apache-airflow-providers-arenadata-ozone`` +============================================ + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Basics + + Home + Changelog + Security + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Guides + + Connections + Operators + Sensors + Transfers + Hooks + Example DAGs + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Resources + + Installing from sources + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: System tests + + System Tests <_api/tests/system/arenadata/ozone/index> + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Commits + + Detailed list of commits + + +Package overview +---------------- + +`Apache Ozone `__ provider package for Airflow. + +Release: ``1.1.0`` + +Python package path: +``airflow.providers.arenadata.ozone`` + +Current provider scope +---------------------- + +This provider is Native CLI-first and supports Ozone workflows through +``ozone sh`` and ``ozone fs`` commands: + +* Ozone administration (volume/bucket lifecycle, quotas, admin extras) +* Ozone filesystem operations (create/list/delete/move/copy/upload/download) +* Ozone path sensors +* HDFS -> Ozone DistCp transfer +* Ozone bucket snapshot backup operator +* Connection-driven SSL and Kerberos runtime wiring + +Runtime deployment model +------------------------ + +This provider does not provision Ozone client runtime by itself. +It is expected to run in one of the following environments: + +* on a shared worker host or container image where Ozone client runtime is already + installed and maintained together with Airflow; +* on a worker host or container image prepared manually with the required + Ozone runtime files. + +For manual preparation, the worker runtime should be provisioned with: + +* Java runtime installed on the worker host or available in the container image; +* Ozone CLI runtime copied to a stable location such as ``/opt/ozone``; +* Ozone client configuration copied to a stable directory such as + ``/etc/ozone/conf``; +* at minimum, ``core-site.xml`` and ``ozone-site.xml`` available in that + configuration directory; +* environment variables wired for task runtime, typically + ``JAVA_HOME``, ``JAVA``, ``OZONE_HOME``, ``OZONE_LIBEXEC_DIR``, + ``OZONE_CONF_DIR``, and ``HADOOP_CONF_DIR``; +* ``ozone`` available in ``PATH`` for interactive shell runs and for Airflow + service runtime; +* working hostname resolution for Ozone Manager and related cluster endpoints, + whether via DNS or explicit host mappings. + +In practical terms, this means the provider works best when Airflow runs close +to an Ozone-enabled runtime, or when the same runtime is reproduced explicitly +on the worker host or in the container image. + +Quick start checklist +--------------------- + +For a minimal working setup: + +1. Prepare Ozone client runtime on the worker host or in the container image. +2. Make ``ozone`` available in ``PATH`` together with Java runtime. +3. Place ``core-site.xml`` and ``ozone-site.xml`` in a directory available to + the worker, then point ``ozone_conf_dir`` to that directory in the Airflow + connection. +4. Create an ``ozone`` connection in Airflow and verify it with + ``test_connection()`` or an example DAG. + +Not in scope +------------ + +* No built-in S3 layer in this provider +* No compatibility aliases for connection keys + +Connection contract +------------------- + +Runtime connection parsing is centralized in: +``airflow/providers/arenadata/ozone/utils/connection_schema.py`` +(``OzoneConnSnapshot``). + +Built-in provider runtime uses typed snapshot fields. +Raw ``Connection.extra`` is exposed as ``snapshot.raw_extra`` only for +custom DAG-level extensions. + +Requirements +------------ + +* ``apache-airflow`` >= ``2.11.0`` + +Example DAGs +------------ + +Example DAGs are located in: +``airflow/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags``. + +.. THE REMAINDER OF THE FILE IS AUTOMATICALLY GENERATED. IT WILL BE OVERWRITTEN AT RELEASE TIME! + + +.. toctree:: + :hidden: + :maxdepth: 1 + :caption: Commits + + Detailed list of commits + + +apache-airflow-providers-arenadata-ozone package +------------------------------------------------------ + +`Apache Ozone `__ + +Provider package for interacting with Apache Ozone via: + +- Native CLI (``ozone``) for ``ofs://`` / ``o3fs://`` paths +- Native Ozone administration and filesystem workflows + + +Release: 1.1.0 + +Provider package +---------------- + +This package is for the ``arenadata.ozone`` provider. +All classes for this package are included in the ``airflow.providers.arenadata.ozone`` python package. + +Installation +------------ + +You can install this package on top of an existing Airflow installation via +``pip install apache-airflow-providers-arenadata-ozone``. +For the minimum Airflow version supported, see ``Requirements`` below. + +Requirements +------------ + +The minimum Apache Airflow version supported by this provider distribution is ``2.11.0``. + +========================================== ================== +PIP package Version required +========================================== ================== +``apache-airflow`` ``>=2.11.0`` +``apache-airflow-providers-common-compat`` ``>=1.12.0`` +========================================== ================== + +Cross provider package dependencies +----------------------------------- + +Those are dependencies that might be needed in order to use all the features of the package. +You need to install the specified provider distributions in order to use them. + +You can install such cross-provider dependencies when installing from PyPI. For example: + +.. code-block:: bash + + pip install apache-airflow-providers-arenadata-ozone[common.compat] + + +================================================================================================================== ================= +Dependent package Extra +================================================================================================================== ================= +`apache-airflow-providers-common-compat `_ ``common.compat`` +================================================================================================================== ================= + +Downloading official packages +----------------------------- + +You can download officially released packages and verify their checksums and signatures from the +`Official Apache Download site `_ + +* `The apache-airflow-providers-arenadata-ozone 1.1.0 sdist package `_ (`asc `__, `sha512 `__) +* `The apache-airflow-providers-arenadata-ozone 1.1.0 wheel package `_ (`asc `__, `sha512 `__) diff --git a/providers/arenadata/ozone/docs/installing-providers-from-sources.rst b/providers/arenadata/ozone/docs/installing-providers-from-sources.rst new file mode 100644 index 0000000000000..085f5d55c1b29 --- /dev/null +++ b/providers/arenadata/ozone/docs/installing-providers-from-sources.rst @@ -0,0 +1,19 @@ + + .. 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. + +.. include:: ../exts/includes/installing-providers-from-sources.rst diff --git a/providers/arenadata/ozone/docs/operators.rst b/providers/arenadata/ozone/docs/operators.rst new file mode 100644 index 0000000000000..20198c82ec29d --- /dev/null +++ b/providers/arenadata/ozone/docs/operators.rst @@ -0,0 +1,188 @@ +.. 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. + +Apache Ozone Operators +====================== + +The Ozone provider is CLI-first: operators delegate most runtime work to +``OzoneFsHook`` or ``OzoneAdminHook`` and execute native ``ozone`` commands. + +Prerequisites +------------- + +To use these operators, configure an :doc:`Ozone Connection ` +and make sure the worker runtime has: + +* Java runtime available to the worker. +* ``ozone`` CLI available in ``PATH``. +* ``ozone-site.xml`` and ``core-site.xml`` reachable through ``ozone_conf_dir`` + or otherwise available in the worker runtime. + +Common operator parameters +-------------------------- + +Most operators accept: + +* ``ozone_conn_id``: Airflow connection ID. Defaults to ``ozone_default`` for + filesystem operations and ``ozone_admin_default`` for administrative + operations. +* ``retry_attempts``: number of CLI retries for transient failures. +* ``timeout``: timeout in seconds for the underlying CLI command. +* ``if_exists`` on create/upload/copy/move operators: explicit + ``ExistingTargetPolicy`` for an already existing destination. DAG code may pass + the enum values or their string forms. Supported values are ``error`` and + ``ignore``; upload operators also support ``overwrite``. + +Administrative operators +------------------------ + +These operators use ``OzoneAdminHook`` and work with Ozone volumes, buckets, +and quotas. + +* ``OzoneCreateVolumeOperator``: create a volume if it does not exist. +* ``OzoneCreateBucketOperator``: create a bucket inside a target volume. +* ``OzoneSetQuotaOperator``: apply a quota to a volume or bucket. +* ``OzoneDeleteVolumeOperator``: delete a volume, optionally recursively. +* ``OzoneDeleteBucketOperator``: delete a bucket, optionally recursively. + +Example: + +.. code-block:: python + + from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreateBucketOperator, + OzoneCreateVolumeOperator, + OzoneSetQuotaOperator, + ) + + create_volume = OzoneCreateVolumeOperator( + task_id="create_volume", + volume_name="analytics", + ozone_conn_id="ozone_admin_default", + ) + + create_bucket = OzoneCreateBucketOperator( + task_id="create_bucket", + volume_name="analytics", + bucket_name="landing", + ozone_conn_id="ozone_admin_default", + ) + + set_quota = OzoneSetQuotaOperator( + task_id="set_quota", + volume="analytics", + bucket="landing", + quota="20GB", + ozone_conn_id="ozone_admin_default", + ) + +Filesystem operators +-------------------- + +These operators use ``OzoneFsHook`` and work with ``ofs://`` or ``o3fs://`` +paths. + +Path and object creation: + +* ``OzoneCreatePathOperator``: create a directory path in Ozone FS. +* ``OzoneUploadContentOperator``: write inline text content to a temporary file + and upload it to Ozone. +* ``OzoneUploadFileOperator``: upload a local file to Ozone. + +Existing-target behavior is intentionally explicit: + +* ``if_exists="error"`` fails before running a potentially expensive CLI write + when the destination already exists. +* ``if_exists="ignore"`` treats an existing destination as success, which is + useful for idempotent marker files and setup paths. +* ``if_exists="overwrite"`` is available for uploads only. The provider deletes + the existing key first and then writes through ``ozone sh key put``. Directory + creation, copy, and move operations do not support overwrite to avoid + accidental data loss. + +For hook-level path creation, new code should use +``OzoneFsHook.make_path(..., if_exists=...)``. The older +``OzoneFsHook.create_path(..., fail_if_exists=...)`` API is kept for backwards +compatibility only and logs a deprecation warning before delegating to +``make_path``. + +Path and object removal: + +* ``OzoneDeleteKeyOperator``: delete one key or a wildcard-selected key set. +* ``OzoneDeletePathOperator``: delete a path, optionally recursively. + +Path inspection: + +* ``OzonePathExistsOperator``: return ``True`` or ``False`` via XCom depending + on path existence. +* ``OzoneListOperator``: list keys/paths and return the result via XCom. + +In-cluster file movement: + +* ``OzoneMoveOperator``: move or rename a key/path inside Ozone. +* ``OzoneCopyOperator``: copy a key/path inside Ozone. +* ``OzoneDownloadFileOperator``: download a remote key to local filesystem. + +Example: + +.. code-block:: python + + from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreatePathOperator, + OzoneListOperator, + OzoneUploadContentOperator, + ) + + create_path = OzoneCreatePathOperator( + task_id="create_path", + path="ofs://om-service/analytics/landing/incoming", + ozone_conn_id="ozone_default", + if_exists="ignore", + ) + + upload_marker = OzoneUploadContentOperator( + task_id="upload_marker", + content="ready", + remote_path="ofs://om-service/analytics/landing/incoming/_SUCCESS", + ozone_conn_id="ozone_default", + if_exists="error", + ) + + list_files = OzoneListOperator( + task_id="list_files", + path="ofs://om-service/analytics/landing/incoming/*", + ozone_conn_id="ozone_default", + ) + +Size-aware upload and download +------------------------------ + +``OzoneUploadContentOperator``, ``OzoneUploadFileOperator``, and +``OzoneDownloadFileOperator`` support ``max_content_size_bytes``. + +If this parameter is not passed explicitly, the operator uses the +connection-level limit from ``OzoneConnSnapshot.max_content_size_bytes``. +This helps protect workers from unexpectedly large payloads. + +Design notes +------------ + +The operators intentionally stay thin: + +* validation and CLI execution live in hooks and utils; +* SSL and Kerberos wiring come from the connection contract; +* operator code focuses on task arguments and Airflow integration. diff --git a/providers/arenadata/ozone/docs/security.rst b/providers/arenadata/ozone/docs/security.rst new file mode 100644 index 0000000000000..b8f95e6ecfa29 --- /dev/null +++ b/providers/arenadata/ozone/docs/security.rst @@ -0,0 +1,19 @@ + + .. 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. + +.. include:: ../exts/includes/security.rst diff --git a/providers/arenadata/ozone/docs/sensors.rst b/providers/arenadata/ozone/docs/sensors.rst new file mode 100644 index 0000000000000..fdc603bec41f6 --- /dev/null +++ b/providers/arenadata/ozone/docs/sensors.rst @@ -0,0 +1,68 @@ +.. 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. + +Apache Ozone Sensors +==================== + +Prerequisites +------------- + +To use sensors, configure an :doc:`Ozone Connection ` and +make sure the worker runtime can execute the native ``ozone`` CLI. + +OzoneKeySensor +-------------- + +``OzoneKeySensor`` waits for a file or directory to appear in Ozone FS. +It is useful for: + +* waiting for marker files such as ``_SUCCESS``; +* waiting for upstream ingestion into a landing path; +* coordinating DAGs that exchange data through Ozone. + +The sensor uses ``OzoneFsHook`` internally and checks a target ``ofs://`` or +``o3fs://`` path with the provider's standard connection-driven runtime +configuration. + +Behavior notes +-------------- + +* ``path`` is templated, so Jinja expressions can be used in DAGs. +* ``timeout`` is the timeout in seconds for a single Ozone CLI existence check. + The sensor stores this value internally as ``cli_timeout`` before calling + ``OzoneFsHook.key_exists(...)``. +* ``timeout`` is intentionally not forwarded to ``BaseSensorOperator`` and does + not control the overall waiting horizon of the sensor. Use standard Airflow + sensor arguments, such as ``poke_interval`` and other ``BaseSensorOperator`` + options, for sensor scheduling behavior. +* retryable CLI errors are treated as "not yet ready" and do not fail the task + immediately; +* non-retryable CLI errors are raised to Airflow as task failures. + +Example: + +.. code-block:: python + + from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor + + wait_for_marker = OzoneKeySensor( + task_id="wait_for_marker", + path="ofs://om-service/analytics/landing/{{ ds_nodash }}/_SUCCESS", + ozone_conn_id="ozone_default", + timeout=60, + poke_interval=30, + ) diff --git a/providers/arenadata/ozone/docs/transfers.rst b/providers/arenadata/ozone/docs/transfers.rst new file mode 100644 index 0000000000000..6b90b36a8955f --- /dev/null +++ b/providers/arenadata/ozone/docs/transfers.rst @@ -0,0 +1,102 @@ +.. 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. + +Apache Ozone Transfers +====================== + +The provider currently exposes two transfer-oriented building blocks: + +* ``HdfsToOzoneOperator`` for bulk migration from HDFS to Ozone with + ``hadoop distcp``; +* ``OzoneBackupOperator`` for bucket snapshot creation as a backup or + lifecycle checkpoint step. + +HdfsToOzoneOperator +------------------- + +``HdfsToOzoneOperator`` migrates data from HDFS to Ozone using native +``hadoop distcp``. + +Typical use cases: + +* initial migration from HDFS to Ozone; +* ingestion pipelines where Ozone is the destination storage; +* bulk copy of large path trees that should stay outside Python memory. + +Runtime notes: + +* the worker must have ``hadoop`` in ``PATH``; +* the operator validates runtime dependencies before starting DistCp; +* HDFS SSL and Kerberos settings are read from ``hdfs_conn_id`` and applied + only to the DistCp subprocess environment; +* secure workers without YARN/MapReduce cluster configuration can set + ``hdfs_distcp_mapreduce_local`` and ``hdfs_distcp_renewer_principal`` in + ``hdfs_conn_id`` extra so DistCp can acquire HDFS/Ozone delegation tokens + and run through the local MapReduce runner; +* if ``hdfs_conn_id`` is omitted, no HDFS-specific Airflow connection settings + are injected and DistCp relies on the worker's local Hadoop runtime setup; +* a dedicated ``apache-airflow-providers-apache-hdfs`` package is not required + by this operator implementation; +* Ozone destination access is resolved by the Hadoop/Ozone client runtime + available on the worker, not by a dedicated ``ozone_conn_id`` argument in + this operator. + +Example: + +.. code-block:: python + + from airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone import HdfsToOzoneOperator + + migrate_to_ozone = HdfsToOzoneOperator( + task_id="migrate_to_ozone", + source_path="hdfs:///warehouse/source_table", + dest_path="ofs://om-service/analytics/landing/source_table", + hdfs_conn_id="hdfs_admin_default", + ) + +OzoneBackupOperator +------------------- + +``OzoneBackupOperator`` creates a bucket snapshot through native Ozone admin +CLI commands. + +Typical use cases: + +* create a recovery point before cleanup; +* preserve a landing bucket before archive/move operations; +* add an idempotent safety checkpoint to lifecycle DAGs. + +Behavior notes: + +* ``volume``, ``bucket``, and ``snapshot_name`` are templated; +* if the target snapshot already exists, the operator treats it as success; +* runtime and authentication are handled through ``ozone_admin_default`` or a + custom admin connection. + +Example: + +.. code-block:: python + + from airflow.providers.arenadata.ozone.transfers.ozone_backup import OzoneBackupOperator + + create_snapshot = OzoneBackupOperator( + task_id="create_snapshot", + volume="analytics", + bucket="landing", + snapshot_name="landing_{{ ts_nodash }}", + ozone_conn_id="ozone_admin_default", + ) diff --git a/providers/arenadata/ozone/provider.yaml b/providers/arenadata/ozone/provider.yaml new file mode 100644 index 0000000000000..c2aa85bdda4c9 --- /dev/null +++ b/providers/arenadata/ozone/provider.yaml @@ -0,0 +1,78 @@ +# 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. + +--- +package-name: apache-airflow-providers-arenadata-ozone +name: Apache Ozone +description: | + `Apache Ozone `__ + + Provider package for interacting with Apache Ozone via: + + - Native CLI (``ozone``) for ``ofs://`` / ``o3fs://`` paths + - Native Ozone administration and filesystem workflows + +state: ready +lifecycle: production +source-date-epoch: 1751688000 +# note that those versions are maintained by release manager - do not update them manually +versions: + - 1.1.0 + - 1.0.1 + - 1.0.0 + +dependencies: + - apache-airflow>=2.11.0 + - apache-airflow-providers-common-compat>=1.12.0 + +integrations: + - integration-name: Apache Ozone + external-doc-url: https://ozone.apache.org/ + tags: [apache] + +hooks: + - integration-name: Apache Ozone + python-modules: + - airflow.providers.arenadata.ozone.hooks.ozone + +operators: + - integration-name: Apache Ozone + python-modules: + - airflow.providers.arenadata.ozone.operators.ozone + +sensors: + - integration-name: Apache Ozone + python-modules: + - airflow.providers.arenadata.ozone.sensors.ozone + +transfers: + - source-integration-name: Hadoop Distributed File System (HDFS) + target-integration-name: Apache Ozone + python-module: airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone + - source-integration-name: Apache Ozone + target-integration-name: Apache Ozone + python-module: airflow.providers.arenadata.ozone.transfers.ozone_backup + +connection-types: + - hook-class-name: airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliHook + connection-type: ozone + +example-dags: + - airflow.providers.arenadata.ozone.example_dags.example_ozone_usage + - airflow.providers.arenadata.ozone.example_dags.example_ozone_copy_from_hdfs + - airflow.providers.arenadata.ozone.example_dags.example_ozone_data_lifecycle + - airflow.providers.arenadata.ozone.example_dags.example_ozone_multi_tenant_management diff --git a/providers/arenadata/ozone/pyproject.toml b/providers/arenadata/ozone/pyproject.toml new file mode 100644 index 0000000000000..a8dfd96cc5fde --- /dev/null +++ b/providers/arenadata/ozone/pyproject.toml @@ -0,0 +1,113 @@ +# 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. + +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! + +# IF YOU WANT TO MODIFY THIS FILE EXCEPT DEPENDENCIES, YOU SHOULD MODIFY THE TEMPLATE +# `pyproject_TEMPLATE.toml.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY +[build-system] +requires = ["flit_core==3.12.0"] +build-backend = "flit_core.buildapi" + +[project] +name = "apache-airflow-providers-arenadata-ozone" +version = "1.1.0" +description = "Provider package apache-airflow-providers-arenadata-ozone for Apache Airflow" +readme = "README.rst" +license = "Apache-2.0" +license-files = ['LICENSE', 'NOTICE'] +authors = [ + {name="Apache Software Foundation", email="dev@airflow.apache.org"}, +] +maintainers = [ + {name="Apache Software Foundation", email="dev@airflow.apache.org"}, +] +keywords = [ "airflow-provider", "arenadata.ozone", "airflow", "integration" ] +classifiers = [ + "Development Status :: 5 - Production/Stable", + "Environment :: Console", + "Environment :: Web Environment", + "Intended Audience :: Developers", + "Intended Audience :: System Administrators", + "Framework :: Apache Airflow", + "Framework :: Apache Airflow :: Provider", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: System :: Monitoring", +] +requires-python = ">=3.10" + +# The dependencies should be modified in place in the generated file. +# Any change in the dependencies is preserved when the file is regenerated +# Make sure to run ``prek update-providers-dependencies --all-files`` +# After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` +dependencies = [ + "apache-airflow>=2.11.0", + "apache-airflow-providers-common-compat>=1.12.0", +] + +[dependency-groups] +dev = [ + "apache-airflow", + "apache-airflow-task-sdk", + "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", + # Additional devel dependencies (do not remove this line and add extra development dependencies) +] + +# To build docs: +# +# uv run --group docs build-docs +# +# To enable auto-refreshing build with server: +# +# uv run --group docs build-docs --autobuild +# +# To see more options: +# +# uv run --group docs build-docs --help +# +docs = [ + "apache-airflow-devel-common[docs]" +] + +[tool.uv.sources] +# These names must match the names as defined in the pyproject.toml of the workspace items, +# *not* the workspace folder paths +apache-airflow = {workspace = true} +apache-airflow-devel-common = {workspace = true} +apache-airflow-task-sdk = {workspace = true} +apache-airflow-providers-common-sql = {workspace = true} +apache-airflow-providers-standard = {workspace = true} + +[project.urls] +"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-arenadata-ozone/1.1.0" +"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-arenadata-ozone/1.1.0/changelog.html" +"Bug Tracker" = "https://github.com/apache/airflow/issues" +"Source Code" = "https://github.com/apache/airflow" +"Slack Chat" = "https://s.apache.org/airflow-slack" +"Mastodon" = "https://fosstodon.org/@airflow" +"YouTube" = "https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/" + +[project.entry-points."apache_airflow_provider"] +provider_info = "airflow.providers.arenadata.ozone.get_provider_info:get_provider_info" + +[tool.flit.module] +name = "airflow.providers.arenadata.ozone" diff --git a/providers/arenadata/ozone/src/airflow/__init__.py b/providers/arenadata/ozone/src/airflow/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/__init__.py @@ -0,0 +1,17 @@ +# 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. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/arenadata/ozone/src/airflow/providers/__init__.py b/providers/arenadata/ozone/src/airflow/providers/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/__init__.py @@ -0,0 +1,17 @@ +# 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. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/__init__.py @@ -0,0 +1,17 @@ +# 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. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/__init__.py new file mode 100644 index 0000000000000..92e582b1b452d --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/__init__.py @@ -0,0 +1,39 @@ +# 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. +# +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE +# OVERWRITTEN WHEN PREPARING DOCUMENTATION FOR THE PACKAGES. +# +# IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE +# `PROVIDER__INIT__PY_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY +# +from __future__ import annotations + +import packaging.version + +from airflow import __version__ as airflow_version + +__all__ = ["__version__"] + +__version__ = "1.1.0" + +if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse( + "2.11.0" +): + raise RuntimeError( + f"The package `apache-airflow-providers-arenadata-ozone:{__version__}` needs Apache Airflow 2.11.0+" + ) diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_copy_from_hdfs.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_copy_from_hdfs.py new file mode 100644 index 0000000000000..d2bea754ab3a4 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_copy_from_hdfs.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python +# 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. + +""" +Ozone Copy From HDFS Example DAG. + +This DAG demonstrates a small HDFS -> Ozone copy workflow: +1. Waits for an Ozone trigger file to appear (using OzoneKeySensor). +2. Creates the same trigger file after a short delay for standalone demos. +3. Sets storage quota for the destination bucket. +4. Copies data from HDFS to Ozone using HdfsToOzoneOperator (DistCp). + +This example is import-safe even when the HDFS provider is not installed. +HdfsToOzoneOperator relies on Hadoop DistCp runtime, and runtime validation +happens only when task `copy_hdfs_to_ozone` executes. + +Runtime values can be overridden from the Trigger UI or via `dag_run.conf`. +The `OZONE_EXAMPLE_COPY_FROM_HDFS_*` environment variables are used as +process-level defaults. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from pathlib import PurePosixPath + +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneSetQuotaOperator, + OzoneUploadContentOperator, +) +from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor +from airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone import HdfsToOzoneOperator +from airflow.providers.arenadata.ozone.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.providers.standard.operators.bash import BashOperator + from airflow.sdk import DAG, timezone + from airflow.sdk.definitions.param import Param +else: + from airflow import DAG + from airflow.models.param import Param + from airflow.operators.bash import BashOperator + from airflow.utils import timezone + +DEFAULT_OM_HOST = os.getenv("OZONE_EXAMPLE_OM_HOST") or "om" +DEFAULT_CONN_ID = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_CONN_ID") or "ozone_admin_default" +DEFAULT_HDFS_CONN_ID = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_HDFS_CONN_ID") or "hdfs_admin_default" +DEFAULT_VOLUME = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_VOLUME") or "vol1" +DEFAULT_BUCKET = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_BUCKET") or "bucket1" +DEFAULT_TRIGGER_FILE = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_TRIGGER_FILE") or "trigger.lck" +DEFAULT_QUOTA = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_QUOTA") or "5GB" +DEFAULT_SOURCE_PATH = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_SOURCE_PATH") or "hdfs:///user/data/legacy/" +DEFAULT_DEST_SUBPATH = os.getenv("OZONE_EXAMPLE_COPY_FROM_HDFS_DEST_SUBPATH") or "migrated/" + +with DAG( + "example_ozone_copy_from_hdfs", + start_date=timezone.datetime(2025, 1, 1), + schedule=None, + catchup=False, + tags=["ozone", "hdfs", "example"], + params={ + "om_host": Param(DEFAULT_OM_HOST, type="string", title="OM host / Ozone authority"), + "pipeline_conn_id": Param(DEFAULT_CONN_ID, type="string", title="Ozone connection ID"), + "hdfs_conn_id": Param(DEFAULT_HDFS_CONN_ID, type="string", title="HDFS connection ID"), + "volume": Param(DEFAULT_VOLUME, type="string", title="Destination volume"), + "bucket": Param(DEFAULT_BUCKET, type="string", title="Destination bucket"), + "trigger_file": Param(DEFAULT_TRIGGER_FILE, type="string", title="Trigger file name"), + "quota": Param(DEFAULT_QUOTA, type="string", title="Destination bucket quota"), + "source_path": Param(DEFAULT_SOURCE_PATH, type="string", title="HDFS source path"), + "dest_subpath": Param(DEFAULT_DEST_SUBPATH, type="string", title="Destination subpath in bucket"), + }, +) as dag: + PIPELINE_TRIGGER_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.volume }}', '{{ params.bucket }}', '{{ params.trigger_file }}')}" + ) + PIPELINE_DEST_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.volume }}', '{{ params.bucket }}', '{{ params.dest_subpath }}')}" + ) + + delay_trigger_file_creation = BashOperator( + task_id="delay_trigger_file_creation", + bash_command="sleep 10", + execution_timeout=timedelta(seconds=30), + ) + + create_trigger_file = OzoneUploadContentOperator( + task_id="create_trigger_file", + content="HDFS to Ozone trigger generated by Airflow at {{ ts }}\n", + remote_path=PIPELINE_TRIGGER_PATH, + ozone_conn_id="{{ params.pipeline_conn_id }}", + retry_attempts=1, + timeout=60, + if_exists="ignore", + execution_timeout=timedelta(minutes=2), + ) + + check_trigger = OzoneKeySensor( + task_id="wait_for_trigger_file", + path=PIPELINE_TRIGGER_PATH, + ozone_conn_id="{{ params.pipeline_conn_id }}", + mode="reschedule", + execution_timeout=timedelta(minutes=1), + ) + + set_bucket_quota = OzoneSetQuotaOperator( + task_id="prepare_bucket_quota", + volume="{{ params.volume }}", + bucket="{{ params.bucket }}", + quota="{{ params.quota }}", + ozone_conn_id="{{ params.pipeline_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + copy_hdfs_to_ozone = HdfsToOzoneOperator( + task_id="copy_hdfs_to_ozone", + source_path="{{ params.source_path }}", + dest_path=PIPELINE_DEST_PATH, + hdfs_conn_id="{{ params.hdfs_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + delay_trigger_file_creation >> create_trigger_file + [check_trigger, create_trigger_file] >> set_bucket_quota >> copy_hdfs_to_ozone diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_data_lifecycle.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_data_lifecycle.py new file mode 100644 index 0000000000000..7f33ee10b3320 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_data_lifecycle.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python +# 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. + +""" +Complete Data Lifecycle Management Example DAG + +This DAG demonstrates a complete data lifecycle workflow: +1. Lists all files in a landing directory (OzoneListOperator) +2. Runs a processing step over the discovered list of files +3. Archives processed files to a new location (OzoneMoveOperator) +4. Creates a disaster-recovery snapshot (OzoneBackupOperator) +5. Cleans up original files from the landing zone + +This example showcases: +- Data archiving and lifecycle management +- Backup and disaster recovery workflows +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from pathlib import PurePosixPath + +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneDeleteKeyOperator, + OzoneListOperator, + OzoneMoveOperator, +) +from airflow.providers.arenadata.ozone.transfers.ozone_backup import OzoneBackupOperator +from airflow.providers.arenadata.ozone.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.providers.standard.operators.bash import BashOperator + from airflow.sdk import DAG, TaskGroup, timezone + from airflow.sdk.definitions.param import Param +else: + from airflow import DAG + from airflow.models.param import Param + from airflow.operators.bash import BashOperator + from airflow.utils import timezone + from airflow.utils.task_group import TaskGroup + +DEFAULT_OM_HOST = os.getenv("OZONE_EXAMPLE_OM_HOST") or "om" +DEFAULT_CONN_ID = os.getenv("OZONE_EXAMPLE_LIFECYCLE_CONN_ID") or "ozone_admin_default" +DEFAULT_LANDING_VOLUME = os.getenv("OZONE_EXAMPLE_LIFECYCLE_LANDING_VOLUME") or "landing" +DEFAULT_LANDING_BUCKET = os.getenv("OZONE_EXAMPLE_LIFECYCLE_LANDING_BUCKET") or "raw" +DEFAULT_ARCHIVE_VOLUME = os.getenv("OZONE_EXAMPLE_LIFECYCLE_ARCHIVE_VOLUME") or "archive" +DEFAULT_ARCHIVE_BUCKET = os.getenv("OZONE_EXAMPLE_LIFECYCLE_ARCHIVE_BUCKET") or "processed" +DEFAULT_SNAPSHOT_PREFIX = os.getenv("OZONE_EXAMPLE_LIFECYCLE_SNAPSHOT_PREFIX") or "snap" + +with DAG( + dag_id="example_ozone_data_lifecycle", + start_date=timezone.datetime(2025, 1, 1), + catchup=False, + schedule=None, + tags=["ozone", "example"], + params={ + "om_host": Param(DEFAULT_OM_HOST, type="string", title="OM host / Ozone authority"), + "lifecycle_conn_id": Param(DEFAULT_CONN_ID, type="string", title="Ozone connection ID"), + "landing_volume": Param(DEFAULT_LANDING_VOLUME, type="string", title="Landing volume"), + "landing_bucket": Param(DEFAULT_LANDING_BUCKET, type="string", title="Landing bucket"), + "archive_volume": Param(DEFAULT_ARCHIVE_VOLUME, type="string", title="Archive volume"), + "archive_bucket": Param(DEFAULT_ARCHIVE_BUCKET, type="string", title="Archive bucket"), + "snapshot_prefix": Param( + DEFAULT_SNAPSHOT_PREFIX, + type="string", + title="Snapshot name prefix", + description="Run date suffix is added automatically as ds_nodash.", + ), + }, + doc_md=""" + ### Data Lifecycle Example (with Hive Registration) + + This DAG demonstrates a complete data lifecycle management workflow: + 1. **List**: Finds all files in a landing directory using `OzoneListOperator`. + 2. **Process**: Runs a processing step over the file list from XCom. + 3. **Archive**: Moves the processed source files to an archive path using `OzoneMoveOperator`. + 4. **Backup**: Creates a disaster-recovery snapshot of the bucket using `OzoneBackupOperator`. + 5. **Cleanup**: Deletes the original files from the landing directory. + + Runtime values can be overridden from the Trigger UI or via `dag_run.conf`. + """, +) as dag: + LANDING_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.landing_volume }}', '{{ params.landing_bucket }}')}" + ) + ARCHIVE_BASE_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.archive_volume }}', '{{ params.archive_bucket }}')}" + ) + + # 0. Ensure required volumes and buckets exist. + create_landing_volume = OzoneCreateVolumeOperator( + task_id="create_landing_volume", + volume_name="{{ params.landing_volume }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + create_archive_volume = OzoneCreateVolumeOperator( + task_id="create_archive_volume", + volume_name="{{ params.archive_volume }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + create_landing_bucket = OzoneCreateBucketOperator( + task_id="create_landing_bucket", + volume_name="{{ params.landing_volume }}", + bucket_name="{{ params.landing_bucket }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + create_archive_bucket = OzoneCreateBucketOperator( + task_id="create_archive_bucket", + volume_name="{{ params.archive_volume }}", + bucket_name="{{ params.archive_bucket }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 1. List all files in the landing directory. The result is pushed to XComs. + list_files_in_landing_zone = OzoneListOperator( + task_id="list_files_in_landing_zone", + path=LANDING_PATH, + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 2. This TaskGroup emulates file processing in a single task. + with TaskGroup(group_id="dynamic_file_processing") as processing_group: + process_files = BashOperator( + task_id="process_files", + bash_command=( + 'echo "Processing files: {{ ti.xcom_pull(task_ids="list_files_in_landing_zone") }}"' + ), + execution_timeout=timedelta(minutes=1), + ) + + # 3. Create archive directory for date-based partitioning + create_archive_dir = OzoneCreatePathOperator( + task_id="create_archive_dir", + path=ARCHIVE_BASE_PATH + "/ds={{ ds }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 4. After processing, move the original files to an archive directory. + archive_files = OzoneMoveOperator( + task_id="archive_landing_files", + source_path=f"{LANDING_PATH}/*", # Move all files + dest_path=ARCHIVE_BASE_PATH + "/ds={{ ds }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + # 5. Create a snapshot of the entire archive volume for backup. + backup_archive = OzoneBackupOperator( + task_id="backup_archive_via_snapshot", + volume="{{ params.archive_volume }}", + bucket="{{ params.archive_bucket }}", + snapshot_name="{{ params.snapshot_prefix }}-{{ ds_nodash }}", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + # 6. Cleanup original files from landing zone. + cleanup_landing_zone = OzoneDeleteKeyOperator( + task_id="cleanup_landing_zone", + path=f"{LANDING_PATH}/*", + ozone_conn_id="{{ params.lifecycle_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + create_landing_volume >> create_landing_bucket + create_archive_volume >> create_archive_bucket + + create_landing_bucket >> list_files_in_landing_zone + + list_files_in_landing_zone >> processing_group + create_archive_bucket >> create_archive_dir + processing_group >> create_archive_dir >> archive_files + archive_files >> backup_archive + backup_archive >> cleanup_landing_zone diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_multi_tenant_management.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_multi_tenant_management.py new file mode 100644 index 0000000000000..a6a9a4840e503 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_multi_tenant_management.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python +# 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. + +""" +Multi-Tenant Management Example DAG + +This DAG demonstrates how to automate the setup of storage resources for new projects/tenants: +1. Creates a dedicated volume for a new project +2. Sets quota limits to prevent resource overuse +3. Creates standard directory structure (landing, processed, etc.) + +This example shows how Airflow can be used for administrative tasks, +automating the provisioning of storage resources for new teams or projects. +Useful for multi-tenant environments where each project needs isolated storage. + +Runtime values can be overridden from the Trigger UI or via `dag_run.conf`; +legacy `OZONE_EXAMPLE_*` environment variables are kept as defaults for +compatibility. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from pathlib import PurePosixPath + +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneSetQuotaOperator, +) +from airflow.providers.arenadata.ozone.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import DAG, timezone + from airflow.sdk.definitions.param import Param +else: + from airflow import DAG + from airflow.models.param import Param + from airflow.utils import timezone + +DEFAULT_OM_HOST = os.getenv("OZONE_EXAMPLE_OM_HOST") or "om" +DEFAULT_CONN_ID = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_CONN_ID") or "ozone_admin_default" +DEFAULT_PROJECT_VOLUME = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_PROJECT_VOLUME") or "project-alpha" +DEFAULT_PROJECT_QUOTA = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_PROJECT_QUOTA") or "10GB" +DEFAULT_LANDING_BUCKET = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_LANDING_BUCKET") or "landing" +DEFAULT_PROCESSED_BUCKET = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_PROCESSED_BUCKET") or "processed" +DEFAULT_BUCKET_QUOTA = os.getenv("OZONE_EXAMPLE_MULTI_TENANT_BUCKET_QUOTA") or "1GB" + +with DAG( + dag_id="example_ozone_multi_tenant_management", + start_date=timezone.datetime(2025, 1, 1), + catchup=False, + schedule=None, + tags=["ozone", "example"], + params={ + "om_host": Param(DEFAULT_OM_HOST, type="string", title="OM host / Ozone authority"), + "multi_tenant_conn_id": Param(DEFAULT_CONN_ID, type="string", title="Ozone connection ID"), + "project_volume": Param(DEFAULT_PROJECT_VOLUME, type="string", title="Project volume"), + "project_quota": Param(DEFAULT_PROJECT_QUOTA, type="string", title="Project quota"), + "landing_bucket": Param(DEFAULT_LANDING_BUCKET, type="string", title="Landing bucket"), + "processed_bucket": Param(DEFAULT_PROCESSED_BUCKET, type="string", title="Processed bucket"), + "bucket_quota": Param(DEFAULT_BUCKET_QUOTA, type="string", title="Bucket quota"), + }, +) as dag: + LANDING_DATA_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.project_volume }}', '{{ params.landing_bucket }}', 'data')}" + ) + PROCESSED_DATA_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.project_volume }}', '{{ params.processed_bucket }}', 'data')}" + ) + + # 1. Create a dedicated volume for the new project + create_volume = OzoneCreateVolumeOperator( + task_id="create_project_volume", + volume_name="{{ params.project_volume }}", + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 2. Set a hard limit on the volume to prevent overuse + set_quota = OzoneSetQuotaOperator( + task_id="set_project_quota", + volume="{{ params.project_volume }}", + quota="{{ params.project_quota }}", + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 3. Create buckets with quotas (required when volume has quota) + create_landing_bucket = OzoneCreateBucketOperator( + task_id="create_landing_bucket", + volume_name="{{ params.project_volume }}", + bucket_name="{{ params.landing_bucket }}", + quota="{{ params.bucket_quota }}", + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + create_processed_bucket = OzoneCreateBucketOperator( + task_id="create_processed_bucket", + volume_name="{{ params.project_volume }}", + bucket_name="{{ params.processed_bucket }}", + quota="{{ params.bucket_quota }}", + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + # 4. Create standard subdirectories inside the buckets + create_landing_dir = OzoneCreatePathOperator( + task_id="create_landing_dir", + path=LANDING_DATA_PATH, + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + create_processed_dir = OzoneCreatePathOperator( + task_id="create_processed_dir", + path=PROCESSED_DATA_PATH, + ozone_conn_id="{{ params.multi_tenant_conn_id }}", + execution_timeout=timedelta(minutes=1), + ) + + create_volume >> set_quota + set_quota >> [create_landing_bucket, create_processed_bucket] + create_landing_bucket >> create_landing_dir + create_processed_bucket >> create_processed_dir diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_usage.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_usage.py new file mode 100644 index 0000000000000..51fa39e0950d1 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/example_dags/example_ozone_usage.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python +# 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. + +""" +Basic Ozone Usage Example DAG + +This DAG demonstrates fundamental native Ozone operations: +- Creates a volume and bucket via Native Admin CLI +- Creates directories and uploads files via Ozone Filesystem (ofs://) +- Uses a sensor to wait for a file to appear + +The same DAG works for plain, SSL, Kerberos, and SSL+Kerberos modes. +The active mode is defined by the selected Airflow connection and its +``Connection Extra``. Runtime values can be overridden from the Trigger UI or +via ``dag_run.conf``; legacy ``OZONE_EXAMPLE_*`` environment variables are kept +as defaults for compatibility. +""" + +from __future__ import annotations + +import os +from datetime import timedelta +from pathlib import PurePosixPath + +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneAdminHook +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneUploadContentOperator, +) +from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.providers.standard.operators.python import PythonOperator + from airflow.sdk import DAG, timezone + from airflow.sdk.definitions.param import Param +else: + from airflow import DAG + from airflow.models.param import Param + from airflow.operators.python import PythonOperator + from airflow.utils import timezone + +DEFAULT_OM_HOST = os.getenv("OZONE_EXAMPLE_OM_HOST") or "om" +DEFAULT_VOLUME = os.getenv("OZONE_EXAMPLE_USAGE_VOLUME") or "vol1" +DEFAULT_BUCKET = os.getenv("OZONE_EXAMPLE_USAGE_BUCKET") or "bucket-native" +DEFAULT_DIR = os.getenv("OZONE_EXAMPLE_USAGE_DIR") or "data_dir" +DEFAULT_FILE = os.getenv("OZONE_EXAMPLE_USAGE_FILE") or "file.txt" +DEFAULT_CONN_ID = os.getenv("OZONE_EXAMPLE_USAGE_ADMIN_CONN_ID") or "ozone_admin_default" +DEFAULT_VOLUME_QUOTA = os.getenv("OZONE_EXAMPLE_USAGE_VOLUME_QUOTA") or "10GB" +DEFAULT_BUCKET_QUOTA = os.getenv("OZONE_EXAMPLE_USAGE_BUCKET_QUOTA") or "10GB" + + +def _check_ozone_auth(**context) -> None: + conn_id = context["params"]["admin_conn_id"] + hook = OzoneAdminHook(ozone_conn_id=conn_id) + ok, message = hook.test_connection() + if not ok: + raise OzoneProviderError(message) + + +default_args = { + "owner": "airflow", + "retries": 1, + "retry_delay": timedelta(minutes=1), +} + +with DAG( + "example_ozone_usage", + start_date=timezone.datetime(2024, 1, 1), + default_args=default_args, + schedule=None, + catchup=False, + tags=["ozone", "example"], + params={ + "om_host": Param(DEFAULT_OM_HOST, type="string", title="OM host / Ozone authority"), + "admin_conn_id": Param(DEFAULT_CONN_ID, type="string", title="Ozone admin connection ID"), + "volume": Param(DEFAULT_VOLUME, type="string", title="Volume name"), + "bucket": Param(DEFAULT_BUCKET, type="string", title="Bucket name"), + "directory": Param(DEFAULT_DIR, type="string", title="Directory path inside bucket"), + "file_name": Param(DEFAULT_FILE, type="string", title="File name"), + "volume_quota": Param( + DEFAULT_VOLUME_QUOTA, + type="string", + title="Volume quota", + description="Quota string accepted by Ozone CLI, for example 10GB.", + ), + "bucket_quota": Param( + DEFAULT_BUCKET_QUOTA, + type="string", + title="Bucket quota", + description="Quota string accepted by Ozone CLI, for example 10GB.", + ), + }, +) as dag: + FS_DIR_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.volume }}', '{{ params.bucket }}', '{{ params.directory }}')}" + ) + FS_FILE_PATH = ( + f"ofs://{{{{ params.om_host }}}}/" + f"{PurePosixPath('{{ params.volume }}', '{{ params.bucket }}', '{{ params.directory }}', '{{ params.file_name }}')}" + ) + + check_ozone_auth = PythonOperator( + task_id="check_ozone_auth", + python_callable=_check_ozone_auth, + execution_timeout=timedelta(seconds=30), + ) + + create_vol = OzoneCreateVolumeOperator( + task_id="create_volume", + volume_name="{{ params.volume }}", + quota="{{ params.volume_quota }}", + ozone_conn_id="{{ params.admin_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + create_bucket_native = OzoneCreateBucketOperator( + task_id="create_bucket_native", + volume_name="{{ params.volume }}", + bucket_name="{{ params.bucket }}", + quota="{{ params.bucket_quota }}", + ozone_conn_id="{{ params.admin_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + fs_mkdir = OzoneCreatePathOperator( + task_id="fs_mkdir", + path=FS_DIR_PATH, + ozone_conn_id="{{ params.admin_conn_id }}", + execution_timeout=timedelta(minutes=5), + ) + + fs_put = OzoneUploadContentOperator( + task_id="fs_put_file", + content="Hello from FS Layer", + remote_path=FS_FILE_PATH, + ozone_conn_id="{{ params.admin_conn_id }}", + if_exists="overwrite", + execution_timeout=timedelta(minutes=5), + ) + + wait_fs_file = OzoneKeySensor( + task_id="wait_fs_file", + path=FS_FILE_PATH, + ozone_conn_id="{{ params.admin_conn_id }}", + mode="reschedule", + timeout=60, + ) + + check_ozone_auth >> create_vol >> create_bucket_native >> fs_mkdir >> fs_put >> wait_fs_file diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/get_provider_info.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/get_provider_info.py new file mode 100644 index 0000000000000..b75fb08531617 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/get_provider_info.py @@ -0,0 +1,72 @@ +# 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. + +# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN! +# +# IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE +# `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY + + +def get_provider_info(): + return { + "package-name": "apache-airflow-providers-arenadata-ozone", + "name": "Apache Ozone", + "description": "`Apache Ozone `__\n\nProvider package for interacting with Apache Ozone via:\n\n- Native CLI (``ozone``) for ``ofs://`` / ``o3fs://`` paths\n- Native Ozone administration and filesystem workflows\n", + "integrations": [ + { + "integration-name": "Apache Ozone", + "external-doc-url": "https://ozone.apache.org/", + "tags": ["apache"], + } + ], + "hooks": [ + { + "integration-name": "Apache Ozone", + "python-modules": ["airflow.providers.arenadata.ozone.hooks.ozone"], + } + ], + "operators": [ + { + "integration-name": "Apache Ozone", + "python-modules": ["airflow.providers.arenadata.ozone.operators.ozone"], + } + ], + "sensors": [ + { + "integration-name": "Apache Ozone", + "python-modules": ["airflow.providers.arenadata.ozone.sensors.ozone"], + } + ], + "transfers": [ + { + "source-integration-name": "Hadoop Distributed File System (HDFS)", + "target-integration-name": "Apache Ozone", + "python-module": "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone", + }, + { + "source-integration-name": "Apache Ozone", + "target-integration-name": "Apache Ozone", + "python-module": "airflow.providers.arenadata.ozone.transfers.ozone_backup", + }, + ], + "connection-types": [ + { + "hook-class-name": "airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliHook", + "connection-type": "ozone", + } + ], + } diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/__init__.py new file mode 100644 index 0000000000000..633f9deb60d75 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/__init__.py @@ -0,0 +1,37 @@ +# +# 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. + +"""Ozone hooks.""" + +from __future__ import annotations + +from airflow.providers.arenadata.ozone.hooks.ozone import ( + OzoneAdminHook, + OzoneCliHook, + OzoneFsHook, + OzoneResource, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError + +__all__ = [ + "OzoneCliHook", + "OzoneFsHook", + "OzoneAdminHook", + "OzoneResource", + "OzoneCliError", +] diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/ozone.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/ozone.py new file mode 100644 index 0000000000000..30679f2989c28 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/hooks/ozone.py @@ -0,0 +1,1692 @@ +# 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 + +import os +import shlex +import subprocess +from enum import Enum +from functools import cached_property +from pathlib import Path + +from airflow.providers.arenadata.ozone.utils.cli_runner import ( + CliRunner, + OzoneCliRunner, + ProcessOutputAnalysis, +) +from airflow.providers.arenadata.ozone.utils.connection_schema import ( + OZONE_CONNECTION_UI_FIELD_BEHAVIOUR, + OzoneConnSnapshot, +) +from airflow.providers.arenadata.ozone.utils.errors import ( + ADMIN_RESOURCE_SPECS, + OzoneCliError, + OzoneCliErrors, + OzoneProviderError, +) +from airflow.providers.arenadata.ozone.utils.helpers import ( + FileHelper, + JsonDict, + JsonValue, + TypeNormalizationHelper, + URIHelper, +) +from airflow.providers.arenadata.ozone.utils.security import KerberosConfig, SSLConfig +from airflow.providers.arenadata.ozone.version_compat import AirflowException, BaseHook, redact + +RETRY_ATTEMPTS = 3 +FAST_TIMEOUT_SECONDS = 5 * 60 +SLOW_TIMEOUT_SECONDS = 60 * 60 + + +class ExistingTargetPolicy(str, Enum): + """Policy for operations that write to an existing Ozone target.""" + + ERROR = "error" + IGNORE = "ignore" + OVERWRITE = "overwrite" + + +class OzoneCliHook(BaseHook): + """Base hook for Ozone CLI commands with retry and auth handling.""" + + conn_name_attr = "ozone_conn_id" + default_conn_name = "ozone_default" + conn_type = "ozone" + hook_name = "Ozone" + + def __init__( + self, + ozone_conn_id: str = default_conn_name, + *, + retry_attempts: int = RETRY_ATTEMPTS, + ) -> None: + super().__init__() + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self._kerberos_ticket_ready = False + + @cached_property + def connection_snapshot(self) -> OzoneConnSnapshot: + """Validate and cache required host, port and extra connection fields.""" + return OzoneConnSnapshot.from_connection( + self.get_connection(self.ozone_conn_id), + conn_id=self.ozone_conn_id, + require_host_port=True, + ) + + @cached_property + def _cached_ssl_env(self) -> dict[str, str] | None: + """SSL/TLS configuration from snapshot.""" + try: + ssl_env_vars = SSLConfig.from_snapshot( + self.connection_snapshot, + conn_id=self.ozone_conn_id, + scope="ozone", + ).as_env() + if not ssl_env_vars: + self.log.debug("No SSL/TLS configuration found in connection snapshot") + return None + ssl_env = SSLConfig.apply_ssl_env_vars(ssl_env_vars) + self.log.debug("SSL/TLS configuration loaded from connection: %s", list(ssl_env_vars.keys())) + if self.connection_snapshot.ozone_security_enabled: + self.log.info("SSL/TLS enabled for connection") + return ssl_env + except AirflowException as err: + self.log.debug("Could not load SSL configuration: %s", str(err)) + return None + + @cached_property + def _cached_kerberos_env(self) -> dict[str, str] | None: + """Kerberos configuration from connection Extra.""" + return KerberosConfig.load_ozone_env( + snapshot=self.connection_snapshot, + conn_id=self.ozone_conn_id, + ) + + @cached_property + def _cached_effective_config_dir(self) -> str | None: + """Config dir used for CLI --config and subprocess environment.""" + if self.connection_snapshot.ozone_conf_dir: + self.log.info( + "Using ozone_conf_dir from connection extra: %s", self.connection_snapshot.ozone_conf_dir + ) + return self.connection_snapshot.ozone_conf_dir + + for env_name in ("OZONE_CONF_DIR", "HADOOP_CONF_DIR"): + env_value = os.environ.get(env_name) + if env_value: + self.log.info("Using ozone_conf_dir from %s environment variable: %s", env_name, env_value) + return env_value + + if self.connection_snapshot.kerberos_enabled: + raise OzoneProviderError( + "Kerberos is enabled but ozone_conf_dir is not configured. " + "Set ozone_conf_dir in the connection extra, or provide OZONE_CONF_DIR/HADOOP_CONF_DIR " + "in the worker environment." + ) + return None + + def _prepared_cli_env(self) -> dict[str, str]: + """Build subprocess environment for Ozone CLI calls.""" + env = CliRunner.merge_env(None) + env["OZONE_OM_ADDRESS"] = f"{self.connection_snapshot.host}:{self.connection_snapshot.port}" + + if self._cached_ssl_env: + env.update(self._cached_ssl_env) + if self._cached_kerberos_env: + env.update(self._cached_kerberos_env) + + config_dir = self._cached_effective_config_dir + if config_dir: + env["OZONE_CONF_DIR"] = config_dir + env["HADOOP_CONF_DIR"] = config_dir + + self.log.debug("Prepared Ozone CLI env for connection '%s'", self.ozone_conn_id) + return env + + @classmethod + def get_ui_field_behaviour(cls) -> JsonDict: + """Describe Ozone connection extras in Airflow UI.""" + return OZONE_CONNECTION_UI_FIELD_BEHAVIOUR + + def test_connection(self) -> tuple[bool, str]: + """Run a minimal CLI command to verify Ozone auth and connectivity.""" + try: + result = self.run_cli( + ["ozone", "sh", "volume", "list", "/"], + timeout=FAST_TIMEOUT_SECONDS, + retry_attempts=0, + check=False, + log_output=False, + return_result=True, + ) + except OzoneCliError as err: + return False, f"Ozone CLI connection test failed: {err}" + + if result.returncode == 0: + return True, "Ozone CLI connection test succeeded." + + error_text = ProcessOutputAnalysis.pick_process_output(result) or "Unknown CLI error" + return False, f"Ozone CLI connection test failed: {error_text}" + + def _prepare_cli_command(self, cmd: list[str]) -> list[str]: + """Add --config for Kerberos-enabled commands when available.""" + if not self.connection_snapshot.kerberos_enabled: + return cmd + + if "--config" in cmd: + return cmd + + config_dir = self._cached_effective_config_dir + config_files_exist = KerberosConfig.check_config_files_exist( + config_dir, + snapshot=self.connection_snapshot, + ) + if not config_files_exist: + self.log.warning( + "Kerberos enabled but configuration files (core-site.xml, ozone-site.xml) " + "not found in %s. Adding --config flag anyway as it is critical for Kerberos authentication.", + config_dir, + ) + else: + self.log.debug("Configuration files found in %s, adding --config flag", config_dir) + + if cmd[:1] == ["ozone"]: + new_cmd = [cmd[0], "--config", config_dir, *cmd[1:]] + self.log.debug("Added --config flag to Ozone CLI command: %s", redact(shlex.join(new_cmd))) + return new_cmd + + return cmd + + def runtime_mode_label(self) -> str: + """Return connection runtime mode label for operator/hook logs.""" + ssl_enabled = self.connection_snapshot.ssl_enabled + kerberos_enabled = self.connection_snapshot.kerberos_enabled + + if ssl_enabled and kerberos_enabled: + return "ssl+kerberos" + if ssl_enabled: + return "ssl" + if kerberos_enabled: + return "kerberos" + return "plain" + + def run_cli( + self, + cmd: list[str], + *, + timeout: int = FAST_TIMEOUT_SECONDS, + retry_attempts: int | None = None, + input_text: str | None = None, + check: bool = True, + log_output: bool = True, + return_result: bool = False, + return_json_result: bool = False, + ) -> str | subprocess.CompletedProcess[str] | JsonValue: + """Execute Ozone CLI command with common auth, retry and logging handling.""" + if return_result and return_json_result: + raise ValueError("return_result and return_json_result cannot be enabled together") + self.log.info( + "Executing Ozone CLI command (connection: %s, mode: %s)", + self.ozone_conn_id, + self.runtime_mode_label(), + ) + self._kerberos_ticket_ready = KerberosConfig.ensure_ticket( + snapshot=self.connection_snapshot, + conn_id=self.ozone_conn_id, + kerberos_ticket_ready=self._kerberos_ticket_ready, + ) + prepared_cmd = self._prepare_cli_command(cmd) + effective_retry_attempts = self.retry_attempts if retry_attempts is None else retry_attempts + result = OzoneCliRunner.run_ozone( + prepared_cmd, + env_overrides=self._prepared_cli_env(), + timeout=timeout, + input_text=input_text, + check=check, + log_output=log_output, + retry_attempts=effective_retry_attempts, + ) + if return_result: + return result + output = ProcessOutputAnalysis.pick_process_output(result) + if return_json_result: + return TypeNormalizationHelper.parse_json_output(output) + return output + + +class OzoneFsHook(OzoneCliHook): + """Interact with Ozone through the Hadoop-compatible ozone fs interface.""" + + hook_name = "Ozone FS" + + @staticmethod + def _validate_if_exists( + value: ExistingTargetPolicy | str, + *, + allow_overwrite: bool = True, + ) -> ExistingTargetPolicy: + """Validate existing-target policy used by create/copy/upload helpers.""" + allowed: set[ExistingTargetPolicy] = {ExistingTargetPolicy.ERROR, ExistingTargetPolicy.IGNORE} + if allow_overwrite: + allowed.add(ExistingTargetPolicy.OVERWRITE) + try: + policy = ExistingTargetPolicy(value) + except ValueError: + expected = ", ".join(sorted(allowed)) + raise ValueError(f"if_exists must be one of: {expected}. Got: {value!r}") + if policy not in allowed: + expected = ", ".join(sorted(allowed)) + raise ValueError(f"if_exists must be one of: {expected}. Got: {value!r}") + return policy + + # ============================== + # Key operations + # ============================== + def create_key( + self, + path: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + ) -> None: + """Create an empty key in Ozone FS.""" + policy = self._validate_if_exists(if_exists) + if self.exists(path, timeout=timeout): + if policy == ExistingTargetPolicy.IGNORE: + self.log.info("Ozone key %s already exists, treating as success.", path) + return + if policy == ExistingTargetPolicy.OVERWRITE: + self.delete_key(path, timeout=timeout) + else: + raise OzoneProviderError(f"Ozone key already exists: {path}") + self.run_cli(self._fs_cmd("-touchz", path), timeout=timeout) + + def key_exists(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> bool: + """Return True when a key exists in Ozone FS.""" + return self.exists(path, timeout=timeout) + + def list_keys(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[str]: + """List keys for a plain path or wildcard pattern.""" + if URIHelper.contains_wildcards(path): + matched = self.list_wildcard_matches(path, timeout=timeout, fallback_action="listing") + if matched is None: + return self.list_paths(path, timeout=timeout) + return matched + return self.list_paths(path, timeout=timeout) + + def glob(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[str]: + """Return wildcard matches for a key or path pattern.""" + return self.list_keys(path, timeout=timeout) + + def get_key_info(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> JsonDict: + """Return full metadata for a key using ozone sh key info.""" + parsed = self.run_cli( + ["ozone", "sh", "key", "info", URIHelper.to_key_uri(path)], + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from key info {path}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def get_key_property(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> JsonDict: + """Return the most useful key properties as a compact dictionary.""" + info = self.get_key_info(path, timeout=timeout) + replication_config = info.get("replicationConfig") + if replication_config is None: + self.log.warning( + "Key info for %s does not contain 'replicationConfig'; falling back to top-level fields", + path, + ) + if isinstance(replication_config, dict): + replication_type = replication_config.get("replicationType", info.get("replicationType")) + replication_value = replication_config.get("requiredNodes", info.get("replicationFactor")) + else: + replication_type = info.get("replicationType") + replication_value = info.get("replicationFactor") + + return { + "volume_name": info.get("volumeName"), + "bucket_name": info.get("bucketName"), + "name": info.get("name"), + "data_size": info.get("dataSize"), + "creation_time": info.get("creationTime"), + "modification_time": info.get("modificationTime"), + "replication_type": replication_type, + "replication": replication_value, + "metadata": info.get("metadata"), + "file_encryption_info": info.get("fileEncryptionInfo"), + } + + def read_text(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> str: + """Read a small text key from Ozone FS.""" + return self.run_cli(self._fs_cmd("-cat", path), timeout=timeout) + + def set_key_property( + self, + path: str, + replication_factor: int | None = None, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """ + Set writable key properties. + + Currently only replication factor is exposed because it is the only key-level + property already used in the provider codebase. + """ + if replication_factor is None: + self.log.debug("No key properties were specified to be set for path: %s", path) + return + + self.run_cli( + self._fs_cmd("-setrep", str(replication_factor), path), + timeout=timeout, + ) + + def delete_key(self, path: str, *, timeout: int = SLOW_TIMEOUT_SECONDS) -> None: + """Delete a key or a wildcard-matched group of keys.""" + self._delete_fs( + path, + recursive=False, + timeout=timeout, + fallback_action="direct delete", + ) + + # ============================== + # Path operations + # ============================== + def create_path( + self, + path: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + recursive: bool = True, + fail_if_exists: bool = False, + ) -> None: + """ + Create a directory tree in Ozone FS. + + Deprecated. Use `make_path(..., if_exists=...)` instead. + """ + self.log.warning( + "OzoneFsHook.create_path(..., fail_if_exists=...) is deprecated and will be removed " + "in a future version. Use make_path(..., if_exists=...) instead." + ) + policy = ExistingTargetPolicy.ERROR if fail_if_exists else ExistingTargetPolicy.IGNORE + self.make_path(path, timeout=timeout, recursive=recursive, if_exists=policy) + + def make_path( + self, + path: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + recursive: bool = True, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + ) -> None: + """Create a directory tree in Ozone FS.""" + policy = self._validate_if_exists(if_exists, allow_overwrite=False) + target_path = path.rstrip("/") + if not target_path: + return + if self.exists(target_path, timeout=timeout): + if policy == ExistingTargetPolicy.IGNORE: + self.log.info("Ozone path %s already exists, treating as success.", target_path) + return + raise OzoneProviderError(f"Destination path already exists: {target_path}") + + cmd = self._fs_cmd("-mkdir") + if recursive: + cmd.append("-p") + cmd.append(target_path) + + self.run_cli(cmd, timeout=timeout) + + def exists(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> bool: + """Return True when a file or directory exists in Ozone FS.""" + result = self.run_cli( + self._fs_cmd("-test", "-e", path), + timeout=timeout, + retry_attempts=0, + check=False, + log_output=False, + return_result=True, + ) + if result.returncode == 0: + return True + + output = ProcessOutputAnalysis.from_completed_process(result) + normalized = output.normalized_meaningful_output + if any(marker in normalized for marker in ("not found", "does not exist", "no such file")): + return False + if result.returncode == 1 and not output.meaningful_lines: + self.log.debug( + "Treating `ozone fs -test -e` return code 1 for %s as missing path because only noise output was produced.", + path, + ) + return False + + error_text = output.preferred_output or "Unknown CLI error" + raise OzoneCliError( + f"Ozone FS existence check failed (return code: {result.returncode}): {redact(error_text)}", + command=self._fs_cmd("-test", "-e", path), + stderr=result.stderr, + returncode=result.returncode, + retryable=OzoneCliErrors.is_retryable_failure(result.returncode, error_text), + ) + + def path_exists(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> bool: + """Return True when a path exists in Ozone FS.""" + return self.exists(path, timeout=timeout) + + def list_paths(self, path: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[str]: + """List file system paths under the provided Ozone URI.""" + output = self.run_cli(self._fs_cmd("-ls", "-C", path), timeout=timeout) + paths, skipped_lines = ProcessOutputAnalysis.extract_meaningful_output_lines( + output, + accept_line=lambda line: ( + URIHelper.parse_ozone_uri(line)[0] in {"ofs", "o3fs", "o3"} or line.startswith("/") + ), + ) + + if skipped_lines: + self.log.debug( + "Ignored %d non-path line(s) from `ozone fs -ls` output: %s", + len(skipped_lines), + skipped_lines[:3], + ) + + return paths + + def delete_path( + self, + path: str, + *, + recursive: bool = True, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> None: + """Delete a file or directory path, optionally recursively.""" + self._delete_fs( + path, + recursive=recursive, + timeout=timeout, + fallback_action="direct path delete", + ) + + # ============================== + # Transfer operations + # ============================== + def upload_key( + self, + local_path: str, + remote_path: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + ) -> None: + """Upload a local file to an Ozone URI.""" + policy = self._validate_if_exists(if_exists) + local_file = Path(local_path) + if not FileHelper.is_readable_file(local_file): + raise OzoneProviderError(f"Local file does not exist or is not readable: {local_path}") + + if self.exists(remote_path, timeout=timeout): + if policy == ExistingTargetPolicy.IGNORE: + self.log.info("Remote path %s already exists, treating as success.", remote_path) + return + if policy == ExistingTargetPolicy.ERROR: + raise OzoneProviderError(f"Remote path already exists: {remote_path}") + self.log.info("Remote path %s already exists, deleting it before upload.", remote_path) + self.delete_key(remote_path, timeout=timeout) + + cmd = ["ozone", "sh", "key", "put", URIHelper.to_key_uri(remote_path), str(local_file)] + self.run_cli(cmd, timeout=timeout) + + def download_key( + self, + remote_path: str, + local_path: str, + *, + overwrite: bool = False, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> None: + """Download one file from Ozone to the local filesystem.""" + local_file = Path(local_path) + if local_file.exists() and not overwrite: + raise OzoneProviderError( + f"Local file already exists: {local_path}. Set overwrite=True to replace it." + ) + local_file.parent.mkdir(parents=True, exist_ok=True) + if not FileHelper.is_writable_target(local_file): + raise OzoneProviderError(f"Local target path is not writable: {local_path}") + + cmd = self._fs_cmd("-get") + if overwrite: + cmd.append("-f") + cmd.extend([remote_path, str(local_file)]) + self.run_cli(cmd, timeout=timeout) + + def move( + self, + source_path: str, + dest_path: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + ) -> None: + """Move or rename one key or wildcard-selected keys within Ozone FS.""" + self._transfer_path( + "-mv", + source_path, + dest_path, + timeout=timeout, + fallback_action="direct move", + if_exists=if_exists, + ) + + def copy_path( + self, + source_path: str, + dest_path: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + ) -> None: + """Copy one key or wildcard-selected keys within Ozone FS.""" + self._transfer_path( + "-cp", + source_path, + dest_path, + timeout=timeout, + fallback_action="direct copy", + if_exists=if_exists, + ) + + # ============================== + # Shared FS helpers + # ============================== + def _fs_cmd(self, action: str, *args: str) -> list[str]: + """Build an ozone fs command.""" + return ["ozone", "fs", action, *args] + + def list_wildcard_matches( + self, + path: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + fallback_action: str = "direct operation", + ) -> list[str] | None: + """Resolve wildcard matches, or return None when listing should fall back.""" + try: + _, matched = URIHelper.resolve_wildcard_matches( + path, + lambda current_path: self.list_paths(current_path, timeout=timeout), + ) + except OzoneCliError as err: + if not err.retryable: + raise + source_dir, _ = URIHelper.split_ozone_wildcard_path(path) + self.log.warning( + "Could not list wildcard source directory %s (%s); falling back to %s for %s", + source_dir, + type(err).__name__, + fallback_action, + path, + ) + return None + return matched + + def _delete_fs( + self, + path: str, + *, + recursive: bool, + timeout: int, + fallback_action: str, + ) -> None: + """Shared delete implementation for key and path deletion.""" + rm_cmd = self._fs_cmd("-rm") + if recursive: + rm_cmd.append("-r") + + if not URIHelper.contains_wildcards(path): + if self.exists(path, timeout=timeout): + self.run_cli([*rm_cmd, path], timeout=timeout) + return + + matched = self.list_wildcard_matches(path, timeout=timeout, fallback_action=fallback_action) + if matched is None: + self.run_cli([*rm_cmd, path], timeout=timeout) + return + + for matched_path in matched: + self.run_cli([*rm_cmd, matched_path], timeout=timeout) + + def _transfer_path( + self, + action: str, + source_path: str, + dest_path: str, + *, + timeout: int, + fallback_action: str, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + ) -> None: + """Shared move and copy implementation with wildcard handling.""" + policy = self._validate_if_exists(if_exists, allow_overwrite=False) + if URIHelper.contains_wildcards(source_path): + matched = self.list_wildcard_matches( + source_path, timeout=timeout, fallback_action=fallback_action + ) + if matched is None: + self.make_path(dest_path, timeout=timeout) + self.run_cli(self._fs_cmd(action, source_path, dest_path), timeout=timeout) + return + + if not matched: + return + + self.make_path(dest_path, timeout=timeout) + dest_dir = dest_path.rstrip("/") + for matched_path in matched: + _, filename = URIHelper.split_ozone_path(matched_path) + dest_file_path = URIHelper.join_ozone_path(dest_dir, filename) + if self.exists(dest_file_path, timeout=timeout): + if policy == ExistingTargetPolicy.IGNORE: + self.log.info("Destination path %s already exists, skipping.", dest_file_path) + continue + raise OzoneProviderError(f"Destination path already exists: {dest_file_path}") + self.run_cli(self._fs_cmd(action, matched_path, dest_file_path), timeout=timeout) + return + + parent_path, _ = URIHelper.split_ozone_path(dest_path) + if parent_path: + self.make_path(parent_path, timeout=timeout) + if self.exists(dest_path, timeout=timeout): + if policy == ExistingTargetPolicy.IGNORE: + self.log.info("Destination path %s already exists, treating as success.", dest_path) + return + raise OzoneProviderError(f"Destination path already exists: {dest_path}") + if self.exists(source_path, timeout=timeout): + self.run_cli(self._fs_cmd(action, source_path, dest_path), timeout=timeout) + + +class OzoneResource(str, Enum): + """Supported Ozone admin resource types.""" + + VOLUME = "volume" + BUCKET = "bucket" + + +class OzoneAdminHook(OzoneCliHook): + """Interact with core namespace admin operations through ozone sh.""" + + default_conn_name = "ozone_admin_default" + hook_name = "Ozone Admin" + + # ============================== + # Volume operations + # ============================== + def create_volume( + self, + volume_name: str, + quota: str | None = None, + *, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + owner: str | None = None, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + ) -> None: + """Create a volume with optional quotas and owner.""" + self._create_resource( + resource=OzoneResource.VOLUME, + resource_path=self._resource_path(volume_name), + quota=quota, + space_quota=space_quota, + namespace_quota=namespace_quota, + owner=owner, + timeout=timeout, + if_exists=if_exists, + ) + + def list_volumes(self, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[JsonDict]: + """List all volumes visible to the current user.""" + return self._list_resources( + resource=OzoneResource.VOLUME, + parent_path="/", + timeout=timeout, + ) + + def get_volume_info(self, volume_name: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> JsonDict: + """Return detailed metadata for a volume.""" + return self._get_resource_info(volume=volume_name, timeout=timeout) + + def volume_exists(self, volume_name: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> bool: + """Return True when the target volume exists.""" + return self._resource_exists(volume=volume_name, timeout=timeout) + + def delete_volume( + self, + volume_name: str, + recursive: bool = False, + force: bool = False, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> None: + """Delete a volume, optionally recursively.""" + self._delete_resource( + resource=OzoneResource.VOLUME, + resource_path=self._resource_path(volume_name), + recursive=recursive, + force=force, + timeout=timeout, + ) + + def set_volume_owner( + self, + volume_name: str, + owner: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Change the owner of an existing volume.""" + parsed = self.run_cli( + ["ozone", "sh", "volume", "update", self._resource_path(volume_name), "--user", owner], + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from volume update {volume_name}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + # ============================== + # Bucket operations + # ============================== + def create_bucket( + self, + volume_name: str, + bucket_name: str, + quota: str | None = None, + *, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + owner: str | None = None, + layout: str | None = None, + replication_type: str | None = None, + replication: str | None = None, + encryption_key: str | None = None, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + ) -> None: + """Create a bucket with optional owner, layout, quotas, replication and encryption key.""" + extra_args: list[str] = [] + if owner: + extra_args.extend(["--user", owner]) + if layout: + extra_args.extend(["--layout", layout]) + if replication_type: + extra_args.extend(["--type", replication_type]) + if replication: + extra_args.extend(["--replication", replication]) + if encryption_key: + extra_args.extend(["--bucketkey", encryption_key]) + + self._create_resource( + resource=OzoneResource.BUCKET, + resource_path=self._resource_path(volume_name, bucket_name), + quota=quota, + space_quota=space_quota, + namespace_quota=namespace_quota, + extra_args=extra_args, + timeout=timeout, + if_exists=if_exists, + ) + + def list_buckets(self, volume_name: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[JsonDict]: + """List buckets under the given volume.""" + return self._list_resources( + resource=OzoneResource.BUCKET, + parent_path=self._resource_path(volume_name), + timeout=timeout, + ) + + def get_bucket_info( + self, + volume_name: str, + bucket_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Return detailed metadata for a bucket.""" + return self._get_resource_info(volume=volume_name, bucket=bucket_name, timeout=timeout) + + def bucket_exists( + self, + volume_name: str, + bucket_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> bool: + """Return True when the target bucket exists.""" + return self._resource_exists(volume=volume_name, bucket=bucket_name, timeout=timeout) + + def delete_bucket( + self, + volume_name: str, + bucket_name: str, + recursive: bool = False, + force: bool = False, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> None: + """Delete a bucket, optionally recursively.""" + self._delete_resource( + resource=OzoneResource.BUCKET, + resource_path=self._resource_path(volume_name, bucket_name), + recursive=recursive, + force=force, + timeout=timeout, + ) + + def set_bucket_owner( + self, + volume_name: str, + bucket_name: str, + owner: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Change the owner of an existing bucket.""" + parsed = self.run_cli( + [ + "ozone", + "sh", + "bucket", + "update", + self._resource_path(volume_name, bucket_name), + "--user", + owner, + ], + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + "Unexpected JSON payload from bucket update " + f"{volume_name}/{bucket_name}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + # ============================== + # Quota operations + # ============================== + def set_quota( + self, + *, + volume: str, + bucket: str | None = None, + quota: str | None = None, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Set space and/or namespace quota for a volume or bucket.""" + if quota is None and space_quota is None and namespace_quota is None: + raise ValueError("At least one quota must be provided.") + + self.run_cli( + self._quota_cmd( + "setquota", + volume=volume, + bucket=bucket, + quota=quota, + space_quota=space_quota, + namespace_quota=namespace_quota, + ), + timeout=timeout, + ) + + def clear_quota( + self, + *, + volume: str, + bucket: str | None = None, + clear_space_quota: bool = True, + clear_namespace_quota: bool = False, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Clear space and/or namespace quota for a volume or bucket.""" + if not clear_space_quota and not clear_namespace_quota: + raise ValueError("At least one quota flag must be set to True.") + + resource, target = self._resolve_admin_target(volume=volume, bucket=bucket) + cmd = self._admin_cmd(resource, "clrquota") + if clear_space_quota: + cmd.append("--space-quota") + if clear_namespace_quota: + cmd.append("--namespace-quota") + cmd.append(target) + self.run_cli(cmd, timeout=timeout) + + def get_quota( + self, + *, + volume: str, + bucket: str | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Return quota and usage fields from volume or bucket info.""" + info = self._get_resource_info(volume=volume, bucket=bucket, timeout=timeout) + return { + "quota_in_bytes": info.get("quotaInBytes"), + "quota_in_namespace": info.get("quotaInNamespace"), + "used_bytes": info.get("usedBytes"), + "used_namespace": info.get("usedNamespace"), + } + + def set_volume_quota( + self, + volume_name: str, + *, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Set one or both quota types for a volume.""" + self.set_quota( + volume=volume_name, + space_quota=space_quota, + namespace_quota=namespace_quota, + timeout=timeout, + ) + + def set_bucket_quota( + self, + volume_name: str, + bucket_name: str, + *, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Set one or both quota types for a bucket.""" + self.set_quota( + volume=volume_name, + bucket=bucket_name, + space_quota=space_quota, + namespace_quota=namespace_quota, + timeout=timeout, + ) + + def clear_volume_quota( + self, + volume_name: str, + *, + clear_space_quota: bool = True, + clear_namespace_quota: bool = False, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Clear one or both quota types for a volume.""" + self.clear_quota( + volume=volume_name, + clear_space_quota=clear_space_quota, + clear_namespace_quota=clear_namespace_quota, + timeout=timeout, + ) + + def clear_bucket_quota( + self, + volume_name: str, + bucket_name: str, + *, + clear_space_quota: bool = True, + clear_namespace_quota: bool = False, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> None: + """Clear one or both quota types for a bucket.""" + self.clear_quota( + volume=volume_name, + bucket=bucket_name, + clear_space_quota=clear_space_quota, + clear_namespace_quota=clear_namespace_quota, + timeout=timeout, + ) + + def get_volume_quota(self, volume_name: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> JsonDict: + """Return quota and usage fields for a volume.""" + return self.get_quota(volume=volume_name, timeout=timeout) + + def get_bucket_quota( + self, + volume_name: str, + bucket_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Return quota and usage fields for a bucket.""" + return self.get_quota(volume=volume_name, bucket=bucket_name, timeout=timeout) + + # ============================== + # Shared admin helpers + # ============================== + def _admin_cmd(self, resource: OzoneResource, action: str, *args: str) -> list[str]: + """Build an ozone sh command for a resource.""" + return ["ozone", "sh", resource.value, action, *args] + + def _resource_path(self, volume: str, bucket: str | None = None) -> str: + """Build a native Ozone resource path.""" + if bucket: + return f"/{volume}/{bucket}" + return f"/{volume}" + + def _resolve_admin_target( + self, + *, + volume: str, + bucket: str | None = None, + ) -> tuple[OzoneResource, str]: + """Resolve resource type and path for volume-or-bucket methods.""" + if bucket: + return OzoneResource.BUCKET, self._resource_path(volume, bucket) + return OzoneResource.VOLUME, self._resource_path(volume) + + def _normalize_quota_inputs( + self, + *, + quota: str | None = None, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + ) -> list[str]: + """Build quota CLI args while keeping old quota as a compatibility alias.""" + if quota is not None and space_quota is not None: + raise ValueError("Use either quota or space_quota, but not both.") + + effective_space_quota = space_quota if space_quota is not None else quota + args: list[str] = [] + if effective_space_quota is not None: + args.extend(["--space-quota", str(effective_space_quota)]) + if namespace_quota is not None: + args.extend(["--namespace-quota", str(namespace_quota)]) + return args + + def _create_resource( + self, + *, + resource: OzoneResource, + resource_path: str, + quota: str | None = None, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + owner: str | None = None, + extra_args: list[str] | None = None, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + ) -> None: + """Create a volume or bucket and treat already-exists errors as success.""" + policy = OzoneFsHook._validate_if_exists(if_exists, allow_overwrite=False) + spec = ADMIN_RESOURCE_SPECS[resource.value] + cmd = self._admin_cmd( + resource, + "create", + *self._normalize_quota_inputs( + quota=quota, + space_quota=space_quota, + namespace_quota=namespace_quota, + ), + ) + if owner: + cmd.extend(["--user", owner]) + if extra_args: + cmd.extend(extra_args) + cmd.append(resource_path) + + result = self.run_cli( + cmd, + timeout=timeout, + retry_attempts=0, + check=False, + log_output=False, + return_result=True, + ) + if result.returncode == 0: + return + + output = ProcessOutputAnalysis.from_completed_process(result) + error_message = output.preferred_output or "No error message provided" + if spec.already_exists_marker.lower() in error_message.lower(): + if policy == ExistingTargetPolicy.ERROR: + raise OzoneProviderError(f"{resource.value.capitalize()} already exists: {resource_path}") + self.log.info( + "%s %s already exists, treating as success.", resource.value.capitalize(), resource_path + ) + return + + raise OzoneProviderError( + f"Ozone command failed (return code: {result.returncode}): {redact(error_message)}" + ) + + def _delete_resource( + self, + *, + resource: OzoneResource, + resource_path: str, + recursive: bool = False, + force: bool = False, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> None: + """Delete a volume or bucket with idempotent not-found handling.""" + spec = ADMIN_RESOURCE_SPECS[resource.value] + cmd = self._admin_cmd(resource, "delete") + if recursive: + cmd.append("-r") + cmd.append(resource_path) + input_text = "yes\n" if recursive and force else None + + result = self.run_cli( + cmd, + timeout=timeout, + input_text=input_text, + retry_attempts=0, + check=False, + log_output=False, + return_result=True, + ) + if result.returncode == 0: + return + + output = ProcessOutputAnalysis.from_completed_process(result) + error_text = output.preferred_output or "Unknown error" + normalized = output.normalized_preferred_output + if ( + any(marker.lower() in normalized for marker in spec.not_found_markers) + or "does not exist" in normalized + ): + self.log.info( + "%s %s does not exist, treating as success.", resource.value.capitalize(), resource_path + ) + return + + raise OzoneProviderError(f"Failed to delete {resource.value} {resource_path}: {redact(error_text)}") + + def _get_resource_info( + self, + *, + volume: str, + bucket: str | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Return JSON info for a volume or bucket.""" + resource, target = self._resolve_admin_target(volume=volume, bucket=bucket) + parsed = self.run_cli( + self._admin_cmd(resource, "info", target), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from {resource.value} info {target}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def _resource_exists( + self, + *, + volume: str, + bucket: str | None = None, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> bool: + """Check existence of a volume or bucket using the corresponding info command.""" + resource, target = self._resolve_admin_target(volume=volume, bucket=bucket) + spec = ADMIN_RESOURCE_SPECS[resource.value] + + result = self.run_cli( + self._admin_cmd(resource, "info", target), + timeout=timeout, + retry_attempts=0, + check=False, + log_output=False, + return_result=True, + ) + if result.returncode == 0: + return True + + output = ProcessOutputAnalysis.from_completed_process(result) + error_text = output.preferred_output + normalized = output.normalized_preferred_output + if ( + any(marker.lower() in normalized for marker in spec.not_found_markers) + or "does not exist" in normalized + ): + return False + + raise OzoneProviderError( + f"Failed to check {resource.value} existence for {target}: {redact(error_text or 'Unknown error')}" + ) + + def _list_resources( + self, + *, + resource: OzoneResource, + parent_path: str, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> list[JsonDict]: + """Return parsed list output for volume or bucket commands.""" + parsed = self.run_cli( + self._admin_cmd(resource, "list", parent_path), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, list): + raise OzoneProviderError( + f"Unexpected JSON payload from {resource.value} list {parent_path}: expected list, got {type(parsed).__name__}." + ) + return parsed + + def _quota_cmd( + self, + action: str, + *, + volume: str, + bucket: str | None = None, + quota: str | None = None, + space_quota: str | None = None, + namespace_quota: str | int | None = None, + ) -> list[str]: + """Build a quota-related command for a volume or bucket.""" + resource, target = self._resolve_admin_target(volume=volume, bucket=bucket) + return self._admin_cmd( + resource, + action, + *self._normalize_quota_inputs( + quota=quota, + space_quota=space_quota, + namespace_quota=namespace_quota, + ), + target, + ) + + +class OzoneAdminExtraHook(OzoneAdminHook): + """Interact with advanced admin features: snapshots, tenants and cluster reports.""" + + hook_name = "Ozone Admin Extra" + + # ============================== + # Advanced bucket operations + # ============================== + def set_bucket_replication_config( + self, + volume_name: str, + bucket_name: str, + replication_type: str, + replication: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Set default replication config for new keys in a bucket.""" + return self.run_cli( + [ + "ozone", + "sh", + "bucket", + "set-replication-config", + self._resource_path(volume_name, bucket_name), + "--type", + replication_type, + "--replication", + replication, + ], + timeout=timeout, + ) + + def create_bucket_link( + self, + source_volume: str, + source_bucket: str, + target_volume: str, + target_bucket: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Create a bucket link pointing to another bucket in the same cluster.""" + return self.run_cli( + [ + "ozone", + "sh", + "bucket", + "link", + self._resource_path(source_volume, source_bucket), + self._resource_path(target_volume, target_bucket), + ], + timeout=timeout, + ) + + # ============================== + # Snapshot operations + # ============================== + def create_snapshot( + self, + volume_name: str, + bucket_name: str, + snapshot_name: str | None = None, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> str: + """Create a point-in-time snapshot for a bucket.""" + cmd = self._snapshot_cmd("create", self._resource_path(volume_name, bucket_name)) + if snapshot_name: + cmd.append(snapshot_name) + return self.run_cli(cmd, timeout=timeout) + + def list_snapshots( + self, + volume_name: str, + bucket_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> list[str]: + """List snapshots for a bucket.""" + output = self.run_cli( + self._snapshot_cmd("list", self._resource_path(volume_name, bucket_name)), + timeout=timeout, + ) + return output.splitlines() if output else [] + + def get_snapshot_info( + self, + volume_name: str, + bucket_name: str, + snapshot_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Return raw snapshot information for a named snapshot.""" + return self.run_cli( + self._snapshot_cmd("info", self._resource_path(volume_name, bucket_name), snapshot_name), + timeout=timeout, + ) + + def diff_snapshots( + self, + volume_name: str, + bucket_name: str, + from_snapshot: str, + to_snapshot_or_bucket: str, + *, + page_size: int | None = None, + continuation_token: str | None = None, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> str: + """Return diff output between two snapshots or a snapshot and the live bucket.""" + cmd = self._snapshot_cmd("diff") + if page_size is not None: + cmd.extend(["-p", str(page_size)]) + if continuation_token: + cmd.extend(["-t", continuation_token]) + cmd.extend( + [ + self._resource_path(volume_name, bucket_name), + from_snapshot, + to_snapshot_or_bucket, + ] + ) + return self.run_cli(cmd, timeout=timeout) + + def list_snapshot_diff_jobs( + self, + volume_name: str, + bucket_name: str, + *, + job_status: str | None = None, + all_status: bool = False, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """List snapshot diff jobs for a bucket.""" + cmd = self._snapshot_cmd("listDiff", self._resource_path(volume_name, bucket_name)) + if job_status: + cmd.extend(["--job-status", job_status]) + if all_status: + cmd.append("--all-status") + return self.run_cli(cmd, timeout=timeout) + + def cancel_snapshot_diff(self, job_id: str, *, timeout: int = FAST_TIMEOUT_SECONDS) -> str: + """Cancel an in-progress snapshot diff job by its job ID.""" + return self.run_cli(self._snapshot_cmd("cancelDiff", job_id), timeout=timeout) + + def delete_snapshot( + self, + volume_name: str, + bucket_name: str, + snapshot_name: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> str: + """Delete a named snapshot from a bucket.""" + return self.run_cli( + self._snapshot_cmd("delete", self._resource_path(volume_name, bucket_name), snapshot_name), + timeout=timeout, + ) + + # ============================== + # Tenant operations + # ============================== + def create_tenant( + self, + tenant_name: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> JsonDict: + """Create a new tenant and return verbose JSON metadata.""" + parsed = self.run_cli( + self._tenant_cmd("--verbose", "create", tenant_name), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant create {tenant_name}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def list_tenants(self, *, timeout: int = FAST_TIMEOUT_SECONDS) -> list[JsonDict]: + """List all tenants in JSON form.""" + parsed = self.run_cli( + self._tenant_cmd("list", "--json"), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, list): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant list: expected list, got {type(parsed).__name__}." + ) + return parsed + + def delete_tenant( + self, + tenant_name: str, + *, + timeout: int = SLOW_TIMEOUT_SECONDS, + ) -> JsonDict: + """Delete an empty tenant and return verbose JSON metadata.""" + parsed = self.run_cli( + self._tenant_cmd("--verbose", "delete", tenant_name), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant delete {tenant_name}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def tenant_assign_user( + self, + user_name: str, + tenant_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Assign a user to a tenant and let Ozone generate access credentials.""" + return self.run_cli( + self._tenant_cmd("user", "assign", user_name, f"--tenant={tenant_name}"), + timeout=timeout, + ) + + def tenant_revoke_user( + self, + access_id: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Revoke tenant access for an access ID.""" + return self.run_cli(self._tenant_cmd("user", "revoke", access_id), timeout=timeout) + + def tenant_assign_admin( + self, + access_id: str, + tenant_name: str, + *, + delegated: bool = False, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Promote a tenant access ID to tenant admin and return verbose JSON metadata.""" + cmd = self._tenant_cmd("--verbose", "user", "assignadmin", access_id) + if delegated: + cmd.append("--delegated") + cmd.append(f"--tenant={tenant_name}") + parsed = self.run_cli( + cmd, + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant assignadmin {access_id}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def tenant_revoke_admin( + self, + access_id: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Revoke tenant admin privileges and return verbose JSON metadata.""" + parsed = self.run_cli( + self._tenant_cmd("--verbose", "user", "revokeadmin", access_id), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant revokeadmin {access_id}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def list_tenant_users( + self, + tenant_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> list[JsonDict]: + """List users assigned to a tenant in JSON form.""" + parsed = self.run_cli( + self._tenant_cmd("user", "list", "--json", tenant_name), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, list): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant user list {tenant_name}: expected list, got {type(parsed).__name__}." + ) + return parsed + + def get_tenant_user_info( + self, + user_name: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> JsonDict: + """Return all tenant assignments for a user in JSON form.""" + parsed = self.run_cli( + self._tenant_cmd("user", "info", "--json", user_name), + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, dict): + raise OzoneProviderError( + f"Unexpected JSON payload from tenant user info {user_name}: expected dict, got {type(parsed).__name__}." + ) + return parsed + + def get_tenant_user_secret( + self, + access_id: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Return the current secret for a tenant access ID without generating a new one.""" + return self.run_cli(self._tenant_cmd("user", "get-secret", access_id), timeout=timeout) + + def set_tenant_user_secret( + self, + access_id: str, + secret_key: str, + *, + timeout: int = FAST_TIMEOUT_SECONDS, + ) -> str: + """Set the secret key for a tenant access ID.""" + return self.run_cli( + self._tenant_cmd("user", "set-secret", access_id, "--secret", secret_key), + timeout=timeout, + ) + + # ============================== + # Cluster report operations + # ============================== + def get_container_report(self, *, timeout: int = SLOW_TIMEOUT_SECONDS) -> JsonDict | list[JsonDict]: + """Fetch and parse the JSON container report from SCM.""" + parsed = self.run_cli( + ["ozone", "admin", "container", "report", "--json"], + timeout=timeout, + return_json_result=True, + ) + if not isinstance(parsed, (dict, list)): + raise OzoneProviderError( + "Unexpected JSON payload from container report: " + f"expected dict or list, got {type(parsed).__name__}." + ) + return parsed + + # ============================== + # Shared extra helpers + # ============================== + def _tenant_cmd(self, *args: str) -> list[str]: + """Build an ozone tenant command.""" + return ["ozone", "tenant", *args] + + def _snapshot_cmd(self, action: str, *args: str) -> list[str]: + """Build an ozone sh snapshot command.""" + return ["ozone", "sh", "snapshot", action, *args] diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/__init__.py new file mode 100644 index 0000000000000..14ab3a1aa5bc7 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/__init__.py @@ -0,0 +1,57 @@ +# +# 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. + +"""Ozone operators.""" + +from __future__ import annotations + +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCopyOperator, + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneDeleteBucketOperator, + OzoneDeleteKeyOperator, + OzoneDeletePathOperator, + OzoneDeleteVolumeOperator, + OzoneDownloadFileOperator, + OzoneListOperator, + OzoneMoveOperator, + OzonePathExistsOperator, + OzoneSetQuotaOperator, + OzoneUploadContentOperator, + OzoneUploadFileOperator, +) + +__all__ = [ + "OzoneCreateVolumeOperator", + "OzoneCreateBucketOperator", + "OzoneSetQuotaOperator", + "OzoneDeleteVolumeOperator", + "OzoneDeleteBucketOperator", + "OzoneCreatePathOperator", + "OzoneUploadContentOperator", + "OzoneUploadFileOperator", + "OzoneDeleteKeyOperator", + "OzoneDeletePathOperator", + "OzonePathExistsOperator", + "OzoneListOperator", + "OzoneCopyOperator", + "OzoneMoveOperator", + "OzoneDownloadFileOperator", +] diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/ozone.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/ozone.py new file mode 100644 index 0000000000000..7711228d1b48f --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/operators/ozone.py @@ -0,0 +1,619 @@ +# 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 + +import tempfile +from pathlib import Path +from typing import TYPE_CHECKING + +from airflow.providers.arenadata.ozone.hooks.ozone import ( + FAST_TIMEOUT_SECONDS, + RETRY_ATTEMPTS, + SLOW_TIMEOUT_SECONDS, + ExistingTargetPolicy, + OzoneAdminHook, + OzoneFsHook, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.utils.helpers import FileHelper, TypeNormalizationHelper +from airflow.providers.arenadata.ozone.version_compat import BaseOperator + +if TYPE_CHECKING: + from airflow.providers.arenadata.ozone.version_compat import Context + + +class OzoneCreateVolumeOperator(BaseOperator): + """Create an Ozone Volume using Native Admin CLI.""" + + template_fields = ("volume_name", "quota", "ozone_conn_id") + + def __init__( + self, + volume_name: str, + quota: str | None = None, + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume_name = volume_name + self.quota = quota + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Create the target volume if it does not exist.""" + quota = TypeNormalizationHelper.require_optional_non_empty( + self.quota, "quota parameter cannot be an empty string (use None for unlimited)" + ) + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.create_volume(self.volume_name, quota, timeout=self.timeout, if_exists=self.if_exists) + + +class OzoneCreateBucketOperator(BaseOperator): + """Create an Ozone Bucket using Native Admin CLI.""" + + template_fields = ("volume_name", "bucket_name", "quota", "ozone_conn_id") + + def __init__( + self, + volume_name: str, + bucket_name: str, + quota: str | None = None, + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume_name = volume_name + self.bucket_name = bucket_name + self.quota = quota + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Create the target bucket if it does not exist.""" + quota = TypeNormalizationHelper.require_optional_non_empty( + self.quota, "quota parameter cannot be an empty string (use None for unlimited)" + ) + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.create_bucket( + self.volume_name, + self.bucket_name, + quota, + timeout=self.timeout, + if_exists=self.if_exists, + ) + + +class OzoneSetQuotaOperator(BaseOperator): + """Dynamically adjust volume or bucket quotas.""" + + template_fields = ("volume", "quota", "bucket", "ozone_conn_id") + + def __init__( + self, + volume: str, + quota: str, + bucket: str | None = None, + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = FAST_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume = volume + self.quota = quota + self.bucket = bucket + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Apply a quota to a volume or bucket.""" + bucket = TypeNormalizationHelper.require_optional_non_empty( + self.bucket, "bucket parameter cannot be an empty string (use None for volume quota)" + ) + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.set_quota( + volume=self.volume, + quota=self.quota, + bucket=bucket, + timeout=self.timeout, + ) + + +class OzoneDeleteVolumeOperator(BaseOperator): + """Delete an Ozone Volume using Native Admin CLI.""" + + template_fields = ("volume_name", "ozone_conn_id") + + def __init__( + self, + volume_name: str, + recursive: bool = False, + force: bool = False, + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume_name = volume_name + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + if force and not recursive: + raise ValueError("force=True requires recursive=True") + self.recursive = recursive + self.force = force + + def execute(self, context: Context) -> None: + """Delete a volume with optional recursive mode.""" + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.delete_volume( + self.volume_name, + self.recursive, + self.force, + timeout=self.timeout, + ) + + +class OzoneDeleteBucketOperator(BaseOperator): + """Delete an Ozone Bucket using Native Admin CLI.""" + + template_fields = ("volume_name", "bucket_name", "ozone_conn_id") + + def __init__( + self, + volume_name: str, + bucket_name: str, + recursive: bool = False, + force: bool = False, + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume_name = volume_name + self.bucket_name = bucket_name + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + if force and not recursive: + raise ValueError("force=True requires recursive=True") + self.recursive = recursive + self.force = force + + def execute(self, context: Context) -> None: + """Delete a bucket with optional recursive mode.""" + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.delete_bucket( + self.volume_name, + self.bucket_name, + self.recursive, + self.force, + timeout=self.timeout, + ) + + +class OzoneCreatePathOperator(BaseOperator): + """Create directory path via FS interface.""" + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = FAST_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.IGNORE, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Create a directory path in Ozone FS.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.make_path(self.path, timeout=self.timeout, if_exists=self.if_exists) + + +class OzoneUploadContentOperator(BaseOperator): + """Put content string to a file via FS interface.""" + + template_fields = ("content", "remote_path", "ozone_conn_id") + + def __init__( + self, + content: str, + remote_path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + max_content_size_bytes: int | None = None, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.remote_path = remote_path + self.ozone_conn_id = ozone_conn_id + self.content = content + self.max_content_size_bytes = TypeNormalizationHelper.require_optional_positive_int( + max_content_size_bytes, + field_name="max_content_size_bytes", + ) + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Write string content to a temporary file and upload it.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + size_limit_bytes = self.max_content_size_bytes or hook.connection_snapshot.max_content_size_bytes + with tempfile.TemporaryDirectory(prefix="ozone_fs_put_") as tmp_dir: + tmp_path = Path(tmp_dir) / "payload.txt" + tmp_path.write_text(self.content, encoding="utf-8") + content_size_bytes = FileHelper.get_file_size_bytes(tmp_path) + if content_size_bytes <= 0: + raise OzoneProviderError( + "Unable to determine payload size for Ozone upload " + f"({content_size_bytes} bytes) from temporary file: {tmp_path}" + ) + if content_size_bytes > size_limit_bytes: + raise OzoneProviderError( + f"Content size ({content_size_bytes} bytes) exceeds configured limit " + f"({size_limit_bytes} bytes) for Ozone upload." + ) + hook.upload_key( + str(tmp_path), + self.remote_path, + timeout=self.timeout, + if_exists=self.if_exists, + ) + + +class OzoneDeleteKeyOperator(BaseOperator): + """Deletes a key (file) from Ozone FS.""" + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Delete a single key or a wildcard-matched group of keys.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.delete_key(self.path, timeout=self.timeout) + + +class OzoneDeletePathOperator(BaseOperator): + """Delete file/directory path from Ozone FS.""" + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + recursive: bool = True, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.recursive = recursive + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Delete file/directory path with optional recursive mode.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.delete_path(self.path, recursive=self.recursive, timeout=self.timeout) + + +class OzonePathExistsOperator(BaseOperator): + """Check whether Ozone path exists and return boolean via XCom.""" + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = FAST_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> bool: + """Return True when path exists, False otherwise.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + return hook.path_exists(self.path, timeout=self.timeout) + + +class OzoneListOperator(BaseOperator): + """Lists keys/prefixes in an Ozone path and returns them via XCom.""" + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + *, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = FAST_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> list[str]: + """Return a list of keys, optionally filtered by wildcard pattern.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + return hook.list_keys(self.path, timeout=self.timeout) + + +class OzoneUploadFileOperator(BaseOperator): + """Upload a file from the local filesystem to Ozone.""" + + template_fields = ("local_path", "remote_path", "ozone_conn_id") + + def __init__( + self, + local_path: str, + remote_path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + overwrite: bool = False, + max_content_size_bytes: int | None = None, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.local_path = local_path + self.remote_path = remote_path + self.ozone_conn_id = ozone_conn_id + self.overwrite = overwrite + self.max_content_size_bytes = TypeNormalizationHelper.require_optional_positive_int( + max_content_size_bytes, + field_name="max_content_size_bytes", + ) + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Upload a local file to Ozone, with optional overwrite.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + local_path_obj = Path(self.local_path) + if not FileHelper.is_readable_file(local_path_obj): + raise OzoneProviderError(f"Local file not found or is not readable: {self.local_path}") + + size_limit_bytes = self.max_content_size_bytes or hook.connection_snapshot.max_content_size_bytes + file_size_bytes = FileHelper.get_file_size_bytes(local_path_obj) + if file_size_bytes <= 0: + raise OzoneProviderError( + f"Local file size is unavailable or non-positive ({file_size_bytes} bytes): {self.local_path}" + ) + if file_size_bytes > size_limit_bytes: + raise OzoneProviderError( + f"Local file size ({file_size_bytes} bytes) exceeds configured limit " + f"({size_limit_bytes} bytes) for Ozone upload: {self.local_path}" + ) + + hook.upload_key( + str(local_path_obj), + self.remote_path, + timeout=self.timeout, + if_exists="overwrite" if self.overwrite else "error", + ) + + +class OzoneMoveOperator(BaseOperator): + """Move or rename a key within the same Ozone cluster.""" + + template_fields = ("source_path", "dest_path", "ozone_conn_id") + + def __init__( + self, + *, + source_path: str, + dest_path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.source_path = source_path + self.dest_path = dest_path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Move one key or a wildcard-selected batch within Ozone.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.move(self.source_path, self.dest_path, timeout=self.timeout, if_exists=self.if_exists) + + +class OzoneCopyOperator(BaseOperator): + """Copy one key or wildcard-selected keys within Ozone.""" + + template_fields = ("source_path", "dest_path", "ozone_conn_id") + + def __init__( + self, + *, + source_path: str, + dest_path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + if_exists: ExistingTargetPolicy | str = ExistingTargetPolicy.ERROR, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.source_path = source_path + self.dest_path = dest_path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + self.if_exists = if_exists + + def execute(self, context: Context) -> None: + """Copy one key or a wildcard-selected batch within Ozone.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + hook.copy_path(self.source_path, self.dest_path, timeout=self.timeout, if_exists=self.if_exists) + + +class OzoneDownloadFileOperator(BaseOperator): + """Download a file from Ozone to local filesystem.""" + + template_fields = ("remote_path", "local_path", "ozone_conn_id") + + def __init__( + self, + *, + remote_path: str, + local_path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + overwrite: bool = False, + max_content_size_bytes: int | None = None, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.remote_path = remote_path + self.local_path = local_path + self.ozone_conn_id = ozone_conn_id + self.overwrite = overwrite + self.max_content_size_bytes = TypeNormalizationHelper.require_optional_positive_int( + max_content_size_bytes, + field_name="max_content_size_bytes", + ) + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Download remote file to local path.""" + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + size_limit_bytes = self.max_content_size_bytes or hook.connection_snapshot.max_content_size_bytes + key_properties = hook.get_key_property(self.remote_path, timeout=self.timeout) + remote_size_bytes = TypeNormalizationHelper.parse_non_negative_int(key_properties.get("data_size")) + if remote_size_bytes is None: + raise OzoneProviderError( + f"Unable to determine remote key size for {self.remote_path}; " + "missing or invalid 'data_size' in key properties." + ) + if remote_size_bytes > size_limit_bytes: + raise OzoneProviderError( + f"Remote file size ({remote_size_bytes} bytes) exceeds configured limit " + f"({size_limit_bytes} bytes) for Ozone download: {self.remote_path}" + ) + + hook.download_key( + self.remote_path, + self.local_path, + overwrite=self.overwrite, + timeout=self.timeout, + ) diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/__init__.py new file mode 100644 index 0000000000000..31a0c2de234aa --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/__init__.py @@ -0,0 +1,25 @@ +# +# 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. + +"""Ozone sensors.""" + +from __future__ import annotations + +from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor + +__all__ = ["OzoneKeySensor"] diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/ozone.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/ozone.py new file mode 100644 index 0000000000000..878dbb88e04ef --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/sensors/ozone.py @@ -0,0 +1,76 @@ +# +# 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 TYPE_CHECKING + +from airflow.providers.arenadata.ozone.hooks.ozone import ( + FAST_TIMEOUT_SECONDS, + RETRY_ATTEMPTS, + OzoneFsHook, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError +from airflow.providers.arenadata.ozone.version_compat import BaseSensorOperator + +if TYPE_CHECKING: + from airflow.providers.arenadata.ozone.version_compat import Context + + +class OzoneKeySensor(BaseSensorOperator): + """ + Wait for a file or directory (key) to exist on Ozone FS (ofs:// or o3fs://). + + Uses OzoneFsHook; retryable CLI errors are treated as "not yet" and trigger retry. + """ + + template_fields = ("path", "ozone_conn_id") + + def __init__( + self, + path: str, + ozone_conn_id: str = OzoneFsHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = FAST_TIMEOUT_SECONDS, # Timeout for object existence check operation in Ozone CLI + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.path = path + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + # Keep the Ozone CLI timeout separate from BaseSensorOperator.timeout. + self.cli_timeout = timeout + self.log.debug("OzoneKeySensor initialized (path=%s, conn_id=%s)", self.path, self.ozone_conn_id) + + def poke(self, context: Context) -> bool: + """Return True when the target path appears in Ozone.""" + self.log.debug("Checking if key exists in Ozone: %s", self.path) + hook = OzoneFsHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + try: + if hook.key_exists(self.path, timeout=self.cli_timeout): + self.log.info("Key found in Ozone: %s", self.path) + return True + return False + except OzoneCliError as e: + if e.retryable: + self.log.warning("Retryable error while checking key existence (will retry): %s", str(e)) + return False + raise diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/__init__.py new file mode 100644 index 0000000000000..6f5757b3e4a57 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/__init__.py @@ -0,0 +1,29 @@ +# +# 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. + +"""Ozone transfer operators.""" + +from __future__ import annotations + +from airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone import HdfsToOzoneOperator +from airflow.providers.arenadata.ozone.transfers.ozone_backup import OzoneBackupOperator + +__all__ = [ + "HdfsToOzoneOperator", + "OzoneBackupOperator", +] diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/hdfs_to_ozone.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/hdfs_to_ozone.py new file mode 100644 index 0000000000000..c28a9242224b4 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/hdfs_to_ozone.py @@ -0,0 +1,227 @@ +# 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 + +import logging +import shutil +from functools import cached_property +from typing import TYPE_CHECKING + +from airflow.providers.arenadata.ozone.hooks.ozone import ( + RETRY_ATTEMPTS, + SLOW_TIMEOUT_SECONDS, +) +from airflow.providers.arenadata.ozone.utils.cli_runner import CliRunner +from airflow.providers.arenadata.ozone.utils.connection_schema import ( + OzoneConnSnapshot, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.utils.security import ( + KerberosConfig, + SSLConfig, +) +from airflow.providers.arenadata.ozone.version_compat import ( + AirflowException, + BaseHook, + BaseOperator, +) + +if TYPE_CHECKING: + from airflow.providers.arenadata.ozone.version_compat import Context + +log = logging.getLogger(__name__) +DISTCP_BASE_COMMAND = ["hadoop", "distcp"] +DISTCP_COPY_OPTIONS = ["-update", "-skipcrccheck"] +DISTCP_MAPREDUCE_LOCAL_OPTION = "-Dmapreduce.framework.name=local" +DISTCP_YARN_RENEWER_PRINCIPAL_OPTION = "-Dyarn.resourcemanager.principal={principal}" +DISTCP_JOBTRACKER_RENEWER_PRINCIPAL_OPTION = "-Dmapreduce.jobtracker.kerberos.principal={principal}" + + +class HdfsToOzoneOperator(BaseOperator): + """ + Migrate data from HDFS to Ozone using DistCp. + + Optional HDFS SSL/TLS and Kerberos settings are read from the + ``hdfs_conn_id`` connection extra. + """ + + template_fields = ("source_path", "dest_path", "hdfs_conn_id") + + def __init__( + self, + source_path: str, + dest_path: str, + hdfs_conn_id: str | None = None, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if hdfs_conn_id is not None and not isinstance(hdfs_conn_id, str): + raise ValueError(f"hdfs_conn_id must be a string or None, got {type(hdfs_conn_id).__name__}") + self.source_path = source_path + self.dest_path = dest_path + self.hdfs_conn_id = hdfs_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + self.log.debug( + "Initializing HdfsToOzoneOperator - source: %s, destination: %s", self.source_path, self.dest_path + ) + + @cached_property + def _hdfs_ssl_env(self) -> dict[str, str] | None: + """Load SSL/TLS configuration from HDFS snapshot lazily.""" + if not self._hdfs_connection_snapshot: + return None + try: + ssl_env_vars = SSLConfig.from_snapshot( + self._hdfs_connection_snapshot, + conn_id=self.hdfs_conn_id, + scope="hdfs", + ).as_env() + if not ssl_env_vars: + self.log.debug("No HDFS SSL/TLS configuration found in connection snapshot") + return None + ssl_env = SSLConfig.apply_ssl_env_vars(ssl_env_vars) + self.log.debug("HDFS SSL/TLS configuration loaded from connection: %s", list(ssl_env_vars.keys())) + if self._hdfs_connection_snapshot.hdfs_ssl_enabled: + self.log.info("HDFS SSL/TLS enabled for connection") + return ssl_env + except AirflowException as err: + # Connection might not exist yet, that's OK + self.log.debug("Could not load HDFS SSL configuration (connection may not exist): %s", str(err)) + return None + + @cached_property + def _hdfs_connection_snapshot(self) -> OzoneConnSnapshot | None: + """Read HDFS connection snapshot once for Kerberos/SSL helpers.""" + if not self.hdfs_conn_id: + return None + try: + conn = BaseHook.get_connection(self.hdfs_conn_id) + return OzoneConnSnapshot.from_connection( + conn, + conn_id=self.hdfs_conn_id, + require_host_port=False, + ) + except AirflowException as err: + self.log.debug("Could not load HDFS connection extra (connection may not exist): %s", str(err)) + return None + + @cached_property + def _hdfs_kerberos_env(self) -> dict[str, str] | None: + """Load HDFS Kerberos subprocess env from the connection snapshot.""" + if not self.hdfs_conn_id: + return None + if not self._hdfs_connection_snapshot: + return None + return KerberosConfig.load_hdfs_env( + snapshot=self._hdfs_connection_snapshot, + conn_id=self.hdfs_conn_id, + ) + + def _build_distcp_env(self) -> dict[str, str] | None: + """Build DistCp env and initialize HDFS Kerberos when configured.""" + env: dict[str, str] = {} + if self._hdfs_ssl_env: + self.log.debug("Applying SSL environment variables for HDFS DistCp") + env.update(self._hdfs_ssl_env) + + if self._hdfs_kerberos_env: + snapshot = self._hdfs_connection_snapshot + if snapshot is None: + raise OzoneProviderError( + "HDFS Kerberos environment exists but HDFS connection snapshot is missing." + ) + if not KerberosConfig.kinit_hdfs_from_snapshot( + snapshot=snapshot, + conn_id=self.hdfs_conn_id, + ): + raise OzoneProviderError( + f"HDFS Kerberos authentication failed for connection '{self.hdfs_conn_id}' " + f"using principal '{snapshot.hdfs_kerberos_principal}'." + ) + env.update(self._hdfs_kerberos_env) + env["HADOOP_SECURITY_AUTHENTICATION"] = "kerberos" + if snapshot.krb5_conf: + env["KRB5_CONFIG"] = snapshot.krb5_conf + + return env or None + + def _build_distcp_command(self) -> list[str]: + """Build the DistCp command for the current transfer.""" + options: list[str] = [] + snapshot = self._hdfs_connection_snapshot + if snapshot: + if snapshot.hdfs_distcp_mapreduce_local: + options.append(DISTCP_MAPREDUCE_LOCAL_OPTION) + if snapshot.hdfs_distcp_renewer_principal: + principal = snapshot.hdfs_distcp_renewer_principal + options.extend( + [ + DISTCP_YARN_RENEWER_PRINCIPAL_OPTION.format(principal=principal), + DISTCP_JOBTRACKER_RENEWER_PRINCIPAL_OPTION.format(principal=principal), + ] + ) + return [*DISTCP_BASE_COMMAND, *options, *DISTCP_COPY_OPTIONS, self.source_path, self.dest_path] + + def _validate_runtime_inputs(self) -> None: + """Validate operator inputs right before DistCp execution.""" + if not isinstance(self.source_path, str): + raise OzoneProviderError("HdfsToOzoneOperator requires source_path to be a string") + if not self.source_path.strip(): + raise OzoneProviderError("HdfsToOzoneOperator requires non-empty source_path") + + if not isinstance(self.dest_path, str): + raise OzoneProviderError("HdfsToOzoneOperator requires dest_path to be a string") + if not self.dest_path.strip(): + raise OzoneProviderError("HdfsToOzoneOperator requires non-empty dest_path") + + def _validate_hadoop_runtime(self) -> str: + """Check that the Hadoop DistCp runtime binary is available.""" + hadoop_bin = shutil.which("hadoop") + if not hadoop_bin: + raise OzoneProviderError( + "HdfsToOzoneOperator is loaded, but Hadoop DistCp runtime is unavailable: " + "executable 'hadoop' was not found in PATH. " + "Install Hadoop client tools or provide 'hadoop' in PATH on the worker that runs this task" + ) + return hadoop_bin + + def _validate_runtime_dependencies(self) -> None: + """Run fail-first runtime validation for DistCp prerequisites.""" + self._validate_runtime_inputs() + self._validate_hadoop_runtime() + + def execute(self, context: Context) -> None: + self.log.info("Starting DistCp migration: %s -> %s", self.source_path, self.dest_path) + + self._validate_runtime_dependencies() + cmd = self._build_distcp_command() + distcp_env = self._build_distcp_env() + + CliRunner.run_process( + cmd, + env_overrides=distcp_env, + timeout=self.timeout, + retry_attempts=self.retry_attempts, + check=True, + log_output=True, + ) + self.log.info("DistCp migration completed") diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/ozone_backup.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/ozone_backup.py new file mode 100644 index 0000000000000..536b0203a4938 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/transfers/ozone_backup.py @@ -0,0 +1,83 @@ +# 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 TYPE_CHECKING + +from airflow.providers.arenadata.ozone.hooks.ozone import ( + RETRY_ATTEMPTS, + SLOW_TIMEOUT_SECONDS, + OzoneAdminHook, +) +from airflow.providers.arenadata.ozone.utils.errors import ( + SNAPSHOT_ALREADY_EXISTS_MARKERS, + OzoneCliError, + OzoneProviderError, +) +from airflow.providers.arenadata.ozone.version_compat import BaseOperator + +if TYPE_CHECKING: + from airflow.providers.arenadata.ozone.version_compat import Context + + +class OzoneBackupOperator(BaseOperator): + """Creates a snapshot of a bucket for backup and disaster recovery.""" + + template_fields = ("volume", "bucket", "snapshot_name", "ozone_conn_id") + + def __init__( + self, + *, + volume: str, + bucket: str, + snapshot_name: str, # Please, don't forget add date and time to the snapshot name. + ozone_conn_id: str = OzoneAdminHook.default_conn_name, + retry_attempts: int = RETRY_ATTEMPTS, + timeout: int = SLOW_TIMEOUT_SECONDS, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.volume = volume + self.bucket = bucket + self.snapshot_name = snapshot_name + self.ozone_conn_id = ozone_conn_id + self.retry_attempts = retry_attempts + self.timeout = timeout + + def execute(self, context: Context) -> None: + """Create snapshot and treat existing snapshot as success.""" + path = f"/{self.volume}/{self.bucket}" + cmd = ["ozone", "sh", "snapshot", "create", path, self.snapshot_name] + + hook = OzoneAdminHook( + ozone_conn_id=self.ozone_conn_id, + retry_attempts=self.retry_attempts, + ) + try: + hook.run_cli(cmd, timeout=self.timeout) + self.log.info("Snapshot created: %s (path=%s)", self.snapshot_name, path) + except OzoneCliError as e: + error_stderr = (e.stderr or "").strip() + stderr_upper = error_stderr.upper() + if any(marker in stderr_upper for marker in SNAPSHOT_ALREADY_EXISTS_MARKERS): + self.log.info( + "Snapshot %s already exists, treating as success (idempotent operation).", + self.snapshot_name, + ) + return + raise OzoneProviderError(str(e)) from e diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/__init__.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/cli_runner.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/cli_runner.py new file mode 100644 index 0000000000000..5edfcb31b68ed --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/cli_runner.py @@ -0,0 +1,561 @@ +# 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. + +"""Unified subprocess execution helpers for the Ozone provider.""" + +from __future__ import annotations + +import logging +import os +import shlex +import subprocess +import time +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import TypeAlias, TypeVar + +from tenacity import ( + before_sleep_log, + retry, + retry_if_exception, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) +from typing_extensions import ParamSpec + +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError, OzoneCliErrors +from airflow.providers.arenadata.ozone.version_compat import redact + +log = logging.getLogger(__name__) +P = ParamSpec("P") +R = TypeVar("R") +EnvOverrideValue: TypeAlias = "str | int | float | bool" +EnvOverridesMapping: TypeAlias = "Mapping[str, EnvOverrideValue]" + + +DEFAULT_PROCESS_TIMEOUT_SECONDS = 300 +DEFAULT_RETRY_WAIT = wait_exponential(multiplier=2, min=2, max=60) + + +@dataclass(frozen=True) +class ProcessOutputAnalysis: + """Immutable parsed snapshot of a completed process output.""" + + NOISE_LINE_PREFIXES = ( + "log4j:", + "usage:", + "the general command line syntax is:", + "generic options supported are:", + ) + NOISE_LINE_CONTAINS = ( + "please initialize the log4j system properly", + "see http://logging.apache.org/log4j", + ) + + raw_output: str + meaningful_lines: tuple[str, ...] + skipped_lines: tuple[str, ...] + meaningful_output: str + preferred_output: str + normalized_meaningful_output: str + normalized_preferred_output: str + + @staticmethod + def pick_process_output(process_result: subprocess.CompletedProcess[str]) -> str: + """Return output with stdout-first fallback to stderr.""" + return ((process_result.stdout or "").strip()) or ((process_result.stderr or "").strip()) + + @classmethod + def is_noise_output_line(cls, line: str) -> bool: + """Return True for non-data CLI noise lines.""" + lowered = (line or "").strip().lower() + if not lowered: + return True + if lowered.startswith(cls.NOISE_LINE_PREFIXES): + return True + return any(fragment in lowered for fragment in cls.NOISE_LINE_CONTAINS) + + @classmethod + def extract_meaningful_output_lines( + cls, + output: str | None, + *, + accept_line: Callable[[str], bool] | None = None, + ) -> tuple[list[str], list[str]]: + """ + Split output into (accepted, skipped) lines after noise filtering. + + If ``accept_line`` is not provided, all non-noise lines are accepted. + """ + if not output: + return [], [] + + accepted: list[str] = [] + skipped: list[str] = [] + for raw_line in output.splitlines(): + line = raw_line.strip() + if cls.is_noise_output_line(line): + if line: + skipped.append(line) + continue + if accept_line is None or accept_line(line): + accepted.append(line) + else: + skipped.append(line) + return accepted, skipped + + @classmethod + def from_completed_process( + cls, + process_result: subprocess.CompletedProcess[str], + *, + accept_line: Callable[[str], bool] | None = None, + ) -> ProcessOutputAnalysis: + """Build a normalized analysis snapshot for a completed process.""" + raw_output = "\n".join(part for part in (process_result.stderr, process_result.stdout) if part) + meaningful_lines, skipped_lines = cls.extract_meaningful_output_lines( + raw_output, + accept_line=accept_line, + ) + meaningful_output = "\n".join(meaningful_lines) + preferred_output = meaningful_output or cls.pick_process_output(process_result) + return cls( + raw_output=raw_output, + meaningful_lines=tuple(meaningful_lines), + skipped_lines=tuple(skipped_lines), + meaningful_output=meaningful_output, + preferred_output=preferred_output, + normalized_meaningful_output=meaningful_output.lower(), + normalized_preferred_output=preferred_output.lower(), + ) + + +class CliRunner: + """Subprocess execution helpers grouped by provider use-cases.""" + + TERMINATE_GRACE_PERIOD_SECONDS = 5 + + @staticmethod + def merge_env(env_overrides: EnvOverridesMapping | None) -> dict[str, str]: + """Merge process env with stringified overrides.""" + env = os.environ.copy() + if env_overrides: + env.update({str(key): str(value) for key, value in env_overrides.items()}) + return env + + @staticmethod + def log_streams( + stdout: str | None, + stderr: str | None, + *, + level: int = logging.DEBUG, + ) -> None: + """Log stdout/stderr when they are not empty.""" + if stdout and stdout.strip(): + log.log(level, "stdout: %s", redact(stdout.strip())) + if stderr and stderr.strip(): + log.log(level, "stderr: %s", redact(stderr.strip())) + + @classmethod + def run( + cls, + command: Sequence[str], + *, + env_overrides: EnvOverridesMapping | None = None, + timeout: int = DEFAULT_PROCESS_TIMEOUT_SECONDS, + input_text: str | None = None, + cwd: str | None = None, + check: bool = True, + log_output: bool = False, + ) -> subprocess.CompletedProcess[str]: + """Run command once without retry.""" + cmd = list(command) + masked_command = redact(shlex.join(cmd)) + log.debug("Executing command: %s", masked_command) + + process: subprocess.Popen[str] | None = None + try: + process = subprocess.Popen( + cmd, + stdin=subprocess.PIPE if input_text is not None else None, + env=cls.merge_env(env_overrides), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + cwd=cwd, + ) + stdout, stderr = process.communicate(input=input_text, timeout=timeout) + result = subprocess.CompletedProcess( + args=cmd, + returncode=process.returncode if process.returncode is not None else 1, + stdout=stdout, + stderr=stderr, + ) + if check and result.returncode != 0: + raise subprocess.CalledProcessError( + returncode=result.returncode, + cmd=cmd, + output=result.stdout, + stderr=result.stderr, + ) + except subprocess.CalledProcessError as err: + cls.log_streams(err.stdout, err.stderr, level=logging.ERROR) + log.error("Command failed with return code %s: %s", err.returncode, masked_command) + raise + except subprocess.TimeoutExpired as err: + if process is not None: + cls.terminate_process(process, masked_command=masked_command) + cls.log_streams(err.stdout, err.stderr, level=logging.ERROR) + log.error("Command timed out after %s seconds: %s", timeout, masked_command) + raise + except FileNotFoundError: + log.error("Executable not found: %s", redact(cmd[0] if cmd else "")) + raise + except OSError as err: + log.error("OS error while running command %s: %s", masked_command, redact(str(err))) + raise + except BaseException: + if process is not None: + cls.terminate_process(process, masked_command=masked_command) + raise + + if log_output: + cls.log_streams(result.stdout, result.stderr) + return result + + @classmethod + def terminate_process(cls, process: subprocess.Popen[str], *, masked_command: str) -> None: + """Gracefully stop process and force-kill if it does not exit.""" + if process.poll() is not None: + return + + process.terminate() + try: + process.wait(timeout=cls.TERMINATE_GRACE_PERIOD_SECONDS) + log.debug("Terminated process: %s", masked_command) + return + except subprocess.TimeoutExpired: + process.kill() + try: + process.wait(timeout=cls.TERMINATE_GRACE_PERIOD_SECONDS) + except subprocess.TimeoutExpired: + log.warning("Process did not exit after kill: %s", masked_command) + log.warning("Force-killed process after graceful timeout: %s", masked_command) + + @classmethod + def run_streaming( + cls, + command: Sequence[str], + *, + env: EnvOverridesMapping | None = None, + cwd: str | None = None, + check: bool = True, + on_start: Callable[[subprocess.Popen[str]], None] | None = None, + on_output: Callable[[str], None] | None = None, + ) -> str: + """Run command via ``Popen`` and stream combined output line-by-line.""" + cmd = list(command) + masked_command = redact(shlex.join(cmd)) + log.debug("Executing streaming command: %s", masked_command) + + try: + sub_process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=cwd, + close_fds=True, + env=None if env is None else {str(k): str(v) for k, v in env.items()}, + text=True, + ) + except FileNotFoundError: + log.error("Executable not found: %s", redact(cmd[0] if cmd else "")) + raise + except OSError as err: + log.error("OS error while running command %s: %s", masked_command, redact(str(err))) + raise + + if on_start is not None: + on_start(sub_process) + + if sub_process.stdout is None: + raise RuntimeError("Subprocess has no stdout") + + stdout_lines: list[str] = [] + for line in sub_process.stdout: + stdout_lines.append(line) + if on_output is not None: + on_output(line.rstrip("\n")) + + sub_process.wait() + stdout = "".join(stdout_lines) + if check and sub_process.returncode: + raise subprocess.CalledProcessError( + returncode=sub_process.returncode, + cmd=cmd, + output=stdout, + stderr=stdout, + ) + return stdout + + @classmethod + def retry_for( + cls, + *, + retry_exceptions: type[BaseException] | tuple[type[BaseException], ...] | None = None, + retry_condition: Callable[[BaseException], bool] | None = None, + logger: logging.Logger, + retry_attempts: int = 3, + ) -> Callable[[Callable[P, R]], Callable[P, R]]: + """Build tenacity retry decorator with provider defaults.""" + if retry_exceptions is None and retry_condition is None: + raise ValueError("retry_for requires retry_exceptions or retry_condition") + + if retry_condition is not None: + condition_policy = retry_if_exception(retry_condition) + retry_policy = ( + retry_if_exception_type(retry_exceptions) & condition_policy + if retry_exceptions is not None + else condition_policy + ) + else: + retry_policy = retry_if_exception_type(retry_exceptions) + + return retry( + wait=DEFAULT_RETRY_WAIT, + stop=stop_after_attempt(retry_attempts), + retry=retry_policy, + before_sleep=before_sleep_log(logger, logging.WARNING), + reraise=True, + ) + + @classmethod + def run_process( + cls, + command: Sequence[str], + *, + retry_exceptions: type[BaseException] | tuple[type[BaseException], ...] = ( + subprocess.CalledProcessError, + subprocess.TimeoutExpired, + OSError, + ), + retry_condition: Callable[[BaseException], bool] | None = None, + env_overrides: EnvOverridesMapping | None = None, + timeout: int = DEFAULT_PROCESS_TIMEOUT_SECONDS, + input_text: str | None = None, + cwd: str | None = None, + retry_attempts: int = 3, + check: bool = True, + log_output: bool = False, + ) -> subprocess.CompletedProcess[str]: + """Run generic process command with retry.""" + retry_decorator = cls.retry_for( + retry_exceptions=retry_exceptions, + retry_condition=retry_condition, + retry_attempts=retry_attempts, + logger=log, + ) + + @retry_decorator + def _execute() -> subprocess.CompletedProcess[str]: + return cls.run( + command, + env_overrides=env_overrides, + timeout=timeout, + input_text=input_text, + cwd=cwd, + check=check, + log_output=log_output, + ) + + return _execute() + + +class OzoneCliRunner(CliRunner): + """Ozone CLI specific process helpers and retry policy.""" + + @staticmethod + def _is_retryable_ozone_cli_exception(exc: BaseException) -> bool: + """Return True when Ozone CLI error should be retried.""" + return isinstance(exc, OzoneCliError) and exc.retryable + + @staticmethod + def _as_ozone_cli_error( + err: BaseException, + *, + command: Sequence[str], + timeout: int, + ) -> OzoneCliError: + """Convert low-level process errors to OzoneCliError.""" + cmd = list(command) + if isinstance(err, FileNotFoundError): + return OzoneCliError( + "Ozone CLI not found. Please ensure Ozone is installed and available in PATH.", + command=cmd, + returncode=None, + retryable=False, + ) + if isinstance(err, subprocess.TimeoutExpired): + error_message = (err.stderr or err.stdout or "Ozone command timed out").strip() + return OzoneCliError( + f"Ozone command timed out after {timeout} seconds: {redact(error_message)}", + command=cmd, + stderr=error_message, + returncode=None, + retryable=True, + ) + if isinstance(err, subprocess.CalledProcessError): + error_message = (err.stderr or err.stdout or "No error message provided").strip() + return_code = err.returncode + is_retryable = OzoneCliErrors.is_retryable_failure(return_code, error_message) + error_kind = "retryable" if is_retryable else "non-retryable" + return OzoneCliError( + f"Ozone command failed ({error_kind}, return code: {return_code}): {redact(error_message)}", + command=cmd, + stderr=error_message, + returncode=return_code, + retryable=is_retryable, + ) + if isinstance(err, OSError): + return OzoneCliError( + f"Ozone command execution failed: {redact(str(err))}", + command=cmd, + stderr=str(err), + returncode=None, + retryable=False, + ) + return OzoneCliError( + f"Ozone command execution failed: {redact(str(err))}", + command=cmd, + stderr=str(err), + returncode=None, + retryable=False, + ) + + @classmethod + def run_ozone_once( + cls, + command: Sequence[str], + *, + env_overrides: EnvOverridesMapping | None = None, + timeout: int = DEFAULT_PROCESS_TIMEOUT_SECONDS, + input_text: str | None = None, + cwd: str | None = None, + check: bool = True, + log_output: bool = False, + ) -> subprocess.CompletedProcess[str]: + """Run Ozone CLI command once without retry.""" + start = time.monotonic() + try: + result = cls.run( + command, + env_overrides=env_overrides, + timeout=timeout, + input_text=input_text, + cwd=cwd, + check=check, + log_output=log_output, + ) + elapsed = time.monotonic() - start + log.debug( + "Ozone command finished in %.2f seconds: %s", elapsed, redact(shlex.join(list(command))) + ) + return result + except (FileNotFoundError, subprocess.TimeoutExpired, subprocess.CalledProcessError, OSError) as err: + elapsed = time.monotonic() - start + log.debug("Ozone command failed in %.2f seconds: %s", elapsed, redact(shlex.join(list(command)))) + raise cls._as_ozone_cli_error(err, command=command, timeout=timeout) from err + + @classmethod + def run_ozone( + cls, + command: Sequence[str], + *, + env_overrides: EnvOverridesMapping | None = None, + timeout: int = DEFAULT_PROCESS_TIMEOUT_SECONDS, + input_text: str | None = None, + cwd: str | None = None, + retry_attempts: int = 3, + check: bool = True, + log_output: bool = True, + ) -> subprocess.CompletedProcess[str]: + """Run Ozone CLI command with retry policy and CLI-aware retry condition.""" + effective_retry_attempts = max(1, retry_attempts) + + retry_decorator = cls.retry_for( + retry_exceptions=OzoneCliError, + retry_condition=cls._is_retryable_ozone_cli_exception, + logger=log, + retry_attempts=effective_retry_attempts, + ) + + @retry_decorator + def _execute() -> subprocess.CompletedProcess[str]: + return cls.run_ozone_once( + command, + env_overrides=env_overrides, + timeout=timeout, + input_text=input_text, + cwd=cwd, + check=check, + log_output=log_output, + ) + + return _execute() + + +class KerberosCliRunner(CliRunner): + """Kerberos command execution helper.""" + + @classmethod + def run_kerberos( + cls, + command: Sequence[str], + *, + env_overrides: EnvOverridesMapping | None = None, + timeout: int = DEFAULT_PROCESS_TIMEOUT_SECONDS, + input_text: str | None = None, + cwd: str | None = None, + log_output: bool = False, + ) -> bool: + """Run command and return ``True`` on success, ``False`` on failure.""" + try: + cls.run( + command, + env_overrides=env_overrides, + timeout=timeout, + input_text=input_text, + cwd=cwd, + check=True, + log_output=log_output, + ) + return True + except subprocess.CalledProcessError as err: + log.error("Kerberos command failed with return code %s", err.returncode) + cls.log_streams(err.stdout, err.stderr, level=logging.ERROR) + return False + except subprocess.TimeoutExpired as err: + log.error("Kerberos command timed out after %s seconds", timeout) + cls.log_streams(err.stdout, err.stderr, level=logging.ERROR) + return False + except FileNotFoundError: + log.error("Kerberos executable not found") + return False + except OSError as err: + log.error("Kerberos command OS error: %s", redact(str(err))) + return False diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/connection_schema.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/connection_schema.py new file mode 100644 index 0000000000000..86fbd81cc79f2 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/connection_schema.py @@ -0,0 +1,266 @@ +# 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 dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.utils.helpers import ( + JsonDict, + TypeNormalizationHelper, +) + +if TYPE_CHECKING: + from airflow.providers.arenadata.ozone.version_compat import Connection + +# ============================================================================ +# Ozone connection schema and runtime defaults +# ============================================================================ +# This module is the single source of truth for: +# - connection schema (canonical keys) +# - defaults +# - connection UI behaviour +# - OzoneConnSnapshot typed contract +# ============================================================================ + +# --- Runtime parameters (connection schema scope) ----------------------------- +KINIT_TIMEOUT_SECONDS = 5 * 60 +CORE_SITE_XML = "core-site.xml" +OZONE_SITE_XML = "ozone-site.xml" +MAX_CONTENT_SIZE_BYTES = 1024 * 1024 * 1024 + +OZONE_CONNECTION_UI_FIELD_BEHAVIOUR: JsonDict = { + "hidden_fields": ["schema"], + "relabeling": {"host": "Ozone OM Host", "port": "Ozone OM Port"}, + "placeholders": { + "host": "ozone-om", + "port": "9862", + "extra": ( + '{"ozone_security_enabled": "true", ' + '"hadoop_security_authentication": "kerberos", "kerberos_principal": "user@REALM", ' + '"kerberos_keytab": "/opt/airflow/keytabs/user.keytab", ' + '"krb5_conf": "/opt/airflow/kerberos-config/krb5.conf", ' + '"ozone_conf_dir": "/opt/airflow/ozone-conf"}' + ), + }, +} + + +@dataclass(frozen=True) +class OzoneConnSnapshot: + """Strict typed snapshot of all supported Ozone/HDFS connection parameters.""" + + # Connection extra keys. + OZONE_SECURITY_ENABLED_KEY = "ozone_security_enabled" + HDFS_SSL_ENABLED_KEY = "hdfs_ssl_enabled" + OZONE_OM_HTTPS_PORT_KEY = "ozone_om_https_port" + OZONE_SSL_KEYSTORE_LOCATION_KEY = "ozone_ssl_keystore_location" + OZONE_SSL_KEYSTORE_PASSWORD_KEY = "ozone_ssl_keystore_password" + OZONE_SSL_KEYSTORE_TYPE_KEY = "ozone_ssl_keystore_type" + OZONE_SSL_TRUSTSTORE_LOCATION_KEY = "ozone_ssl_truststore_location" + OZONE_SSL_TRUSTSTORE_PASSWORD_KEY = "ozone_ssl_truststore_password" + OZONE_SSL_TRUSTSTORE_TYPE_KEY = "ozone_ssl_truststore_type" + DFS_ENCRYPT_DATA_TRANSFER_KEY = "dfs_encrypt_data_transfer" + HDFS_SSL_KEYSTORE_LOCATION_KEY = "hdfs_ssl_keystore_location" + HDFS_SSL_KEYSTORE_PASSWORD_KEY = "hdfs_ssl_keystore_password" + HDFS_SSL_KEYSTORE_TYPE_KEY = "hdfs_ssl_keystore_type" + HDFS_SSL_TRUSTSTORE_LOCATION_KEY = "hdfs_ssl_truststore_location" + HDFS_SSL_TRUSTSTORE_PASSWORD_KEY = "hdfs_ssl_truststore_password" + HDFS_SSL_TRUSTSTORE_TYPE_KEY = "hdfs_ssl_truststore_type" + HADOOP_SECURITY_AUTHENTICATION_KEY = "hadoop_security_authentication" + HDFS_KERBEROS_ENABLED_KEY = "hdfs_kerberos_enabled" + KERBEROS_PRINCIPAL_KEY = "kerberos_principal" + KERBEROS_KEYTAB_KEY = "kerberos_keytab" + KERBEROS_PASSWORD_KEY = "kerberos_password" + KRB5_CONF_KEY = "krb5_conf" + OZONE_CONF_DIR_KEY = "ozone_conf_dir" + HDFS_KERBEROS_PRINCIPAL_KEY = "hdfs_kerberos_principal" + HDFS_KERBEROS_KEYTAB_KEY = "hdfs_kerberos_keytab" + HDFS_KERBEROS_PASSWORD_KEY = "hdfs_kerberos_password" + HDFS_DISTCP_MAPREDUCE_LOCAL_KEY = "hdfs_distcp_mapreduce_local" + HDFS_DISTCP_RENEWER_PRINCIPAL_KEY = "hdfs_distcp_renewer_principal" + KINIT_TIMEOUT_SECONDS_KEY = "kinit_timeout_seconds" + CORE_SITE_XML_KEY = "core_site_xml" + OZONE_SITE_XML_KEY = "ozone_site_xml" + MAX_CONTENT_SIZE_BYTES_KEY = "max_content_size_bytes" + OZONE_SSL_ENV_MAPPING = ( + ("ozone_om_https_port", "OZONE_OM_HTTPS_PORT", False), + ("ozone_ssl_keystore_location", "OZONE_SSL_KEYSTORE_LOCATION", False), + ("ozone_ssl_keystore_password", "OZONE_SSL_KEYSTORE_PASSWORD", True), + ("ozone_ssl_keystore_type", "OZONE_SSL_KEYSTORE_TYPE", False), + ("ozone_ssl_truststore_location", "OZONE_SSL_TRUSTSTORE_LOCATION", False), + ("ozone_ssl_truststore_password", "OZONE_SSL_TRUSTSTORE_PASSWORD", True), + ("ozone_ssl_truststore_type", "OZONE_SSL_TRUSTSTORE_TYPE", False), + ) + HDFS_SSL_ENV_MAPPING = ( + ("dfs_encrypt_data_transfer", "DFS_ENCRYPT_DATA_TRANSFER", False), + ("hdfs_ssl_keystore_location", "HDFS_SSL_KEYSTORE_LOCATION", False), + ("hdfs_ssl_keystore_password", "HDFS_SSL_KEYSTORE_PASSWORD", True), + ("hdfs_ssl_keystore_type", "HDFS_SSL_KEYSTORE_TYPE", False), + ("hdfs_ssl_truststore_location", "HDFS_SSL_TRUSTSTORE_LOCATION", False), + ("hdfs_ssl_truststore_password", "HDFS_SSL_TRUSTSTORE_PASSWORD", True), + ("hdfs_ssl_truststore_type", "HDFS_SSL_TRUSTSTORE_TYPE", False), + ) + OZONE_KERBEROS_ENV_MAPPING = ( + ("kerberos_principal", "KERBEROS_PRINCIPAL", False), + ("kerberos_keytab", "KERBEROS_KEYTAB", True), + ("krb5_conf", "KRB5_CONFIG", False), + ("ozone_conf_dir", "OZONE_CONF_DIR", False), + ) + host: str + port: int + ozone_security_enabled: bool = False + hdfs_ssl_enabled: bool = False + ozone_om_https_port: str | None = None + ozone_ssl_keystore_location: str | None = None + ozone_ssl_keystore_password: str | None = None + ozone_ssl_keystore_type: str | None = None + ozone_ssl_truststore_location: str | None = None + ozone_ssl_truststore_password: str | None = None + ozone_ssl_truststore_type: str | None = None + dfs_encrypt_data_transfer: str | None = None + hdfs_ssl_keystore_location: str | None = None + hdfs_ssl_keystore_password: str | None = None + hdfs_ssl_keystore_type: str | None = None + hdfs_ssl_truststore_location: str | None = None + hdfs_ssl_truststore_password: str | None = None + hdfs_ssl_truststore_type: str | None = None + hadoop_security_authentication: str | None = None + hdfs_kerberos_enabled: bool = False + kerberos_principal: str | None = None + kerberos_keytab: str | None = None + kerberos_password: str | None = None + krb5_conf: str | None = None + ozone_conf_dir: str | None = None + hdfs_kerberos_principal: str | None = None + hdfs_kerberos_keytab: str | None = None + hdfs_kerberos_password: str | None = None + hdfs_distcp_mapreduce_local: bool = False + hdfs_distcp_renewer_principal: str | None = None + kinit_timeout_seconds: int = KINIT_TIMEOUT_SECONDS + core_site_xml: str = CORE_SITE_XML + ozone_site_xml: str = OZONE_SITE_XML + max_content_size_bytes: int = MAX_CONTENT_SIZE_BYTES + # Raw connection extra is exposed for DAG-level custom extensions when + # developers need to read additional non-provider keys directly. + raw_extra: JsonDict = field(default_factory=dict) + + @property + def ssl_enabled(self) -> bool: + """Return True when Ozone CLI should run with TLS enabled.""" + return self.ozone_security_enabled + + @property + def kerberos_enabled(self) -> bool: + """Return True when Ozone CLI should run with Kerberos auth enabled.""" + auth = (self.hadoop_security_authentication or "").strip().lower() + return auth == "kerberos" + + @classmethod + def from_connection( + cls, + conn: Connection, + *, + conn_id: str, + require_host_port: bool = True, + ) -> OzoneConnSnapshot: + """Create a typed snapshot from Connection model and strict canonical keys.""" + raw_host = getattr(conn, "host", None) + raw_port = getattr(conn, "port", None) + if require_host_port: + if not raw_host or not str(raw_host).strip(): + raise OzoneProviderError(f"Connection '{conn_id}' must define host for Ozone CLI operations.") + if raw_port in (None, ""): + raise OzoneProviderError(f"Connection '{conn_id}' must define port for Ozone CLI operations.") + + extra = TypeNormalizationHelper.ensure_json_dict(conn.extra_dejson) + normalize_flag_bool = TypeNormalizationHelper.normalize_flag_bool + normalize_optional_str = TypeNormalizationHelper.normalize_optional_str + parse_positive_int = TypeNormalizationHelper.parse_positive_int + return cls( + host=str(raw_host).strip() if raw_host else "", + port=int(raw_port) if raw_port not in (None, "") else 0, + ozone_security_enabled=normalize_flag_bool( + extra.get(cls.OZONE_SECURITY_ENABLED_KEY), + default=False, + ), + hdfs_ssl_enabled=normalize_flag_bool( + extra.get(cls.HDFS_SSL_ENABLED_KEY), + default=False, + ), + ozone_om_https_port=normalize_optional_str(extra.get(cls.OZONE_OM_HTTPS_PORT_KEY)), + ozone_ssl_keystore_location=normalize_optional_str( + extra.get(cls.OZONE_SSL_KEYSTORE_LOCATION_KEY) + ), + ozone_ssl_keystore_password=normalize_optional_str( + extra.get(cls.OZONE_SSL_KEYSTORE_PASSWORD_KEY) + ), + ozone_ssl_keystore_type=normalize_optional_str(extra.get(cls.OZONE_SSL_KEYSTORE_TYPE_KEY)), + ozone_ssl_truststore_location=normalize_optional_str( + extra.get(cls.OZONE_SSL_TRUSTSTORE_LOCATION_KEY) + ), + ozone_ssl_truststore_password=normalize_optional_str( + extra.get(cls.OZONE_SSL_TRUSTSTORE_PASSWORD_KEY) + ), + ozone_ssl_truststore_type=normalize_optional_str(extra.get(cls.OZONE_SSL_TRUSTSTORE_TYPE_KEY)), + dfs_encrypt_data_transfer=normalize_optional_str(extra.get(cls.DFS_ENCRYPT_DATA_TRANSFER_KEY)), + hdfs_ssl_keystore_location=normalize_optional_str(extra.get(cls.HDFS_SSL_KEYSTORE_LOCATION_KEY)), + hdfs_ssl_keystore_password=normalize_optional_str(extra.get(cls.HDFS_SSL_KEYSTORE_PASSWORD_KEY)), + hdfs_ssl_keystore_type=normalize_optional_str(extra.get(cls.HDFS_SSL_KEYSTORE_TYPE_KEY)), + hdfs_ssl_truststore_location=normalize_optional_str( + extra.get(cls.HDFS_SSL_TRUSTSTORE_LOCATION_KEY) + ), + hdfs_ssl_truststore_password=normalize_optional_str( + extra.get(cls.HDFS_SSL_TRUSTSTORE_PASSWORD_KEY) + ), + hdfs_ssl_truststore_type=normalize_optional_str(extra.get(cls.HDFS_SSL_TRUSTSTORE_TYPE_KEY)), + hadoop_security_authentication=normalize_optional_str( + extra.get(cls.HADOOP_SECURITY_AUTHENTICATION_KEY) + ), + hdfs_kerberos_enabled=normalize_flag_bool( + extra.get(cls.HDFS_KERBEROS_ENABLED_KEY), + default=False, + ), + kerberos_principal=normalize_optional_str(extra.get(cls.KERBEROS_PRINCIPAL_KEY)), + kerberos_keytab=normalize_optional_str(extra.get(cls.KERBEROS_KEYTAB_KEY)), + kerberos_password=normalize_optional_str(extra.get(cls.KERBEROS_PASSWORD_KEY)), + krb5_conf=normalize_optional_str(extra.get(cls.KRB5_CONF_KEY)), + ozone_conf_dir=normalize_optional_str(extra.get(cls.OZONE_CONF_DIR_KEY)), + hdfs_kerberos_principal=normalize_optional_str(extra.get(cls.HDFS_KERBEROS_PRINCIPAL_KEY)), + hdfs_kerberos_keytab=normalize_optional_str(extra.get(cls.HDFS_KERBEROS_KEYTAB_KEY)), + hdfs_kerberos_password=normalize_optional_str(extra.get(cls.HDFS_KERBEROS_PASSWORD_KEY)), + hdfs_distcp_mapreduce_local=normalize_flag_bool( + extra.get(cls.HDFS_DISTCP_MAPREDUCE_LOCAL_KEY), + default=False, + ), + hdfs_distcp_renewer_principal=normalize_optional_str( + extra.get(cls.HDFS_DISTCP_RENEWER_PRINCIPAL_KEY) + ), + kinit_timeout_seconds=parse_positive_int( + extra.get(cls.KINIT_TIMEOUT_SECONDS_KEY), + KINIT_TIMEOUT_SECONDS, + ), + core_site_xml=(normalize_optional_str(extra.get(cls.CORE_SITE_XML_KEY)) or CORE_SITE_XML), + ozone_site_xml=(normalize_optional_str(extra.get(cls.OZONE_SITE_XML_KEY)) or OZONE_SITE_XML), + max_content_size_bytes=parse_positive_int( + extra.get(cls.MAX_CONTENT_SIZE_BYTES_KEY), + MAX_CONTENT_SIZE_BYTES, + ), + raw_extra=dict(extra), + ) diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/errors.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/errors.py new file mode 100644 index 0000000000000..860edfcdfafb4 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/errors.py @@ -0,0 +1,127 @@ +# 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 dataclasses import dataclass + +from airflow.providers.arenadata.ozone.version_compat import AirflowException + +SNAPSHOT_ALREADY_EXISTS_MARKERS = ("FILE_ALREADY_EXISTS", "SNAPSHOT ALREADY EXISTS") + + +class OzoneProviderError(AirflowException): + """Base exception for Ozone provider runtime validation errors.""" + + +class OzoneCliError(AirflowException): + """Ozone CLI error with retryable marker.""" + + def __init__( + self, + message: str, + *, + command: list[str] | None = None, + stderr: str | None = None, + returncode: int | None = None, + retryable: bool = False, + ) -> None: + super().__init__(message) + self.command = command + self.stderr = stderr + self.returncode = returncode + self.retryable = retryable + + +class OzoneCliErrors: + """Error classifier for Ozone CLI operations.""" + + non_retryable_return_codes = ( + 2, # CLI usage / invalid arguments + 126, # command invoked cannot execute + 127, # command not found + 130, # interrupted by user (SIGINT) + ) + retryable_return_codes = ( + 143, # terminated (often timeout/orchestration kill) + 255, # common transport/session failures in CLI tools + ) + + non_retryable_errors = ( + "already_exists", + "not_found", + "does not exist", + "accessdenied", + "access denied", + "permission denied", + "invalid", + "illegalargumentexception", + "authentication failed", + "unauthorized", + ) + + retryable_errors = ( + "timed out", + "timeout", + "connection refused", + "connection reset", + "connection closed", + "no route to host", + "temporarily unavailable", + "service unavailable", + "unknownhostexception", + "could not resolve", + "sockettimeoutexception", + "connectexception", + "eofexception", + "broken pipe", + "too many requests", + "throttl", + "503", + ) + + @classmethod + def is_retryable_failure(cls, return_code: int | None, stderr: str) -> bool: + """Return True when CLI failure matches retryable return code or markers.""" + normalized = (stderr or "").lower() + if any(marker in normalized for marker in cls.non_retryable_errors): + return False + if return_code in cls.non_retryable_return_codes: + return False + if return_code in cls.retryable_return_codes: + return True + return any(marker in normalized for marker in cls.retryable_errors) + + +@dataclass(frozen=True) +class AdminResourceSpec: + """CLI markers used for idempotent admin operations.""" + + already_exists_marker: str + not_found_markers: tuple[str, ...] + + +ADMIN_RESOURCE_SPECS: dict[str, AdminResourceSpec] = { + "volume": AdminResourceSpec( + already_exists_marker="VOLUME_ALREADY_EXISTS", + not_found_markers=("VOLUME_NOT_FOUND",), + ), + "bucket": AdminResourceSpec( + already_exists_marker="BUCKET_ALREADY_EXISTS", + not_found_markers=("BUCKET_NOT_FOUND",), + ), +} diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/helpers.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/helpers.py new file mode 100644 index 0000000000000..b9efc59c06ac3 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/helpers.py @@ -0,0 +1,277 @@ +# 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 + +import fnmatch +import json +import os +import posixpath +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import TypeAlias +from urllib.parse import urlsplit, urlunsplit + +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.version_compat import redact + +JsonScalar: TypeAlias = "str | int | float | bool | None" +JsonValue: TypeAlias = "JsonScalar | list[JsonValue] | dict[str, JsonValue]" +JsonDict: TypeAlias = "dict[str, JsonValue]" + + +class URIHelper: + """URI and wildcard helpers for Ozone paths.""" + + WILDCARD_CHARS = ("*", "?", "[") + + @classmethod + def contains_wildcards(cls, value: str) -> bool: + """Return True if path contains glob or wildcard characters.""" + return bool(value) and any(ch in value for ch in cls.WILDCARD_CHARS) + + @staticmethod + def parse_ozone_uri(value: str) -> tuple[str, str, str]: + """Parse Ozone URI and return (scheme, netloc, path_or_raw).""" + parsed = urlsplit(value or "") + if parsed.scheme: + path = parsed.path or "/" + if not path.startswith("/"): + path = "/" + path + return parsed.scheme, parsed.netloc or "", path + return "", "", value or "" + + @staticmethod + def build_ozone_uri(scheme: str, netloc: str, path_or_raw: str) -> str: + """Build Ozone-style URI from components returned by parse_ozone_uri.""" + if not scheme: + return path_or_raw + path = path_or_raw or "/" + if not path.startswith("/"): + path = "/" + path + return urlunsplit((scheme, netloc, path, "", "")) + + @classmethod + def to_key_uri(cls, value: str) -> str: + """Return an object-store key URI accepted by `ozone sh key` commands.""" + scheme, netloc, path_or_raw = cls.parse_ozone_uri(value) + if scheme in {"ofs", "o3fs", "o3"}: + return cls.build_ozone_uri("o3", netloc, path_or_raw) + if scheme: + return value + return path_or_raw.lstrip("/") + + @classmethod + def split_ozone_path(cls, value: str) -> tuple[str, str]: + """ + Split Ozone path into parent directory and basename. + + Returns: + tuple[parent_path, name] + """ + target_value = (value or "").rstrip("/") + if not target_value: + return "", "" + + scheme, netloc, parsed_path = cls.parse_ozone_uri(target_value) + parent_raw, name = posixpath.split(parsed_path) + + if not parent_raw: + return "", name + + if scheme and parent_raw == "/": + return "", name + + return cls.build_ozone_uri(scheme, netloc, parent_raw), name + + @classmethod + def join_ozone_path(cls, dir_path: str, name: str) -> str: + """Join Ozone directory path and child name preserving scheme and netloc.""" + target_dir = (dir_path or "").rstrip("/") + child_name = (name or "").strip("/") + + if not target_dir: + return child_name + if not child_name: + return target_dir + + scheme, netloc, parsed_path = cls.parse_ozone_uri(target_dir) + base_path = parsed_path.rstrip("/") or ("/" if scheme else "") + full_path = posixpath.join(base_path, child_name) + return cls.build_ozone_uri(scheme, netloc, full_path) + + @classmethod + def split_ozone_wildcard_path(cls, path: str) -> tuple[str, str]: + """Split wildcard path into source directory URI and basename pattern.""" + scheme, netloc, uri_path = cls.parse_ozone_uri(path) + source_dir_path, pattern = posixpath.split(uri_path) + source_dir_path = source_dir_path or ("/" if scheme else "") + source_dir = cls.build_ozone_uri(scheme, netloc, source_dir_path or "/") + return source_dir, pattern + + @staticmethod + def filter_paths_by_basename_pattern(paths: Sequence[str], pattern: str) -> list[str]: + """Filter path list by basename using fnmatch semantics.""" + return [path for path in paths if fnmatch.fnmatch(posixpath.basename(path), pattern)] + + @classmethod + def resolve_wildcard_matches( + cls, + wildcard_path: str, + list_paths_func: Callable[[str], Sequence[str]], + ) -> tuple[str, list[str]]: + """Resolve wildcard path into source directory and matching file paths.""" + source_dir, pattern = cls.split_ozone_wildcard_path(wildcard_path) + listed_paths = list_paths_func(source_dir) + return source_dir, cls.filter_paths_by_basename_pattern(listed_paths, pattern) + + +class FileHelper: + """Filesystem access helpers for local runtime paths.""" + + @staticmethod + def is_readable_file(path: str | Path) -> bool: + """Return True when path points to an existing readable regular file.""" + target = Path(path) + return target.is_file() and os.access(target, os.R_OK) + + @staticmethod + def is_writable_target(path: str | Path) -> bool: + """ + Return True when target file can be written. + + If target exists, it must be a writable file. + If target does not exist, nearest existing parent directory must be writable/executable. + """ + target = Path(path) + if target.exists(): + return target.is_file() and os.access(target, os.W_OK) + + probe = target.parent if target.parent != Path("") else Path(".") + while not probe.exists(): + if probe == probe.parent: + return False + probe = probe.parent + return probe.is_dir() and os.access(probe, os.W_OK | os.X_OK) + + @staticmethod + def get_file_size_bytes(path: str | Path) -> int: + """Return file size in bytes for an existing local file path, or -1 when unavailable.""" + try: + return Path(path).stat().st_size + except OSError: + return -1 + + +class TypeNormalizationHelper: + """Normalization and lightweight validation helpers for typed inputs.""" + + @staticmethod + def normalize_optional_str(value: JsonValue) -> str | None: + """Return stripped string value or None for empty or None input.""" + if value is None: + return None + normalized = str(value).strip() + return normalized or None + + @classmethod + def require_optional_non_empty(cls, value: JsonValue, message: str) -> str | None: + """Validate and normalize optional string input.""" + if value is None: + return None + if not isinstance(value, str): + raise ValueError(message) + normalized = cls.normalize_optional_str(value) + if not normalized: + raise ValueError(message) + return normalized + + @staticmethod + def normalize_flag_bool(value: JsonValue, default: bool = False) -> bool: + """Normalize feature-flag style values used in connection extras/env.""" + if value is None: + return default + if isinstance(value, bool): + return value + normalized = str(value).strip().lower() + if normalized in {"1", "true", "yes", "y", "on", "enabled"}: + return True + if normalized in {"0", "false", "no", "n", "off", "disabled"}: + return False + return default + + @staticmethod + def ensure_json_dict(value: object) -> dict[str, JsonValue]: + """Return dict value when input is a JSON object-like mapping, otherwise empty dict.""" + return value if isinstance(value, dict) else {} + + @staticmethod + def parse_positive_int(value: JsonValue, default: int) -> int: + """Parse positive integer value, falling back to default when invalid.""" + if value is None: + return default + try: + parsed = int(str(value).strip()) + except (TypeError, ValueError): + return default + return parsed if parsed > 0 else default + + @staticmethod + def parse_non_negative_int(value: JsonValue) -> int | None: + """Parse non-negative integer value, returning None when invalid.""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, int): + return value if value >= 0 else None + try: + parsed = int(str(value).strip()) + except (TypeError, ValueError): + return None + return parsed if parsed >= 0 else None + + @staticmethod + def require_optional_positive_int(value: int | None, *, field_name: str) -> int | None: + """Validate optional positive int parameter used by operators.""" + if value is None: + return None + if value <= 0: + raise ValueError(f"{field_name} must be a positive integer when provided.") + return value + + @staticmethod + def parse_json_output(output: str) -> JsonValue: + """Parse CLI output that may contain plain JSON or logs followed by JSON.""" + raw_output = (output or "").strip() + if not raw_output: + raise OzoneProviderError("Empty JSON output.") + + try: + return json.loads(raw_output) + except json.JSONDecodeError: + pass + + decoder = json.JSONDecoder() + for index, char in enumerate(raw_output): + if char not in "[{": + continue + try: + parsed, _ = decoder.raw_decode(raw_output, idx=index) + return parsed + except json.JSONDecodeError: + continue + + raise OzoneProviderError(f"Failed to parse JSON output: {redact(raw_output)}") diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/security.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/security.py new file mode 100644 index 0000000000000..4bff721a689a0 --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/utils/security.py @@ -0,0 +1,561 @@ +# 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 + +import logging +import os +from dataclasses import dataclass, field +from pathlib import Path + +from airflow.configuration import ensure_secrets_loaded +from airflow.providers.arenadata.ozone.utils.cli_runner import KerberosCliRunner +from airflow.providers.arenadata.ozone.utils.connection_schema import ( + OzoneConnSnapshot, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneProviderError +from airflow.providers.arenadata.ozone.utils.helpers import FileHelper +from airflow.providers.arenadata.ozone.version_compat import AirflowException, mask_secret + +log = logging.getLogger(__name__) + + +def _normalize_and_validate_scope(scope: str) -> str: + """Normalize security scope and validate supported values.""" + normalized = str(scope).strip().lower() + if normalized not in {"ozone", "hdfs", "all"}: + raise ValueError("scope must be one of: 'ozone', 'hdfs', 'all'") + return normalized + + +class SecretResolver: + """Resolve `secret://` references via Airflow Secrets Backend.""" + + @classmethod + def get_secret_value( + cls, value: str | None, conn_id: str | None = None, *, mask: bool = True + ) -> str | None: + """ + Resolve secret from Airflow Secrets Backend or return value as-is. + + If ``value`` starts with ``secret://``, each configured backend is queried + via ``get_config(value)`` and the first non-None value is returned. + """ + if not value: + return value + + if not value.startswith("secret://"): + if mask: + mask_secret(value) + return value + + for backend in ensure_secrets_loaded(): + try: + secret_value = backend.get_config(value) + except (AttributeError, TypeError) as err: + log.error("Secrets Backend returned invalid response for %s: %s", value, str(err)) + raise ValueError(f"Failed to retrieve secret '{value}': {err}") from err + + if secret_value is not None: + resolved_value = str(secret_value) + if conn_id: + log.debug( + "Resolved secret from Secrets Backend for connection %s: %s", + conn_id, + value.split("/")[-1] if "/" in value else value, + ) + if mask: + mask_secret(resolved_value) + return resolved_value + raise ValueError(f"Secret not found: {value}") + + +@dataclass(frozen=True) +class SSLConfig: + """Structured SSL/TLS configuration split by subsystem.""" + + ozone: dict[str, str] = field(default_factory=dict) + hdfs: dict[str, str] = field(default_factory=dict) + conn_id: str | None = None + + def as_env(self) -> dict[str, str]: + """Merge all subsystem env vars into a single mapping.""" + merged: dict[str, str] = {} + merged.update(self.ozone) + merged.update(self.hdfs) + return merged + + @classmethod + def from_snapshot( + cls, snapshot: OzoneConnSnapshot, *, conn_id: str | None = None, scope: str + ) -> SSLConfig: + """Build SSL configuration from typed snapshot by explicit scope.""" + normalized_scope = _normalize_and_validate_scope(scope) + return cls( + ozone=cls._build_ozone_env(snapshot, conn_id) if normalized_scope in {"ozone", "all"} else {}, + hdfs=cls._build_hdfs_env(snapshot, conn_id) if normalized_scope in {"hdfs", "all"} else {}, + conn_id=conn_id, + ) + + @classmethod + def _build_ozone_env(cls, snapshot: OzoneConnSnapshot, conn_id: str | None) -> dict[str, str]: + """Build SSL env vars for Ozone Native CLI (ozone-site.xml mapping).""" + env: dict[str, str] = {} + if snapshot.ozone_security_enabled: + env["OZONE_SECURITY_ENABLED"] = "true" + for field_name, env_key, is_secret in OzoneConnSnapshot.OZONE_SSL_ENV_MAPPING: + value = getattr(snapshot, field_name) + if value is None: + continue + if is_secret: + value = SecretResolver.get_secret_value(str(value), conn_id=conn_id) + env[env_key] = str(value) + return env + + @classmethod + def _build_hdfs_env(cls, snapshot: OzoneConnSnapshot, conn_id: str | None) -> dict[str, str]: + """Build SSL env vars for HDFS clients (core-site.xml / hdfs-site.xml mapping).""" + env: dict[str, str] = {} + if snapshot.hdfs_ssl_enabled: + env["HDFS_SSL_ENABLED"] = "true" + for field_name, env_key, is_secret in OzoneConnSnapshot.HDFS_SSL_ENV_MAPPING: + value = getattr(snapshot, field_name) + if value is None: + continue + if is_secret: + value = SecretResolver.get_secret_value(str(value), conn_id=conn_id) + env[env_key] = str(value) + return env + + @staticmethod + def apply_ssl_env_vars( + env_vars: dict[str, str], existing_env: dict[str, str] | None = None + ) -> dict[str, str]: + """Apply SSL environment variables to existing environment.""" + if existing_env is None: + return env_vars.copy() + env = existing_env.copy() + env.update(env_vars) + return env + + +@dataclass(frozen=True) +class KerberosConfig: + """Structured Kerberos configuration split by subsystem.""" + + core: dict[str, str] = field(default_factory=dict) + hdfs: dict[str, str] = field(default_factory=dict) + + def as_env(self) -> dict[str, str]: + merged: dict[str, str] = {} + merged.update(self.core) + merged.update(self.hdfs) + return merged + + @staticmethod + def _hdfs_principal_keytab_from_snapshot( + snapshot: OzoneConnSnapshot, + conn_id: str | None, + ) -> None | tuple[str, str]: + """Return HDFS principal and resolved keytab from snapshot when configured.""" + principal_value = snapshot.hdfs_kerberos_principal + keytab_value = snapshot.hdfs_kerberos_keytab + if principal_value is None or keytab_value is None: + return None + return ( + str(principal_value), + str(SecretResolver.get_secret_value(str(keytab_value), conn_id=conn_id)), + ) + + @staticmethod + def _hdfs_principal_password_from_snapshot( + snapshot: OzoneConnSnapshot, + conn_id: str | None, + ) -> None | tuple[str, str]: + """Return HDFS principal and resolved password from snapshot when configured.""" + principal_value = snapshot.hdfs_kerberos_principal + password_value = snapshot.hdfs_kerberos_password + if principal_value is None or password_value is None: + return None + return ( + str(principal_value), + str(SecretResolver.get_secret_value(str(password_value), conn_id=conn_id)), + ) + + @classmethod + def _get_core_env(cls, snapshot: OzoneConnSnapshot, conn_id: str | None) -> dict[str, str]: + """Kerberos env for Ozone/Hadoop core tools.""" + env_vars: dict[str, str] = {} + + if str(snapshot.hadoop_security_authentication or "").lower() == "kerberos": + env_vars["HADOOP_SECURITY_AUTHENTICATION"] = "kerberos" + for field_name, env_key, is_secret in OzoneConnSnapshot.OZONE_KERBEROS_ENV_MAPPING: + value = getattr(snapshot, field_name) + if value is None: + continue + if is_secret: + value = SecretResolver.get_secret_value(str(value), conn_id=conn_id) + env_vars[env_key] = str(value) + + return env_vars + + @classmethod + def _get_hdfs_env(cls, snapshot: OzoneConnSnapshot, conn_id: str | None) -> dict[str, str]: + """Build HDFS Kerberos subprocess env without exporting passwords.""" + env_vars: dict[str, str] = {} + + if snapshot.hdfs_kerberos_enabled: + if snapshot.hdfs_kerberos_principal: + env_vars["HDFS_KERBEROS_PRINCIPAL"] = str(snapshot.hdfs_kerberos_principal) + if snapshot.hdfs_kerberos_keytab: + env_vars["HDFS_KERBEROS_KEYTAB"] = str( + SecretResolver.get_secret_value(str(snapshot.hdfs_kerberos_keytab), conn_id=conn_id) + ) + + return env_vars + + @classmethod + def get_env_vars( + cls, + snapshot: OzoneConnSnapshot, + *, + scope: str, + conn_id: str | None = None, + ) -> dict[str, str]: + """Extract Kerberos environment variables from snapshot by explicit scope.""" + normalized_scope = _normalize_and_validate_scope(scope) + config = cls( + core=cls._get_core_env(snapshot, conn_id) if normalized_scope in {"ozone", "all"} else {}, + hdfs=cls._get_hdfs_env(snapshot, conn_id) if normalized_scope in {"hdfs", "all"} else {}, + ) + return config.as_env() + + @classmethod + def kinit_with_keytab( + cls, + principal: str, + keytab: str, + krb5_conf: str | None = None, + *, + snapshot: OzoneConnSnapshot, + ) -> bool: + """ + Perform Kerberos authentication using keytab file. + + Runs system ``kinit -kt keytab principal`` with optional KRB5_CONFIG. + """ + if not principal or not keytab: + log.warning("Kerberos principal or keytab not provided, skipping kinit") + return False + + if not FileHelper.is_readable_file(keytab): + log.error("Keytab file is missing or not readable: %s", keytab) + return False + + cmd = ["kinit", "-kt", keytab, principal] + env_overrides: dict[str, str] = {} + if krb5_conf: + if not FileHelper.is_readable_file(krb5_conf): + log.error("KRB5 config file is missing or not readable: %s", krb5_conf) + return False + env_overrides["KRB5_CONFIG"] = krb5_conf + + if KerberosCliRunner.run_kerberos( + cmd, + env_overrides=env_overrides or None, + timeout=snapshot.kinit_timeout_seconds, + ): + log.info("Successfully authenticated with Kerberos: %s", principal) + return True + return False + + @classmethod + def kinit_with_password( + cls, + principal: str, + password: str, + krb5_conf: str | None = None, + *, + snapshot: OzoneConnSnapshot, + ) -> bool: + """ + Perform Kerberos authentication using a password. + + Runs system ``kinit principal`` and supplies the password through stdin. + """ + if not principal or not password: + log.warning("Kerberos principal or password not provided, skipping kinit") + return False + + mask_secret(password) + env_overrides: dict[str, str] = {} + if krb5_conf: + if not FileHelper.is_readable_file(krb5_conf): + log.error("KRB5 config file is missing or not readable: %s", krb5_conf) + return False + env_overrides["KRB5_CONFIG"] = krb5_conf + + if KerberosCliRunner.run_kerberos( + ["kinit", principal], + env_overrides=env_overrides or None, + timeout=snapshot.kinit_timeout_seconds, + input_text=f"{password}\n", + ): + log.info("Successfully authenticated with Kerberos: %s", principal) + return True + return False + + @classmethod + def kinit_from_snapshot( + cls, + *, + snapshot: OzoneConnSnapshot, + conn_id: str | None = None, + ) -> bool: + """Run kinit using credentials defined in the connection snapshot.""" + principal = snapshot.kerberos_principal + if not principal: + return False + + krb5_conf = snapshot.krb5_conf + if snapshot.kerberos_keytab: + keytab = SecretResolver.get_secret_value(str(snapshot.kerberos_keytab), conn_id=conn_id) + return cls.kinit_with_keytab(principal, str(keytab), krb5_conf, snapshot=snapshot) + + if snapshot.kerberos_password: + password = SecretResolver.get_secret_value(str(snapshot.kerberos_password), conn_id=conn_id) + return cls.kinit_with_password(principal, str(password), krb5_conf, snapshot=snapshot) + + return False + + @classmethod + def kinit_hdfs_from_snapshot( + cls, + *, + snapshot: OzoneConnSnapshot, + conn_id: str | None = None, + ) -> bool: + """Run HDFS kinit using keytab or password credentials from snapshot.""" + if not snapshot.hdfs_kerberos_enabled: + return False + + pair = cls._hdfs_principal_keytab_from_snapshot(snapshot, conn_id) + if pair: + principal, keytab = pair + return cls.kinit_with_keytab(principal, keytab, snapshot.krb5_conf, snapshot=snapshot) + + password_pair = cls._hdfs_principal_password_from_snapshot(snapshot, conn_id) + if password_pair: + principal, password = password_pair + return cls.kinit_with_password(principal, password, snapshot.krb5_conf, snapshot=snapshot) + + return False + + @staticmethod + def apply_env_vars( + env_vars: dict[str, str], existing_env: dict[str, str] | None = None + ) -> dict[str, str]: + """Build Kerberos-related environment overrides for a subprocess.""" + base_env = (existing_env if existing_env is not None else os.environ).copy() + overrides: dict[str, str] = env_vars.copy() + + enabled = ( + overrides.get( + "HADOOP_SECURITY_AUTHENTICATION", base_env.get("HADOOP_SECURITY_AUTHENTICATION", "") + ).lower() + == "kerberos" + ) + if not enabled: + return overrides + + hadoop_opts = overrides.get("HADOOP_OPTS") or base_env.get("HADOOP_OPTS", "") + hadoop_flag = "-Dhadoop.security.authentication=kerberos" + if hadoop_flag not in hadoop_opts: + hadoop_opts = (hadoop_opts + " " + hadoop_flag).strip() + overrides["HADOOP_OPTS"] = hadoop_opts + + ozone_opts = overrides.get("OZONE_OPTS") or base_env.get("OZONE_OPTS", "") + ozone_flags = [ + "-Dhadoop.security.authentication=kerberos", + "-Dozone.security.enabled=true", + ] + for flag in ozone_flags: + if flag not in ozone_opts: + ozone_opts = (ozone_opts + " " + flag).strip() + overrides["OZONE_OPTS"] = ozone_opts + + ozone_conf_dir = overrides.get("OZONE_CONF_DIR") + if ozone_conf_dir: + overrides["HADOOP_CONF_DIR"] = ozone_conf_dir + else: + base_ozone_conf_dir = base_env.get("OZONE_CONF_DIR") + if base_ozone_conf_dir: + overrides["OZONE_CONF_DIR"] = base_ozone_conf_dir + overrides["HADOOP_CONF_DIR"] = base_ozone_conf_dir + return overrides + + @classmethod + def load_ozone_env( + cls, + *, + snapshot: OzoneConnSnapshot, + conn_id: str, + ) -> dict[str, str] | None: + """Build Kerberos env for Ozone CLI with unified logging and errors.""" + try: + kerberos_env_vars = cls.get_env_vars(snapshot, scope="ozone", conn_id=conn_id) + if kerberos_env_vars: + kerberos_env = cls.apply_env_vars(kerberos_env_vars) + log.debug("Kerberos configuration loaded from connection: %s", list(kerberos_env_vars.keys())) + if str(snapshot.hadoop_security_authentication or "").lower() == "kerberos": + log.debug("Kerberos authentication enabled for Ozone Native CLI") + return kerberos_env + log.debug("No Kerberos configuration found in connection Extra") + return None + except AirflowException as err: + log.debug("Could not load Kerberos configuration (connection may not exist): %s", str(err)) + return None + except ValueError as err: + raise OzoneProviderError( + f"Invalid Kerberos configuration in connection '{conn_id}': {err}" + ) from err + + @classmethod + def load_hdfs_env( + cls, + *, + snapshot: OzoneConnSnapshot, + conn_id: str, + ) -> dict[str, str] | None: + """Build HDFS Kerberos subprocess env with unified logging and errors.""" + try: + kerberos_env_vars = cls.get_env_vars(snapshot, scope="hdfs", conn_id=conn_id) + if kerberos_env_vars: + log.debug( + "HDFS Kerberos configuration loaded from connection: %s", + list(kerberos_env_vars.keys()), + ) + return kerberos_env_vars + log.debug("No HDFS Kerberos configuration found in connection Extra") + return None + except AirflowException as err: + log.debug("Could not load HDFS Kerberos configuration (connection may not exist): %s", str(err)) + return None + except ValueError as err: + raise OzoneProviderError( + f"Invalid HDFS Kerberos configuration in connection '{conn_id}': {err}" + ) from err + + @classmethod + def ensure_ticket( + cls, + *, + snapshot: OzoneConnSnapshot, + conn_id: str, + kerberos_ticket_ready: bool, + ) -> bool: + """Ensure Kerberos ticket is initialized; returns updated readiness flag.""" + if kerberos_ticket_ready and cls.has_valid_ticket(snapshot=snapshot): + return True + if kerberos_ticket_ready: + log.warning( + "Cached Kerberos ticket flag is set but no valid ticket found; re-initializing for connection '%s'", + conn_id, + ) + if not snapshot.kerberos_enabled: + return False + + principal = snapshot.kerberos_principal + keytab_configured = bool(snapshot.kerberos_keytab) + password_configured = bool(snapshot.kerberos_password) + if not principal: + if snapshot.kerberos_enabled: + raise OzoneProviderError( + f"Kerberos authentication is enabled for connection '{conn_id}' " + "but 'kerberos_principal' is missing." + ) + return False + if not keytab_configured and not password_configured: + if snapshot.kerberos_enabled: + raise OzoneProviderError( + f"Kerberos authentication is enabled for connection '{conn_id}' " + "but neither 'kerberos_keytab' nor 'kerberos_password' is configured." + ) + return False + + try: + authenticated = cls.kinit_from_snapshot( + snapshot=snapshot, + conn_id=conn_id, + ) + except ValueError as err: + raise OzoneProviderError( + f"Invalid Kerberos configuration in connection '{conn_id}': {err}" + ) from err + if not authenticated: + raise OzoneProviderError( + f"Kerberos authentication failed for connection '{conn_id}' using principal '{principal}'." + ) + log.debug("Kerberos ticket is ready for connection '%s'", conn_id) + return True + + @staticmethod + def is_enabled(kerberos_env: dict[str, str] | None) -> bool: + """Return True if Kerberos is effectively enabled.""" + if not kerberos_env: + return False + return str(kerberos_env.get("HADOOP_SECURITY_AUTHENTICATION", "")).lower() == "kerberos" + + @staticmethod + def resolve_config_dir(kerberos_env: dict[str, str] | None) -> str | None: + """Return effective config dir for Kerberos-enabled Ozone CLI.""" + if not kerberos_env: + return None + return kerberos_env.get("OZONE_CONF_DIR") + + @classmethod + def has_valid_ticket(cls, *, snapshot: OzoneConnSnapshot) -> bool: + """Return True when current default Kerberos ticket cache has non-expired credentials.""" + return KerberosCliRunner.run_kerberos( + ["klist", "-s"], + timeout=snapshot.kinit_timeout_seconds, + log_output=False, + ) + + @staticmethod + def check_config_files_exist( + config_dir: str, + *, + snapshot: OzoneConnSnapshot, + ) -> bool: + """Return True when both core-site.xml and ozone-site.xml exist.""" + if not config_dir: + log.debug("Config directory is None or empty") + return False + + config_path = Path(config_dir) + core_site = config_path / snapshot.core_site_xml + ozone_site = config_path / snapshot.ozone_site_xml + core_exists = FileHelper.is_readable_file(core_site) + ozone_exists = FileHelper.is_readable_file(ozone_site) + if not core_exists: + log.debug("core-site.xml not found or not readable: %s", core_site) + if not ozone_exists: + log.debug("ozone-site.xml not found or not readable: %s", ozone_site) + if core_exists and ozone_exists: + log.debug("Both configuration files found in %s", config_dir) + return core_exists and ozone_exists diff --git a/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/version_compat.py b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/version_compat.py new file mode 100644 index 0000000000000..3ab4c94f1292f --- /dev/null +++ b/providers/arenadata/ozone/src/airflow/providers/arenadata/ozone/version_compat.py @@ -0,0 +1,49 @@ +# 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 airflow.providers.common.compat.version_compat import AIRFLOW_V_3_0_PLUS + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import Connection + from airflow.sdk._shared.secrets_masker import mask_secret, redact + from airflow.sdk.bases.hook import BaseHook + from airflow.sdk.bases.operator import BaseOperator + from airflow.sdk.bases.sensor import BaseSensorOperator + from airflow.sdk.definitions.context import Context +else: + from airflow.hooks.base import BaseHook + from airflow.models import BaseOperator + from airflow.models.connection import Connection + from airflow.sensors.base import BaseSensorOperator + from airflow.utils.context import Context + from airflow.utils.log.secrets_masker import mask_secret, redact + +from airflow.providers.common.compat.sdk import AirflowException + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "AirflowException", + "BaseHook", + "BaseOperator", + "BaseSensorOperator", + "Connection", + "Context", + "mask_secret", + "redact", +] diff --git a/providers/arenadata/ozone/tests/conftest.py b/providers/arenadata/ozone/tests/conftest.py new file mode 100644 index 0000000000000..f56ccce0a3f69 --- /dev/null +++ b/providers/arenadata/ozone/tests/conftest.py @@ -0,0 +1,19 @@ +# 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 + +pytest_plugins = "tests_common.pytest_plugin" diff --git a/providers/arenadata/ozone/tests/integration/__init__.py b/providers/arenadata/ozone/tests/integration/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/__init__.py @@ -0,0 +1,17 @@ +# 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. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/arenadata/ozone/tests/integration/arenadata/__init__.py b/providers/arenadata/ozone/tests/integration/arenadata/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/__init__.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/__init__.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/test_ozone.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/test_ozone.py new file mode 100644 index 0000000000000..b7fb38dfc9a9a --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/hooks/test_ozone.py @@ -0,0 +1,328 @@ +# +# 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 + +import os +import tempfile +from pathlib import Path, PurePosixPath + +import pytest + +from airflow.exceptions import AirflowException +from airflow.models.connection import Connection +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneAdminHook, OzoneFsHook +from airflow.utils.db import merge_conn + +CONN_ID = "ozone_test" +OZONE_HOST = os.environ.get("OZONE_HOST", "om") +OZONE_PORT = int(os.environ.get("OZONE_PORT", "9862")) +OZONE_FS_AUTHORITY = os.environ.get("OZONE_FS_AUTHORITY", "om") +VOLUME = "inttest-volume" +BUCKET = "inttest-bucket" +EDGE_VOLUME = "inttest-edge-volume" +EDGE_BUCKET = "inttest-edge-bucket" +EDGE_MISSING_VOLUME = "inttest-edge-missing-volume" +EDGE_MISSING_BUCKET = "inttest-edge-missing-bucket" +EDGE_MISSING_KEY = "missing/to_delete.txt" +EDGE_EXISTING_KEY = "existing-policy.txt" + + +@pytest.fixture(autouse=True) +def _setup_connection(): + merge_conn( + Connection( + conn_id=CONN_ID, + conn_type="ozone", + host=OZONE_HOST, + port=OZONE_PORT, + ) + ) + + +def _cleanup_volume(hook: OzoneAdminHook) -> None: + """Delete test volume with all buckets.""" + _cleanup_named_volume(hook, VOLUME, BUCKET) + + +def _cleanup_edge_volume(hook: OzoneAdminHook) -> None: + """Delete edge-case test volume with all buckets.""" + _cleanup_named_volume(hook, EDGE_VOLUME, EDGE_BUCKET) + + +def _cleanup_named_volume(hook: OzoneAdminHook, volume: str, bucket: str) -> None: + """Delete named test volume with all buckets.""" + if not hook.volume_exists(volume): + return + if hook.bucket_exists(volume, bucket): + try: + hook.run_cli( + [ + "ozone", + "fs", + "-rm", + "-r", + "-skipTrash", + f"ofs://{OZONE_FS_AUTHORITY}/{volume}/{bucket}/*", + ], + check=False, + log_output=False, + retry_attempts=0, + ) + except Exception: + pass + try: + hook.delete_bucket(volume, bucket) + except Exception: + pass + try: + hook.delete_volume(volume) + except Exception: + pass + + +def _edge_key_path(key: str) -> str: + return f"ofs://{OZONE_FS_AUTHORITY}/{PurePosixPath(EDGE_VOLUME, EDGE_BUCKET, key)}" + + +@pytest.mark.integration("ozone") +class TestOzoneAdminHookIntegration: + def setup_method(self): + self.hook = OzoneAdminHook(ozone_conn_id=CONN_ID, retry_attempts=1) + _cleanup_volume(self.hook) + + def teardown_method(self): + _cleanup_volume(self.hook) + + def test_create_and_delete_volume(self): + self.hook.create_volume(VOLUME) + assert self.hook.volume_exists(VOLUME) + + self.hook.delete_volume(VOLUME) + assert not self.hook.volume_exists(VOLUME) + + def test_create_volume_existing_target_policy(self): + self.hook.create_volume(VOLUME) + + self.hook.create_volume(VOLUME, if_exists="ignore") + + with pytest.raises(AirflowException, match="already exists"): + self.hook.create_volume(VOLUME, if_exists="error") + + def test_volume_not_exists(self): + assert not self.hook.volume_exists("nonexistent_volume") + + def test_list_volumes(self): + self.hook.create_volume(VOLUME) + volumes = self.hook.list_volumes() + volume_names = [v["name"] for v in volumes] + assert VOLUME in volume_names + + def test_get_volume_info(self): + self.hook.create_volume(VOLUME) + info = self.hook.get_volume_info(VOLUME) + assert info["name"] == VOLUME + + def test_create_and_delete_bucket(self): + self.hook.create_volume(VOLUME) + self.hook.create_bucket(VOLUME, BUCKET) + assert self.hook.bucket_exists(VOLUME, BUCKET) + + self.hook.delete_bucket(VOLUME, BUCKET) + assert not self.hook.bucket_exists(VOLUME, BUCKET) + + def test_list_buckets(self): + self.hook.create_volume(VOLUME) + self.hook.create_bucket(VOLUME, BUCKET) + buckets = self.hook.list_buckets(VOLUME) + bucket_names = [b["name"] for b in buckets] + assert BUCKET in bucket_names + + def test_get_bucket_info(self): + self.hook.create_volume(VOLUME) + self.hook.create_bucket(VOLUME, BUCKET) + info = self.hook.get_bucket_info(VOLUME, BUCKET) + assert info["name"] == BUCKET + assert info["volumeName"] == VOLUME + + +@pytest.mark.integration("ozone") +class TestOzoneFsHookIntegration: + def setup_method(self): + self.admin = OzoneAdminHook(ozone_conn_id=CONN_ID, retry_attempts=1) + self.hook = OzoneFsHook(ozone_conn_id=CONN_ID, retry_attempts=1) + _cleanup_volume(self.admin) + self.admin.create_volume(VOLUME) + self.admin.create_bucket(VOLUME, BUCKET, replication_type="RATIS", replication="ONE") + self.base_path = f"ofs://{OZONE_FS_AUTHORITY}/{VOLUME}/{BUCKET}" + + def teardown_method(self): + _cleanup_volume(self.admin) + + def test_create_and_check_path(self): + path = f"{self.base_path}/test_dir" + self.hook.make_path(path) + assert self.hook.path_exists(path) + + def test_create_path_existing_target_policy(self): + path = f"{self.base_path}/existing_dir" + self.hook.make_path(path) + + self.hook.make_path(path, if_exists="ignore") + + with pytest.raises(AirflowException, match="already exists"): + self.hook.make_path(path, if_exists="error") + + def test_upload_and_read_key(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("hello ozone") + local_path = f.name + + try: + remote_path = f"{self.base_path}/test_file.txt" + self.hook.upload_key(local_path, remote_path) + assert self.hook.key_exists(remote_path) + + content = self.hook.read_text(remote_path) + assert "hello ozone" in content + finally: + os.unlink(local_path) + + def test_upload_key_existing_target_policy(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("first") + first_local_path = f.name + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("second") + second_local_path = f.name + + try: + remote_path = f"{self.base_path}/existing_upload.txt" + self.hook.upload_key(first_local_path, remote_path) + + with pytest.raises(AirflowException, match="already exists"): + self.hook.upload_key(second_local_path, remote_path, if_exists="error") + + self.hook.upload_key(second_local_path, remote_path, if_exists="ignore") + assert self.hook.read_text(remote_path) == "first" + + self.hook.upload_key(second_local_path, remote_path, if_exists="overwrite") + assert self.hook.read_text(remote_path) == "second" + finally: + os.unlink(first_local_path) + os.unlink(second_local_path) + + def test_download_key(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("download test") + local_path = f.name + + download_path = local_path + ".downloaded" + try: + remote_path = f"{self.base_path}/download_test.txt" + self.hook.upload_key(local_path, remote_path) + + self.hook.download_key(remote_path, download_path) + with open(download_path) as f: + assert f.read() == "download test" + finally: + os.unlink(local_path) + if os.path.exists(download_path): + os.unlink(download_path) + + def test_delete_key(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("to delete") + local_path = f.name + + try: + remote_path = f"{self.base_path}/to_delete.txt" + self.hook.upload_key(local_path, remote_path) + assert self.hook.key_exists(remote_path) + + self.hook.delete_key(remote_path) + assert not self.hook.key_exists(remote_path) + finally: + os.unlink(local_path) + + def test_list_keys(self): + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("list test") + local_path = f.name + + try: + for i in range(3): + self.hook.upload_key(local_path, f"{self.base_path}/list_{i}.txt") + + keys = self.hook.list_keys(self.base_path) + assert len(keys) == 3 + finally: + os.unlink(local_path) + + +@pytest.mark.integration("ozone") +class TestOzoneHookEdgeCaseIntegration: + def setup_method(self): + self.admin = OzoneAdminHook(ozone_conn_id=CONN_ID, retry_attempts=1) + self.fs = OzoneFsHook(ozone_conn_id=CONN_ID, retry_attempts=1) + _cleanup_edge_volume(self.admin) + self.admin.create_volume(EDGE_VOLUME) + self.admin.create_bucket(EDGE_VOLUME, EDGE_BUCKET, replication_type="RATIS", replication="ONE") + + def teardown_method(self): + _cleanup_edge_volume(self.admin) + + def test_missing_fs_path_checks_return_false(self): + path = _edge_key_path(EDGE_MISSING_KEY) + + assert self.fs.exists(path) is False + assert self.fs.path_exists(path) is False + assert self.fs.key_exists(path) is False + + def test_delete_missing_fs_key_is_idempotent(self): + path = _edge_key_path(EDGE_MISSING_KEY) + + self.fs.delete_key(path) + + assert self.fs.key_exists(path) is False + + def test_missing_bucket_checks_and_delete_are_idempotent(self): + assert self.admin.bucket_exists(EDGE_VOLUME, EDGE_MISSING_BUCKET) is False + + self.admin.delete_bucket(EDGE_VOLUME, EDGE_MISSING_BUCKET) + + assert self.admin.bucket_exists(EDGE_VOLUME, EDGE_MISSING_BUCKET) is False + + def test_missing_volume_checks_and_delete_are_idempotent(self): + assert self.admin.volume_exists(EDGE_MISSING_VOLUME) is False + + self.admin.delete_volume(EDGE_MISSING_VOLUME) + + assert self.admin.volume_exists(EDGE_MISSING_VOLUME) is False + + def test_existing_fs_upload_policy_fails_or_skips_without_overwrite(self): + path = _edge_key_path(EDGE_EXISTING_KEY) + self.fs.create_key(path, if_exists="ignore") + + with tempfile.TemporaryDirectory(prefix="ozone_existing_policy_") as tmp_dir: + local_path = Path(tmp_dir) / "payload.txt" + local_path.write_text("existing object policy probe\n", encoding="utf-8") + + with pytest.raises(AirflowException, match="already exists"): + self.fs.upload_key(str(local_path), path, if_exists="error", timeout=60) + + self.fs.upload_key(str(local_path), path, if_exists="ignore", timeout=60) diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/__init__.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/test_ozone.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/test_ozone.py new file mode 100644 index 0000000000000..391671c5ba61d --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/operators/test_ozone.py @@ -0,0 +1,194 @@ +# +# 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 + +import os +import tempfile +from pathlib import Path + +import pytest + +from airflow.models.connection import Connection +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneAdminHook, OzoneFsHook +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneDeleteBucketOperator, + OzoneDeletePathOperator, + OzoneDeleteVolumeOperator, + OzoneDownloadFileOperator, + OzoneListOperator, + OzonePathExistsOperator, + OzoneUploadContentOperator, +) +from airflow.utils.db import merge_conn + +CONN_ID = "ozone_operator_test" +OZONE_HOST = os.environ.get("OZONE_HOST", "om") +OZONE_PORT = int(os.environ.get("OZONE_PORT", "9862")) +OZONE_FS_AUTHORITY = os.environ.get("OZONE_FS_AUTHORITY", "om") +VOLUME = "inttest-operator-volume" +BUCKET = "inttest-operator-bucket" + + +@pytest.fixture(autouse=True) +def _setup_connection(): + merge_conn( + Connection( + conn_id=CONN_ID, + conn_type="ozone", + host=OZONE_HOST, + port=OZONE_PORT, + ) + ) + + +def _cleanup_volume(hook: OzoneAdminHook) -> None: + """Delete test volume with all buckets.""" + if not hook.volume_exists(VOLUME): + return + if hook.bucket_exists(VOLUME, BUCKET): + try: + hook.run_cli( + [ + "ozone", + "fs", + "-rm", + "-r", + "-skipTrash", + f"ofs://{OZONE_FS_AUTHORITY}/{VOLUME}/{BUCKET}/*", + ], + check=False, + log_output=False, + retry_attempts=0, + ) + except Exception: + pass + try: + hook.delete_bucket(VOLUME, BUCKET) + except Exception: + pass + try: + hook.delete_volume(VOLUME) + except Exception: + pass + + +@pytest.mark.integration("ozone") +class TestOzoneOperatorIntegration: + def setup_method(self): + self.admin = OzoneAdminHook(ozone_conn_id=CONN_ID, retry_attempts=1) + self.fs = OzoneFsHook(ozone_conn_id=CONN_ID, retry_attempts=1) + _cleanup_volume(self.admin) + + def teardown_method(self): + _cleanup_volume(self.admin) + + def test_admin_and_filesystem_operator_smoke_flow(self): + base_path = f"ofs://{OZONE_FS_AUTHORITY}/{VOLUME}/{BUCKET}" + directory = f"{base_path}/incoming" + remote_file = f"{directory}/payload.txt" + + OzoneCreateVolumeOperator( + task_id="create_volume", + volume_name=VOLUME, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert self.admin.volume_exists(VOLUME) + + OzoneCreateBucketOperator( + task_id="create_bucket", + volume_name=VOLUME, + bucket_name=BUCKET, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert self.admin.bucket_exists(VOLUME, BUCKET) + + OzoneCreatePathOperator( + task_id="create_path", + path=directory, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert self.fs.path_exists(directory) + + OzoneUploadContentOperator( + task_id="upload_content", + content="hello from operator integration\n", + remote_path=remote_file, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert self.fs.key_exists(remote_file) + + assert OzonePathExistsOperator( + task_id="path_exists", + path=remote_file, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + + listed_paths = OzoneListOperator( + task_id="list_keys", + path=f"{directory}/*", + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert remote_file in listed_paths + + with tempfile.TemporaryDirectory(prefix="ozone_operator_download_") as tmp_dir: + download_path = Path(tmp_dir) / "payload.txt" + OzoneDownloadFileOperator( + task_id="download_file", + remote_path=remote_file, + local_path=str(download_path), + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert download_path.read_text(encoding="utf-8") == "hello from operator integration\n" + + OzoneDeletePathOperator( + task_id="delete_path", + path=directory, + ozone_conn_id=CONN_ID, + recursive=True, + retry_attempts=1, + ).execute({}) + assert not self.fs.path_exists(directory) + + OzoneDeleteBucketOperator( + task_id="delete_bucket", + volume_name=VOLUME, + bucket_name=BUCKET, + recursive=True, + force=True, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert not self.admin.bucket_exists(VOLUME, BUCKET) + + OzoneDeleteVolumeOperator( + task_id="delete_volume", + volume_name=VOLUME, + ozone_conn_id=CONN_ID, + retry_attempts=1, + ).execute({}) + assert not self.admin.volume_exists(VOLUME) diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/__init__.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/test_ozone.py b/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/test_ozone.py new file mode 100644 index 0000000000000..e5726a68b7e7b --- /dev/null +++ b/providers/arenadata/ozone/tests/integration/arenadata/ozone/sensors/test_ozone.py @@ -0,0 +1,130 @@ +# +# 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 + +import os +import tempfile + +import pytest + +from airflow.models.connection import Connection +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneAdminHook, OzoneFsHook +from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor +from airflow.utils.db import merge_conn + +CONN_ID = "ozone_sensor_test" +OZONE_HOST = os.environ.get("OZONE_HOST", "om") +OZONE_PORT = int(os.environ.get("OZONE_PORT", "9862")) +OZONE_FS_AUTHORITY = os.environ.get("OZONE_FS_AUTHORITY", "om") +VOLUME = "inttest-sensor-volume" +BUCKET = "inttest-sensor-bucket" + + +@pytest.fixture(autouse=True) +def _setup_connection(): + merge_conn( + Connection( + conn_id=CONN_ID, + conn_type="ozone", + host=OZONE_HOST, + port=OZONE_PORT, + ) + ) + + +def _cleanup_volume(hook: OzoneAdminHook) -> None: + """Delete test volume with all buckets.""" + if not hook.volume_exists(VOLUME): + return + if hook.bucket_exists(VOLUME, BUCKET): + try: + hook.run_cli( + [ + "ozone", + "fs", + "-rm", + "-r", + "-skipTrash", + f"ofs://{OZONE_FS_AUTHORITY}/{VOLUME}/{BUCKET}/*", + ], + check=False, + log_output=False, + retry_attempts=0, + ) + except Exception: + pass + try: + hook.delete_bucket(VOLUME, BUCKET) + except Exception: + pass + try: + hook.delete_volume(VOLUME) + except Exception: + pass + + +@pytest.mark.integration("ozone") +class TestOzoneKeySensorIntegration: + def setup_method(self): + self.admin = OzoneAdminHook(ozone_conn_id=CONN_ID, retry_attempts=1) + self.fs = OzoneFsHook(ozone_conn_id=CONN_ID, retry_attempts=1) + _cleanup_volume(self.admin) + self.admin.create_volume(VOLUME) + self.admin.create_bucket(VOLUME, BUCKET, replication_type="RATIS", replication="ONE") + self.base_path = f"ofs://{OZONE_FS_AUTHORITY}/{VOLUME}/{BUCKET}" + + def teardown_method(self): + _cleanup_volume(self.admin) + + def _upload_text(self, remote_path: str, content: str = "sensor payload\n") -> None: + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as file: + file.write(content) + local_path = file.name + try: + self.fs.upload_key(local_path, remote_path) + finally: + os.unlink(local_path) + + def test_poke_returns_true_for_existing_key(self): + remote_path = f"{self.base_path}/ready.txt" + self._upload_text(remote_path) + + sensor = OzoneKeySensor( + task_id="wait_for_ready", + path=remote_path, + ozone_conn_id=CONN_ID, + retry_attempts=1, + timeout=30, + ) + + assert sensor.poke({}) is True + assert sensor.cli_timeout == 30 + assert sensor.timeout != 30 + + def test_poke_returns_false_for_missing_key(self): + sensor = OzoneKeySensor( + task_id="wait_for_missing", + path=f"{self.base_path}/missing.txt", + ozone_conn_id=CONN_ID, + retry_attempts=1, + timeout=30, + ) + + assert sensor.poke({}) is False + assert sensor.cli_timeout == 30 + assert sensor.timeout != 30 diff --git a/providers/arenadata/ozone/tests/unit/__init__.py b/providers/arenadata/ozone/tests/unit/__init__.py new file mode 100644 index 0000000000000..5966d6b1d5261 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/__init__.py @@ -0,0 +1,17 @@ +# 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. +__path__ = __import__("pkgutil").extend_path(__path__, __name__) diff --git a/providers/arenadata/ozone/tests/unit/arenadata/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/test_ozone.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/test_ozone.py new file mode 100644 index 0000000000000..623b6c5053801 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/hooks/test_ozone.py @@ -0,0 +1,566 @@ +# 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 + +import subprocess +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.arenadata.ozone.hooks.ozone import ( + ExistingTargetPolicy, + OzoneAdminExtraHook, + OzoneAdminHook, + OzoneCliHook, + OzoneFsHook, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError + +MOCK_CLI_PATH = "airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliHook.run_cli" +MOCK_RUN_PATH = "airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliRunner.run_ozone_once" +MOCK_RUN_RETRY_PATH = "airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliRunner.run_ozone" + + +@pytest.fixture +def admin_hook(): + hook = OzoneAdminHook(ozone_conn_id="test_admin_conn") + conn = MagicMock() + conn.host = "ozone-om" + conn.port = 9862 + conn.extra_dejson = {} + hook.get_connection = lambda _: conn + return hook + + +@pytest.fixture +def admin_extra_hook(): + hook = OzoneAdminExtraHook(ozone_conn_id="test_admin_conn") + conn = MagicMock() + conn.host = "ozone-om" + conn.port = 9862 + conn.extra_dejson = {} + hook.get_connection = lambda _: conn + return hook + + +@pytest.fixture +def ozone_fs_hook(): + hook = OzoneFsHook(ozone_conn_id="test_conn") + conn = MagicMock() + conn.host = "ozone-om" + conn.port = 9862 + conn.extra_dejson = {} + hook.get_connection = lambda _: conn + return hook + + +class TestOzoneCliHookConnectionSnapshot: + def test_connection_snapshot_requires_host(self): + conn = MagicMock() + conn.host = None + conn.port = 9862 + conn.extra_dejson = {} + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + with pytest.raises(AirflowException, match="must define host"): + _ = hook.connection_snapshot + + def test_connection_snapshot_requires_port(self): + conn = MagicMock() + conn.host = "om-host" + conn.port = None + conn.extra_dejson = {} + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + with pytest.raises(AirflowException, match="must define port"): + _ = hook.connection_snapshot + + def test_prepared_cli_env_uses_ozone_conf_dir_from_connection_extra(self): + conn = MagicMock() + conn.host = "om-host" + conn.port = 9862 + conn.extra_dejson = {"ozone_conf_dir": "/opt/airflow/ozone-conf"} + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + + env = hook._prepared_cli_env() + assert env["OZONE_CONF_DIR"] == "/opt/airflow/ozone-conf" + assert env["HADOOP_CONF_DIR"] == "/opt/airflow/ozone-conf" + + def test_prepared_cli_env_uses_ozone_conf_dir_from_environment(self, monkeypatch, caplog): + conn = MagicMock() + conn.host = "om-host" + conn.port = 9862 + conn.extra_dejson = {} + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + monkeypatch.setenv("OZONE_CONF_DIR", "/env/ozone-conf") + + with caplog.at_level("INFO"): + env = hook._prepared_cli_env() + + assert env["OZONE_CONF_DIR"] == "/env/ozone-conf" + assert env["HADOOP_CONF_DIR"] == "/env/ozone-conf" + assert "Using ozone_conf_dir from OZONE_CONF_DIR environment variable" in caplog.text + + def test_prepared_cli_env_raises_when_kerberos_conf_dir_missing(self, monkeypatch): + conn = MagicMock() + conn.host = "om-host" + conn.port = 9862 + conn.extra_dejson = {"hadoop_security_authentication": "kerberos"} + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + monkeypatch.delenv("OZONE_CONF_DIR", raising=False) + monkeypatch.delenv("HADOOP_CONF_DIR", raising=False) + + with pytest.raises( + AirflowException, match="Kerberos is enabled but ozone_conf_dir is not configured" + ): + hook._prepared_cli_env() + + @pytest.mark.parametrize( + ("extra", "expected_mode"), + [ + ({}, "plain"), + ({"ozone_security_enabled": "true"}, "ssl"), + ( + { + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "airflow@EXAMPLE.COM", + "kerberos_keytab": "/tmp/airflow.keytab", + "krb5_conf": "/etc/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf", + }, + "kerberos", + ), + ( + { + "ozone_security_enabled": "true", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "airflow@EXAMPLE.COM", + "kerberos_keytab": "/tmp/airflow.keytab", + "krb5_conf": "/etc/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf", + }, + "ssl+kerberos", + ), + ], + ) + @patch("airflow.providers.arenadata.ozone.hooks.ozone.OzoneCliRunner.run_ozone") + @patch("airflow.providers.arenadata.ozone.hooks.ozone.KerberosConfig.ensure_ticket") + def test_run_cli_logs_runtime_mode( + self, + mock_ensure_ticket: MagicMock, + mock_run_ozone: MagicMock, + extra: dict[str, str], + expected_mode: str, + caplog, + ) -> None: + conn = MagicMock() + conn.host = "om-host" + conn.port = 9862 + conn.extra_dejson = extra + hook = OzoneCliHook(ozone_conn_id="ozone_default") + hook.get_connection = lambda _: conn + + mock_ensure_ticket.return_value = False + mock_run_ozone.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "volume", "list", "/"], + returncode=0, + stdout="[]", + stderr="", + ) + + with caplog.at_level("INFO"): + hook.run_cli(["ozone", "sh", "volume", "list", "/"], log_output=False) + + assert f"mode: {expected_mode}" in caplog.text + + +class TestOzoneAdminHook: + @patch(MOCK_RUN_RETRY_PATH) + def test_create_volume_uses_run_with_retry( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = MagicMock(returncode=0, stdout="", stderr="") + admin_hook.create_volume(volume_name="test_vol") + mock_run_with_retry.assert_called_once() + assert mock_run_with_retry.call_args.args[0] == ["ozone", "sh", "volume", "create", "/test_vol"] + + @patch(MOCK_RUN_RETRY_PATH) + def test_create_volume_already_exists_is_idempotent( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "volume", "create", "/test_vol"], + returncode=255, + stdout="", + stderr=( + "log4j:WARN No appenders could be found for logger (org.apache.hadoop.util.Shell).\n" + "VOLUME_ALREADY_EXISTS Volume already exists" + ), + ) + admin_hook.create_volume(volume_name="test_vol") + mock_run_with_retry.assert_called_once() + + @patch(MOCK_RUN_RETRY_PATH) + def test_create_volume_already_exists_can_fail_fast( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "volume", "create", "/test_vol"], + returncode=255, + stdout="", + stderr="VOLUME_ALREADY_EXISTS Volume already exists", + ) + with pytest.raises(AirflowException, match="Volume already exists"): + admin_hook.create_volume(volume_name="test_vol", if_exists="error") + mock_run_with_retry.assert_called_once() + + @patch(MOCK_RUN_RETRY_PATH) + def test_create_bucket_raises_on_non_idempotent_failure( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "bucket", "create", "/test_vol/test_bkt", "--space-quota", "100GB"], + returncode=1, + stdout="", + stderr="ACCESS_DENIED", + ) + with pytest.raises(AirflowException, match="Ozone command failed"): + admin_hook.create_bucket(volume_name="test_vol", bucket_name="test_bkt", quota="100GB") + + @patch(MOCK_RUN_RETRY_PATH) + def test_delete_bucket_missing_is_idempotent( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "bucket", "delete", "/test_vol/missing_bkt"], + returncode=255, + stdout="", + stderr=( + "log4j:WARN No appenders could be found for logger (org.apache.hadoop.util.Shell).\n" + "BUCKET_NOT_FOUND Bucket not exists" + ), + ) + admin_hook.delete_bucket(volume_name="test_vol", bucket_name="missing_bkt") + mock_run_with_retry.assert_called_once() + + @patch(MOCK_RUN_RETRY_PATH) + def test_delete_volume_missing_is_idempotent( + self, mock_run_with_retry: MagicMock, admin_hook: OzoneAdminHook + ): + mock_run_with_retry.return_value = subprocess.CompletedProcess( + args=["ozone", "sh", "volume", "delete", "/missing_vol"], + returncode=255, + stdout="", + stderr=( + "log4j:WARN No appenders could be found for logger (org.apache.hadoop.util.Shell).\n" + "VOLUME_NOT_FOUND Volume missing_vol is not found" + ), + ) + admin_hook.delete_volume(volume_name="missing_vol") + mock_run_with_retry.assert_called_once() + + @patch(MOCK_CLI_PATH) + def test_get_container_report(self, mock_run_cli: MagicMock, admin_extra_hook: OzoneAdminExtraHook): + mock_run_cli.return_value = {"total": 1, "containers": [{"id": 1}]} + result = admin_extra_hook.get_container_report() + mock_run_cli.assert_called_once_with( + ["ozone", "admin", "container", "report", "--json"], + timeout=3600, + return_json_result=True, + ) + assert result["total"] == 1 + assert len(result["containers"]) == 1 + + +class TestOzoneFsHook: + @patch(MOCK_RUN_RETRY_PATH) + def test_exists_false(self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook): + mock_run_cli.return_value = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://path/does_not_exist"], + returncode=1, + stdout="", + stderr="not found", + ) + assert ozone_fs_hook.exists("ofs://path/does_not_exist") is False + mock_run_cli.assert_called_once() + command = mock_run_cli.call_args.args[0] + assert command[:4] == ["ozone", "fs", "-test", "-e"] + assert command[-1] == "ofs://path/does_not_exist" + + @patch(MOCK_RUN_RETRY_PATH) + def test_exists_false_when_returncode_one_has_only_noise_stderr( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook + ): + mock_run_cli.return_value = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://path/does_not_exist"], + returncode=1, + stdout="", + stderr=( + "log4j:WARN No appenders could be found for logger " + "(org.apache.hadoop.metrics2.lib.MutableMetricsFactory).\n" + "log4j:WARN Please initialize the log4j system properly.\n" + "log4j:WARN See http://logging.apache.org/log4j/1.2/faq.html#noconfig for more info." + ), + ) + assert ozone_fs_hook.exists("ofs://path/does_not_exist") is False + + @patch(MOCK_RUN_RETRY_PATH) + def test_exists_raises_on_meaningful_non_not_found_error( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook + ): + mock_run_cli.return_value = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://path/forbidden"], + returncode=1, + stdout="", + stderr="Permission denied", + ) + with pytest.raises(OzoneCliError, match="existence check failed"): + ozone_fs_hook.exists("ofs://path/forbidden") + + @patch(MOCK_RUN_RETRY_PATH) + def test_exists_raises_when_cli_missing(self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook): + mock_run_cli.side_effect = OzoneCliError("Ozone CLI not found.", retryable=False) + with pytest.raises(AirflowException): + ozone_fs_hook.exists("ofs://path/any") + + @patch(MOCK_CLI_PATH) + def test_upload_key_raises_when_file_missing( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook, tmp_path + ): + missing_path = tmp_path / "missing.txt" + with pytest.raises(AirflowException): + ozone_fs_hook.upload_key(str(missing_path), "ofs://vol1/bucket1/file.txt") + mock_run_cli.assert_not_called() + + @patch(MOCK_RUN_RETRY_PATH) + def test_upload_key_uses_plain_put_for_new_target( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook, tmp_path + ): + local_path = tmp_path / "payload.txt" + local_path.write_text("payload", encoding="utf-8") + mock_run_cli.side_effect = [ + subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1/file.txt"], + returncode=1, + stdout="", + stderr="not found", + ), + subprocess.CompletedProcess( + args=["ozone", "sh", "key", "put", "o3://vol1/bucket1/file.txt", str(local_path)], + returncode=0, + stdout="", + stderr="", + ), + ] + + ozone_fs_hook.upload_key(str(local_path), "ofs://vol1/bucket1/file.txt") + + assert mock_run_cli.call_count == 2 + assert mock_run_cli.call_args_list[1].args[0] == [ + "ozone", + "sh", + "key", + "put", + "o3://vol1/bucket1/file.txt", + str(local_path), + ] + + @patch(MOCK_RUN_RETRY_PATH) + def test_make_path_existing_target_policy_error( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook + ): + mock_run_cli.return_value = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1/dir"], + returncode=0, + stdout="", + stderr="", + ) + + with pytest.raises(AirflowException, match="Destination path already exists"): + ozone_fs_hook.make_path("ofs://vol1/bucket1/dir", if_exists=ExistingTargetPolicy.ERROR) + + mock_run_cli.assert_called_once() + + @patch(MOCK_RUN_RETRY_PATH) + def test_make_path_existing_target_policy_ignore( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook + ): + mock_run_cli.return_value = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1/dir"], + returncode=0, + stdout="", + stderr="", + ) + + ozone_fs_hook.make_path("ofs://vol1/bucket1/dir", if_exists=ExistingTargetPolicy.IGNORE) + + mock_run_cli.assert_called_once() + + @pytest.mark.parametrize( + ("fail_if_exists", "expected_if_exists"), + [ + (True, ExistingTargetPolicy.ERROR), + (False, ExistingTargetPolicy.IGNORE), + ], + ) + def test_create_path_deprecated_wrapper_warns( + self, + ozone_fs_hook: OzoneFsHook, + caplog, + fail_if_exists: bool, + expected_if_exists: ExistingTargetPolicy, + ): + with patch.object(ozone_fs_hook, "make_path") as mock_make_path: + with caplog.at_level("WARNING"): + ozone_fs_hook.create_path( + "ofs://vol1/bucket1/dir", + timeout=123, + recursive=False, + fail_if_exists=fail_if_exists, + ) + + assert "OzoneFsHook.create_path(..., fail_if_exists=...) is deprecated" in caplog.text + mock_make_path.assert_called_once_with( + "ofs://vol1/bucket1/dir", + timeout=123, + recursive=False, + if_exists=expected_if_exists, + ) + + @pytest.mark.parametrize( + "if_exists", + [ExistingTargetPolicy.ERROR, "ignore", ExistingTargetPolicy.OVERWRITE], + ) + @patch(MOCK_RUN_RETRY_PATH) + def test_upload_key_existing_target_policy( + self, + mock_run_cli: MagicMock, + ozone_fs_hook: OzoneFsHook, + tmp_path, + if_exists: ExistingTargetPolicy | str, + ): + local_path = tmp_path / "payload.txt" + local_path.write_text("payload", encoding="utf-8") + exists_result = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1/file.txt"], + returncode=0, + stdout="", + stderr="", + ) + put_result = subprocess.CompletedProcess( + args=["ozone", "sh", "key", "put", "o3://vol1/bucket1/file.txt", str(local_path)], + returncode=0, + stdout="", + stderr="", + ) + mock_run_cli.side_effect = [exists_result, put_result] + + if if_exists == "error": + with pytest.raises(AirflowException, match="Remote path already exists"): + ozone_fs_hook.upload_key( + str(local_path), + "ofs://vol1/bucket1/file.txt", + if_exists=if_exists, + ) + assert mock_run_cli.call_count == 1 + return + + with patch.object(ozone_fs_hook, "delete_key") as mock_delete_key: + ozone_fs_hook.upload_key( + str(local_path), + "ofs://vol1/bucket1/file.txt", + if_exists=if_exists, + ) + expected_calls = 1 if if_exists == "ignore" else 3 + expected_delete_calls = 0 if if_exists == "ignore" else 1 + assert mock_run_cli.call_count == expected_calls - expected_delete_calls + assert mock_delete_key.call_count == expected_delete_calls + if if_exists == "overwrite": + mock_delete_key.assert_called_once_with("ofs://vol1/bucket1/file.txt", timeout=3600) + assert mock_run_cli.call_args_list[1].args[0] == [ + "ozone", + "sh", + "key", + "put", + "o3://vol1/bucket1/file.txt", + str(local_path), + ] + + @patch(MOCK_RUN_RETRY_PATH) + def test_copy_path_existing_destination_fails_fast( + self, mock_run_cli: MagicMock, ozone_fs_hook: OzoneFsHook + ): + mock_run_cli.side_effect = [ + subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1"], + returncode=0, + stdout="", + stderr="", + ), + subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://vol1/bucket1/dst.txt"], + returncode=0, + stdout="", + stderr="", + ), + ] + + with pytest.raises(AirflowException, match="Destination path already exists"): + ozone_fs_hook.copy_path("ofs://vol1/bucket1/src.txt", "ofs://vol1/bucket1/dst.txt") + + assert mock_run_cli.call_count == 2 + + @patch(MOCK_RUN_RETRY_PATH) + def test_test_connection_failure(self, mock_run_cli_check: MagicMock, ozone_fs_hook: OzoneFsHook): + mock_run_cli_check.return_value = MagicMock(returncode=1, stdout="", stderr="auth failed") + ok, message = ozone_fs_hook.test_connection() + assert ok is False + assert "auth failed" in message + + @patch(MOCK_RUN_RETRY_PATH) + def test_test_connection_timeout(self, mock_run_cli_check: MagicMock, ozone_fs_hook: OzoneFsHook): + mock_run_cli_check.side_effect = OzoneCliError("Ozone command timed out", retryable=True) + ok, message = ozone_fs_hook.test_connection() + assert ok is False + assert "connection test failed" in message + + def test_get_key_info_converts_ofs_uri_for_key_cli(self, ozone_fs_hook: OzoneFsHook): + with patch.object(ozone_fs_hook, "run_cli", return_value={"name": "src.txt"}) as mock_run_cli: + assert ozone_fs_hook.get_key_info("ofs://vol1/b1/src.txt") == {"name": "src.txt"} + + mock_run_cli.assert_called_once_with( + ["ozone", "sh", "key", "info", "o3://vol1/b1/src.txt"], + timeout=300, + return_json_result=True, + ) + + def test_get_key_property_warns_when_replication_config_missing(self, ozone_fs_hook: OzoneFsHook, caplog): + with patch.object( + ozone_fs_hook, + "get_key_info", + return_value={"name": "file.txt", "replicationType": "RATIS", "replicationFactor": 3}, + ): + with caplog.at_level("WARNING"): + result = ozone_fs_hook.get_key_property("ofs://vol1/b1/file.txt") + assert result["replication_type"] == "RATIS" + assert result["replication"] == 3 + assert "does not contain 'replicationConfig'" in caplog.text diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/test_ozone.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/test_ozone.py new file mode 100644 index 0000000000000..5894076d02965 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/operators/test_ozone.py @@ -0,0 +1,320 @@ +# 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 unittest.mock import MagicMock, patch + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneFsHook +from airflow.providers.arenadata.ozone.operators.ozone import ( + OzoneCopyOperator, + OzoneCreateBucketOperator, + OzoneCreatePathOperator, + OzoneCreateVolumeOperator, + OzoneDeleteBucketOperator, + OzoneDeleteKeyOperator, + OzoneDeletePathOperator, + OzoneDeleteVolumeOperator, + OzoneDownloadFileOperator, + OzoneListOperator, + OzoneMoveOperator, + OzonePathExistsOperator, + OzoneSetQuotaOperator, + OzoneUploadContentOperator, + OzoneUploadFileOperator, +) + + +class TestOzoneAdminOperators: + def test_admin_operator_template_fields_cover_runtime_params(self): + assert OzoneCreateVolumeOperator.template_fields == ("volume_name", "quota", "ozone_conn_id") + assert OzoneCreateBucketOperator.template_fields == ( + "volume_name", + "bucket_name", + "quota", + "ozone_conn_id", + ) + assert OzoneSetQuotaOperator.template_fields == ("volume", "quota", "bucket", "ozone_conn_id") + assert OzoneDeleteVolumeOperator.template_fields == ("volume_name", "ozone_conn_id") + assert OzoneDeleteBucketOperator.template_fields == ("volume_name", "bucket_name", "ozone_conn_id") + + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneAdminHook") + def test_admin_operators_execute_core_flows(self, mock_admin_hook: MagicMock): + hook = mock_admin_hook.return_value + + OzoneCreateVolumeOperator(task_id="create_volume", volume_name="test_vol").execute(context={}) + OzoneDeleteVolumeOperator( + task_id="delete_volume", + volume_name="test_vol", + recursive=True, + force=True, + ).execute(context={}) + OzoneSetQuotaOperator(task_id="set_quota", volume="test_vol", quota="1TB").execute(context={}) + + assert mock_admin_hook.call_count == 3 + hook.create_volume.assert_called_once_with("test_vol", None, timeout=3600, if_exists="ignore") + hook.delete_volume.assert_called_once_with("test_vol", True, True, timeout=3600) + hook.set_quota.assert_called_once_with(volume="test_vol", quota="1TB", bucket=None, timeout=300) + + def test_admin_operators_validate_force_recursive_contract(self): + with pytest.raises(ValueError, match="force=True requires recursive=True"): + OzoneDeleteVolumeOperator( + task_id="bad_volume_delete", volume_name="v", force=True, recursive=False + ) + with pytest.raises(ValueError, match="force=True requires recursive=True"): + OzoneDeleteBucketOperator( + task_id="bad_bucket_delete", + volume_name="v", + bucket_name="b", + force=True, + recursive=False, + ) + + +@pytest.mark.parametrize( + ("operator", "hook_method", "expected_args", "expected_kwargs", "hook_return"), + [ + ( + OzoneListOperator(task_id="list_files", path="ofs://vol1/b1/"), + "list_keys", + ("ofs://vol1/b1/",), + {"timeout": 300}, + ["ofs://vol1/b1/f1"], + ), + ( + OzoneDeleteKeyOperator(task_id="delete_key", path="ofs://vol1/b1/f1"), + "delete_key", + ("ofs://vol1/b1/f1",), + {"timeout": 3600}, + None, + ), + ( + OzoneCreatePathOperator(task_id="create_path", path="ofs://vol1/b1/dir"), + "make_path", + ("ofs://vol1/b1/dir",), + {"timeout": 300, "if_exists": "ignore"}, + None, + ), + ( + OzoneDeletePathOperator(task_id="delete_path", path="ofs://vol1/b1/dir", recursive=True), + "delete_path", + ("ofs://vol1/b1/dir",), + {"recursive": True, "timeout": 3600}, + None, + ), + ( + OzonePathExistsOperator(task_id="path_exists", path="ofs://vol1/b1/dir"), + "path_exists", + ("ofs://vol1/b1/dir",), + {"timeout": 300}, + True, + ), + ( + OzoneCopyOperator( + task_id="copy_path", + source_path="ofs://vol1/b1/src.txt", + dest_path="ofs://vol1/b1/dst.txt", + ), + "copy_path", + ("ofs://vol1/b1/src.txt", "ofs://vol1/b1/dst.txt"), + {"timeout": 3600, "if_exists": "error"}, + None, + ), + ( + OzoneMoveOperator( + task_id="move_path", + source_path="ofs://vol1/b1/src.txt", + dest_path="ofs://vol1/b1/dst.txt", + ), + "move", + ("ofs://vol1/b1/src.txt", "ofs://vol1/b1/dst.txt"), + {"timeout": 3600, "if_exists": "error"}, + None, + ), + ], +) +@patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") +def test_fs_operators_delegate_to_hook( + mock_ozone_fs_hook: MagicMock, + operator, + hook_method: str, + expected_args: tuple[str, ...], + expected_kwargs: dict[str, object], + hook_return, +): + hook = mock_ozone_fs_hook.return_value + getattr(hook, hook_method).return_value = hook_return + + result = operator.execute(context={}) + + assert mock_ozone_fs_hook.call_args.kwargs["ozone_conn_id"] == OzoneFsHook.default_conn_name + getattr(hook, hook_method).assert_called_once_with(*expected_args, **expected_kwargs) + if hook_method in {"list_keys", "path_exists"}: + assert result == hook_return + else: + assert result is None + + +def test_fs_operator_template_fields_cover_runtime_params(): + assert OzoneCreatePathOperator.template_fields == ("path", "ozone_conn_id") + assert OzoneDeleteKeyOperator.template_fields == ("path", "ozone_conn_id") + assert OzoneDeletePathOperator.template_fields == ("path", "ozone_conn_id") + assert OzonePathExistsOperator.template_fields == ("path", "ozone_conn_id") + assert OzoneListOperator.template_fields == ("path", "ozone_conn_id") + assert OzoneUploadContentOperator.template_fields == ("content", "remote_path", "ozone_conn_id") + assert OzoneUploadFileOperator.template_fields == ("local_path", "remote_path", "ozone_conn_id") + assert OzoneMoveOperator.template_fields == ("source_path", "dest_path", "ozone_conn_id") + assert OzoneCopyOperator.template_fields == ("source_path", "dest_path", "ozone_conn_id") + assert OzoneDownloadFileOperator.template_fields == ("remote_path", "local_path", "ozone_conn_id") + + +class TestOzoneFileOperators: + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") + def test_upload_content_enforces_size_limit_and_uploads_when_valid(self, mock_ozone_fs_hook: MagicMock): + hook = mock_ozone_fs_hook.return_value + hook.connection_snapshot.max_content_size_bytes = 3 + too_large = OzoneUploadContentOperator( + task_id="upload_content_too_large", + content="1234", + remote_path="ofs://vol1/b1/payload.txt", + ) + with pytest.raises(AirflowException, match="exceeds configured limit"): + too_large.execute(context={}) + hook.upload_key.assert_not_called() + + hook.connection_snapshot.max_content_size_bytes = 10 + valid = OzoneUploadContentOperator( + task_id="upload_content_ok", + content="1234", + remote_path="ofs://vol1/b1/payload.txt", + ) + valid.execute(context={}) + hook.upload_key.assert_called_once() + assert hook.upload_key.call_args.kwargs["if_exists"] == "error" + + @pytest.mark.parametrize( + ("is_readable", "file_size", "expected_error"), + [ + (False, 1, "Local file not found or is not readable"), + (True, -1, "size is unavailable or non-positive"), + (True, 0, "size is unavailable or non-positive"), + (True, 1025, "exceeds configured limit"), + ], + ) + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") + @patch("airflow.providers.arenadata.ozone.operators.ozone.FileHelper.get_file_size_bytes") + @patch("airflow.providers.arenadata.ozone.operators.ozone.FileHelper.is_readable_file") + def test_upload_file_guard_paths( + self, + mock_is_readable_file: MagicMock, + mock_get_file_size_bytes: MagicMock, + mock_ozone_fs_hook: MagicMock, + is_readable: bool, + file_size: int, + expected_error: str, + ): + hook = mock_ozone_fs_hook.return_value + hook.connection_snapshot.max_content_size_bytes = 1024 + mock_is_readable_file.return_value = is_readable + mock_get_file_size_bytes.return_value = file_size + + operator = OzoneUploadFileOperator( + task_id="upload_file_guard", + local_path="/tmp/src.txt", + remote_path="ofs://vol1/b1/src.txt", + ) + with pytest.raises(AirflowException, match=expected_error): + operator.execute(context={}) + hook.upload_key.assert_not_called() + + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") + @patch("airflow.providers.arenadata.ozone.operators.ozone.FileHelper.get_file_size_bytes") + @patch("airflow.providers.arenadata.ozone.operators.ozone.FileHelper.is_readable_file") + def test_upload_file_success( + self, + mock_is_readable_file: MagicMock, + mock_get_file_size_bytes: MagicMock, + mock_ozone_fs_hook: MagicMock, + ): + hook = mock_ozone_fs_hook.return_value + hook.connection_snapshot.max_content_size_bytes = 2048 + mock_is_readable_file.return_value = True + mock_get_file_size_bytes.return_value = 1024 + + operator = OzoneUploadFileOperator( + task_id="upload_file_ok", + local_path="/tmp/src.txt", + remote_path="ofs://vol1/b1/src.txt", + ) + operator.execute(context={}) + + hook.upload_key.assert_called_once_with( + "/tmp/src.txt", + "ofs://vol1/b1/src.txt", + timeout=3600, + if_exists="error", + ) + + @pytest.mark.parametrize( + ("data_size", "expected_error"), + [ + (None, "Unable to determine remote key size"), + ("invalid", "Unable to determine remote key size"), + (101, "exceeds configured limit"), + ], + ) + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") + def test_download_file_guard_paths( + self, + mock_ozone_fs_hook: MagicMock, + data_size, + expected_error: str, + ): + hook = mock_ozone_fs_hook.return_value + hook.connection_snapshot.max_content_size_bytes = 100 + hook.get_key_property.return_value = {"data_size": data_size} + operator = OzoneDownloadFileOperator( + task_id="download_file_guard", + remote_path="ofs://vol1/b1/src.txt", + local_path="/tmp/dst.txt", + ) + with pytest.raises(AirflowException, match=expected_error): + operator.execute(context={}) + hook.download_key.assert_not_called() + + @patch("airflow.providers.arenadata.ozone.operators.ozone.OzoneFsHook") + def test_download_file_success(self, mock_ozone_fs_hook: MagicMock): + hook = mock_ozone_fs_hook.return_value + hook.connection_snapshot.max_content_size_bytes = 1024 * 1024 + hook.get_key_property.return_value = {"data_size": 100} + operator = OzoneDownloadFileOperator( + task_id="download_file_ok", + remote_path="ofs://vol1/b1/src.txt", + local_path="/tmp/dst.txt", + overwrite=True, + ) + operator.execute(context={}) + hook.get_key_property.assert_called_once_with("ofs://vol1/b1/src.txt", timeout=operator.timeout) + hook.download_key.assert_called_once_with( + "ofs://vol1/b1/src.txt", + "/tmp/dst.txt", + overwrite=True, + timeout=operator.timeout, + ) diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/test_ozone.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/test_ozone.py new file mode 100644 index 0000000000000..7ef4f0d4dd0be --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/sensors/test_ozone.py @@ -0,0 +1,51 @@ +# 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 unittest.mock import MagicMock, patch + +from airflow.providers.arenadata.ozone.hooks.ozone import OzoneFsHook +from airflow.providers.arenadata.ozone.sensors.ozone import OzoneKeySensor +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError + + +class TestOzoneKeySensor: + def test_template_fields_cover_runtime_params(self): + assert OzoneKeySensor.template_fields == ("path", "ozone_conn_id") + + @patch("airflow.providers.arenadata.ozone.sensors.ozone.OzoneFsHook") + def test_poke_key_exists(self, mock_ozone_hook: MagicMock): + mock_hook_instance = mock_ozone_hook.return_value + mock_hook_instance.key_exists = MagicMock(return_value=True) + sensor = OzoneKeySensor(task_id="test_sensor", path="ofs://vol1/bucket1/key1", timeout=42) + result = sensor.poke(context={}) + assert result is True + mock_ozone_hook.assert_called_once() + assert mock_ozone_hook.call_args.kwargs["ozone_conn_id"] == OzoneFsHook.default_conn_name + mock_hook_instance.key_exists.assert_called_once_with("ofs://vol1/bucket1/key1", timeout=42) + assert mock_hook_instance.key_exists.call_args.args[0] == "ofs://vol1/bucket1/key1" + assert sensor.cli_timeout == 42 + assert sensor.timeout != 42 + + @patch("airflow.providers.arenadata.ozone.sensors.ozone.OzoneFsHook") + def test_poke_retryable_error_returns_false(self, mock_ozone_hook: MagicMock): + mock_hook_instance = mock_ozone_hook.return_value + mock_hook_instance.key_exists.side_effect = OzoneCliError("Connection reset", retryable=True) + sensor = OzoneKeySensor(task_id="test_sensor", path="ofs://vol1/bucket1/key1") + result = sensor.poke(context={}) + assert result is False diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_hdfs_to_ozone.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_hdfs_to_ozone.py new file mode 100644 index 0000000000000..78c0c12a13cce --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_hdfs_to_ozone.py @@ -0,0 +1,275 @@ +# 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 unittest.mock import MagicMock, patch + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone import HdfsToOzoneOperator + + +class TestHdfsToOzoneOperator: + def test_template_fields_cover_runtime_params(self): + assert HdfsToOzoneOperator.template_fields == ("source_path", "dest_path", "hdfs_conn_id") + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute(self, mock_run_with_retry: MagicMock, _mock_which: MagicMock): + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_test", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + ) + operator.execute(context={}) + expected_cmd = [ + "hadoop", + "distcp", + "-update", + "-skipcrccheck", + "hdfs://nn:8020/user/data", + "ofs://om:9862/vol1/bucket1/data", + ] + mock_run_with_retry.assert_called_once() + call_args = mock_run_with_retry.call_args.args[0] + assert call_args == expected_cmd + + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", return_value=None) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_missing_hadoop_fails_fast(self, mock_run_process: MagicMock, _mock_which: MagicMock): + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_missing_hadoop", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + ) + with pytest.raises(AirflowException, match="executable 'hadoop' was not found in PATH"): + operator.execute(context={}) + mock_run_process.assert_not_called() + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_empty_source_fails_before_runner( + self, mock_run_process: MagicMock, _mock_which: MagicMock + ): + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_empty_source", + source_path=" ", + dest_path="ofs://om:9862/vol1/bucket1/data", + ) + with pytest.raises(AirflowException, match="requires non-empty source_path"): + operator.execute(context={}) + mock_run_process.assert_not_called() + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_empty_dest_fails_before_runner( + self, mock_run_process: MagicMock, _mock_which: MagicMock + ): + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_empty_dest", + source_path="hdfs://nn:8020/user/data", + dest_path="", + ) + with pytest.raises(AirflowException, match="requires non-empty dest_path"): + operator.execute(context={}) + mock_run_process.assert_not_called() + + def test_init_invalid_optional_conn_type(self): + with pytest.raises(ValueError, match="hdfs_conn_id must be a string or None"): + HdfsToOzoneOperator( + task_id="hdfs_to_ozone_test_invalid", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + hdfs_conn_id=123, # type: ignore[arg-type] + ) + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.BaseHook.get_connection") + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_wires_distcp_mapreduce_options( + self, + mock_run_process: MagicMock, + mock_get_connection: MagicMock, + _mock_which: MagicMock, + ): + conn = MagicMock() + conn.extra_dejson = { + "hdfs_distcp_mapreduce_local": "true", + "hdfs_distcp_renewer_principal": "slyubarsky@AD.RANGER-TEST", + } + mock_get_connection.return_value = conn + + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_distcp_options", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + hdfs_conn_id="hdfs_admin_default", + ) + operator.execute(context={}) + + assert mock_run_process.call_args.args[0] == [ + "hadoop", + "distcp", + "-Dmapreduce.framework.name=local", + "-Dyarn.resourcemanager.principal=slyubarsky@AD.RANGER-TEST", + "-Dmapreduce.jobtracker.kerberos.principal=slyubarsky@AD.RANGER-TEST", + "-update", + "-skipcrccheck", + "hdfs://nn:8020/user/data", + "ofs://om:9862/vol1/bucket1/data", + ] + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.KerberosConfig.kinit_hdfs_from_snapshot" + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.BaseHook.get_connection") + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_wires_hdfs_kerberos_env( + self, + mock_run_process: MagicMock, + mock_get_connection: MagicMock, + mock_kinit: MagicMock, + _mock_which: MagicMock, + ): + mock_kinit.return_value = True + conn = MagicMock() + conn.extra_dejson = { + "hdfs_ssl_enabled": "true", + "hdfs_ssl_keystore_location": "/etc/security/server.jks", + "hdfs_kerberos_enabled": "true", + "hdfs_kerberos_principal": "hdfs@EXAMPLE.COM", + "hdfs_kerberos_keytab": "/etc/security/keytabs/hdfs.keytab", + "krb5_conf": "/etc/krb5.conf", + } + mock_get_connection.return_value = conn + + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_kerberos", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + hdfs_conn_id="hdfs_admin_default", + ) + operator.execute(context={}) + + mock_kinit.assert_called_once_with( + snapshot=operator._hdfs_connection_snapshot, + conn_id="hdfs_admin_default", + ) + env_overrides = mock_run_process.call_args.kwargs["env_overrides"] + assert env_overrides["HADOOP_SECURITY_AUTHENTICATION"] == "kerberos" + assert env_overrides["HDFS_KERBEROS_PRINCIPAL"] == "hdfs@EXAMPLE.COM" + assert env_overrides["HDFS_KERBEROS_KEYTAB"] == "/etc/security/keytabs/hdfs.keytab" + assert env_overrides["KRB5_CONFIG"] == "/etc/krb5.conf" + assert env_overrides["HDFS_SSL_ENABLED"] == "true" + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.KerberosConfig.kinit_hdfs_from_snapshot" + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.BaseHook.get_connection") + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_wires_hdfs_kerberos_password_without_exporting_secret( + self, + mock_run_process: MagicMock, + mock_get_connection: MagicMock, + mock_kinit: MagicMock, + _mock_which: MagicMock, + ): + mock_kinit.return_value = True + conn = MagicMock() + conn.extra_dejson = { + "hdfs_kerberos_enabled": "true", + "hdfs_kerberos_principal": "hdfs@EXAMPLE.COM", + "hdfs_kerberos_password": "secret_password", + "krb5_conf": "/etc/krb5.conf", + } + mock_get_connection.return_value = conn + + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_kerberos_password", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + hdfs_conn_id="hdfs_admin_default", + ) + operator.execute(context={}) + + mock_kinit.assert_called_once_with( + snapshot=operator._hdfs_connection_snapshot, + conn_id="hdfs_admin_default", + ) + env_overrides = mock_run_process.call_args.kwargs["env_overrides"] + assert env_overrides["HADOOP_SECURITY_AUTHENTICATION"] == "kerberos" + assert env_overrides["HDFS_KERBEROS_PRINCIPAL"] == "hdfs@EXAMPLE.COM" + assert env_overrides["KRB5_CONFIG"] == "/etc/krb5.conf" + assert "HDFS_KERBEROS_PASSWORD" not in env_overrides + assert "HDFS_KERBEROS_KEYTAB" not in env_overrides + + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.shutil.which", + return_value="/usr/bin/hadoop", + ) + @patch( + "airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.KerberosConfig.kinit_hdfs_from_snapshot" + ) + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.BaseHook.get_connection") + @patch("airflow.providers.arenadata.ozone.transfers.hdfs_to_ozone.CliRunner.run_process") + def test_execute_fails_when_hdfs_kerberos_kinit_fails( + self, + mock_run_process: MagicMock, + mock_get_connection: MagicMock, + mock_kinit: MagicMock, + _mock_which: MagicMock, + ): + mock_kinit.return_value = False + conn = MagicMock() + conn.extra_dejson = { + "hdfs_kerberos_enabled": "true", + "hdfs_kerberos_principal": "hdfs@EXAMPLE.COM", + "hdfs_kerberos_keytab": "/etc/security/keytabs/hdfs.keytab", + } + mock_get_connection.return_value = conn + + operator = HdfsToOzoneOperator( + task_id="hdfs_to_ozone_kerberos_fail", + source_path="hdfs://nn:8020/user/data", + dest_path="ofs://om:9862/vol1/bucket1/data", + hdfs_conn_id="hdfs_admin_default", + ) + with pytest.raises(AirflowException, match="HDFS Kerberos authentication failed"): + operator.execute(context={}) + mock_run_process.assert_not_called() diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_ozone_backup.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_ozone_backup.py new file mode 100644 index 0000000000000..0612a1eac790f --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/transfers/test_ozone_backup.py @@ -0,0 +1,74 @@ +# 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 unittest.mock import MagicMock, patch + +from airflow.providers.arenadata.ozone.transfers.ozone_backup import OzoneBackupOperator +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError + + +class TestOzoneBackupOperator: + def test_template_fields_cover_runtime_params(self): + assert OzoneBackupOperator.template_fields == ("volume", "bucket", "snapshot_name", "ozone_conn_id") + + @patch("airflow.providers.arenadata.ozone.transfers.ozone_backup.OzoneAdminHook") + def test_execute_create_snapshot(self, mock_admin_hook: MagicMock): + mock_hook_instance = mock_admin_hook.return_value + operator = OzoneBackupOperator( + task_id="test_backup", volume="test_vol", bucket="test_bucket", snapshot_name="snapshot_20240101" + ) + operator.execute(context={}) + mock_admin_hook.assert_called_once() + assert mock_admin_hook.call_args.kwargs["ozone_conn_id"] == "ozone_admin_default" + mock_hook_instance.run_cli.assert_called_once() + snapshot_cmd = mock_hook_instance.run_cli.call_args.args[0] + assert snapshot_cmd[:4] == ["ozone", "sh", "snapshot", "create"] + assert "/test_vol/test_bucket" in snapshot_cmd + assert "snapshot_20240101" in snapshot_cmd + + @patch("airflow.providers.arenadata.ozone.transfers.ozone_backup.OzoneAdminHook") + def test_execute_idempotent_snapshot_exists(self, mock_admin_hook: MagicMock): + mock_hook_instance = mock_admin_hook.return_value + mock_hook_instance.run_cli.side_effect = OzoneCliError( + "Command failed: snapshot already exists", + stderr="FILE_ALREADY_EXISTS Snapshot already exists", + returncode=255, + ) + operator = OzoneBackupOperator( + task_id="test_backup", volume="test_vol", bucket="test_bucket", snapshot_name="snapshot_20240101" + ) + operator.execute(context={}) + mock_hook_instance.run_cli.assert_called_once() + + @patch("airflow.providers.arenadata.ozone.transfers.ozone_backup.OzoneAdminHook") + def test_execute_idempotent_snapshot_exists_human_message(self, mock_admin_hook: MagicMock): + mock_hook_instance = mock_admin_hook.return_value + mock_hook_instance.run_cli.side_effect = OzoneCliError( + "Command failed: snapshot already exists", + stderr="snapshot already exists", + returncode=255, + ) + operator = OzoneBackupOperator( + task_id="test_backup", + volume="test_vol", + bucket="test_bucket", + snapshot_name="snapshot_20240101", + ) + operator.execute(context={}) + mock_hook_instance.run_cli.assert_called_once() diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/__init__.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/__init__.py new file mode 100644 index 0000000000000..217e5db960782 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/__init__.py @@ -0,0 +1,17 @@ +# +# 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. diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_cli_runner.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_cli_runner.py new file mode 100644 index 0000000000000..099a6fbc1a46c --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_cli_runner.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import subprocess +import sys + +import pytest + +from airflow.providers.arenadata.ozone.utils.cli_runner import ( + CliRunner, + OzoneCliRunner, + ProcessOutputAnalysis, +) +from airflow.providers.arenadata.ozone.utils.errors import OzoneCliError + + +def test_process_output_analysis_filters_noise_and_prefers_meaningful_lines() -> None: + result = subprocess.CompletedProcess( + args=["ozone", "sh", "bucket", "delete", "/vol/missing"], + returncode=255, + stdout="", + stderr=( + "log4j:WARN No appenders could be found for logger (org.apache.hadoop.util.Shell).\n" + "BUCKET_NOT_FOUND Bucket not exists" + ), + ) + + output = ProcessOutputAnalysis.from_completed_process(result) + + assert output.raw_output == result.stderr + assert output.meaningful_lines == ("BUCKET_NOT_FOUND Bucket not exists",) + assert output.skipped_lines == ( + "log4j:WARN No appenders could be found for logger (org.apache.hadoop.util.Shell).", + ) + assert output.meaningful_output == "BUCKET_NOT_FOUND Bucket not exists" + assert output.preferred_output == "BUCKET_NOT_FOUND Bucket not exists" + assert output.normalized_preferred_output == "bucket_not_found bucket not exists" + + +def test_process_output_analysis_falls_back_to_raw_output_when_only_noise() -> None: + result = subprocess.CompletedProcess( + args=["ozone", "fs", "-test", "-e", "ofs://om/vol/bucket/missing.txt"], + returncode=1, + stdout="", + stderr="log4j:WARN Please initialize the log4j system properly.", + ) + + output = ProcessOutputAnalysis.from_completed_process(result) + + assert output.meaningful_lines == () + assert output.preferred_output == "log4j:WARN Please initialize the log4j system properly." + assert output.normalized_meaningful_output == "" + + +def test_clirunner_run_process_retries_and_returns_result(monkeypatch: pytest.MonkeyPatch) -> None: + attempts = {"count": 0} + + def fake_run( + _cls, + command: list[str], + *, + env_overrides=None, + timeout=300, + input_text=None, + cwd=None, + check=True, + log_output=False, + ) -> subprocess.CompletedProcess[str]: + attempts["count"] += 1 + if attempts["count"] == 1: + raise subprocess.CalledProcessError( + returncode=1, + cmd=command, + output="", + stderr="temporary failure", + ) + return subprocess.CompletedProcess(args=command, returncode=0, stdout="ok", stderr="") + + monkeypatch.setattr(CliRunner, "run", classmethod(fake_run)) + + result = CliRunner.run_process( + ["ozone", "sh", "volume", "list", "/"], + retry_attempts=2, + ) + assert result.returncode == 0 + assert result.stdout == "ok" + assert attempts["count"] == 2 + + +def test_clirunner_run_passes_input_text_to_stdin() -> None: + result = CliRunner.run( + [sys.executable, "-c", "import sys; sys.stdout.write(sys.stdin.read())"], + input_text="secret\n", + check=True, + ) + + assert result.stdout == "secret\n" + + +def test_ozoneclirunner_run_ozone_retries_retryable_and_not_non_retryable( + monkeypatch: pytest.MonkeyPatch, +) -> None: + retryable_attempts = {"count": 0} + + def fake_run_retryable( + _cls, + command: list[str], + *, + env_overrides=None, + timeout=300, + input_text=None, + cwd=None, + check=True, + log_output=False, + ) -> subprocess.CompletedProcess[str]: + retryable_attempts["count"] += 1 + if retryable_attempts["count"] == 1: + raise subprocess.CalledProcessError( + returncode=255, + cmd=command, + output="", + stderr="connection reset by peer", + ) + return subprocess.CompletedProcess(args=command, returncode=0, stdout="done", stderr="") + + monkeypatch.setattr(OzoneCliRunner, "run", classmethod(fake_run_retryable)) + result = OzoneCliRunner.run_ozone(["ozone", "sh", "bucket", "list", "/vol"], retry_attempts=2) + assert result.returncode == 0 + assert retryable_attempts["count"] == 2 + + non_retryable_attempts = {"count": 0} + + def fake_run_non_retryable( + _cls, + command: list[str], + *, + env_overrides=None, + timeout=300, + input_text=None, + cwd=None, + check=True, + log_output=False, + ) -> subprocess.CompletedProcess[str]: + non_retryable_attempts["count"] += 1 + raise subprocess.CalledProcessError( + returncode=1, + cmd=command, + output="", + stderr="ACCESS_DENIED", + ) + + monkeypatch.setattr(OzoneCliRunner, "run", classmethod(fake_run_non_retryable)) + with pytest.raises(OzoneCliError, match="non-retryable"): + OzoneCliRunner.run_ozone( + ["ozone", "sh", "bucket", "create", "/vol/bkt"], + retry_attempts=2, + ) + assert non_retryable_attempts["count"] == 1 diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_connection_schema.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_connection_schema.py new file mode 100644 index 0000000000000..f30481257e2c1 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_connection_schema.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.arenadata.ozone.utils.connection_schema import ( + CORE_SITE_XML, + KINIT_TIMEOUT_SECONDS, + MAX_CONTENT_SIZE_BYTES, + OZONE_SITE_XML, + OzoneConnSnapshot, +) + + +def _conn(host: str | None, port: int | None, extra: object) -> MagicMock: + conn = MagicMock() + conn.host = host + conn.port = port + conn.extra_dejson = extra + return conn + + +@pytest.mark.parametrize( + ("host", "port", "error_substring"), + [ + (None, 9862, "must define host"), + ("om", None, "must define port"), + ], +) +def test_from_connection_requires_host_and_port_in_strict_mode( + host: str | None, + port: int | None, + error_substring: str, +) -> None: + with pytest.raises(AirflowException, match=error_substring): + OzoneConnSnapshot.from_connection( + _conn(host, port, {}), + conn_id="ozone_default", + require_host_port=True, + ) + + +def test_from_connection_relaxed_mode_defaults_and_non_dict_extra() -> None: + snapshot = OzoneConnSnapshot.from_connection( + _conn(None, None, "not-a-dict"), + conn_id="hdfs_admin_default", + require_host_port=False, + ) + assert snapshot.host == "" + assert snapshot.port == 0 + assert snapshot.kinit_timeout_seconds == KINIT_TIMEOUT_SECONDS + assert snapshot.core_site_xml == CORE_SITE_XML + assert snapshot.ozone_site_xml == OZONE_SITE_XML + assert snapshot.max_content_size_bytes == MAX_CONTENT_SIZE_BYTES + + +def test_from_connection_parses_security_contract_and_runtime_overrides() -> None: + snapshot = OzoneConnSnapshot.from_connection( + _conn( + "om", + 9862, + { + "ozone_security_enabled": "true", + "hdfs_ssl_enabled": "1", + "hadoop_security_authentication": "kerberos", + "kerberos_principal": "user@REALM", + "kerberos_keytab": "/tmp/user.keytab", + "kerberos_password": "secret://vault/ozone/kerberos_password", + "krb5_conf": "/tmp/krb5.conf", + "ozone_conf_dir": "/opt/airflow/ozone-conf", + "hdfs_kerberos_principal": "hdfs@REALM", + "hdfs_kerberos_keytab": "/tmp/hdfs.keytab", + "hdfs_kerberos_password": "secret://vault/hdfs/kerberos_password", + "hdfs_distcp_mapreduce_local": "true", + "hdfs_distcp_renewer_principal": "yarn-rm/_HOST@REALM", + "kinit_timeout_seconds": "42", + "core_site_xml": "core-custom.xml", + "ozone_site_xml": "ozone-custom.xml", + "max_content_size_bytes": "2048", + }, + ), + conn_id="ozone_default", + require_host_port=True, + ) + assert snapshot.ozone_security_enabled is True + assert snapshot.hdfs_ssl_enabled is True + assert snapshot.hadoop_security_authentication == "kerberos" + assert snapshot.kerberos_principal == "user@REALM" + assert snapshot.kerberos_keytab == "/tmp/user.keytab" + assert snapshot.kerberos_password == "secret://vault/ozone/kerberos_password" + assert snapshot.krb5_conf == "/tmp/krb5.conf" + assert snapshot.ozone_conf_dir == "/opt/airflow/ozone-conf" + assert snapshot.hdfs_kerberos_principal == "hdfs@REALM" + assert snapshot.hdfs_kerberos_keytab == "/tmp/hdfs.keytab" + assert snapshot.hdfs_kerberos_password == "secret://vault/hdfs/kerberos_password" + assert snapshot.hdfs_distcp_mapreduce_local is True + assert snapshot.hdfs_distcp_renewer_principal == "yarn-rm/_HOST@REALM" + assert snapshot.kinit_timeout_seconds == 42 + assert snapshot.core_site_xml == "core-custom.xml" + assert snapshot.ozone_site_xml == "ozone-custom.xml" + assert snapshot.max_content_size_bytes == 2048 + + +@pytest.mark.parametrize( + ("extra", "expected_timeout", "expected_size_limit"), + [ + ({"kinit_timeout_seconds": "not-int", "max_content_size_bytes": "2048"}, KINIT_TIMEOUT_SECONDS, 2048), + ({"kinit_timeout_seconds": "15", "max_content_size_bytes": "-1"}, 15, MAX_CONTENT_SIZE_BYTES), + ], +) +def test_from_connection_invalid_numeric_overrides_fall_back_to_defaults( + extra: dict[str, str], + expected_timeout: int, + expected_size_limit: int, +) -> None: + snapshot = OzoneConnSnapshot.from_connection( + _conn("om", 9862, extra), + conn_id="ozone_default", + require_host_port=True, + ) + assert snapshot.kinit_timeout_seconds == expected_timeout + assert snapshot.max_content_size_bytes == expected_size_limit diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_helpers.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_helpers.py new file mode 100644 index 0000000000000..c094c711df5ad --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_helpers.py @@ -0,0 +1,70 @@ +# 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 pathlib import Path + +import pytest + +from airflow.providers.arenadata.ozone.utils.helpers import ( + FileHelper, + TypeNormalizationHelper, + URIHelper, +) + + +def test_normalize_optional_str_strips_and_handles_empty(): + assert TypeNormalizationHelper.normalize_optional_str(" value ") == "value" + assert TypeNormalizationHelper.normalize_optional_str(" ") is None + assert TypeNormalizationHelper.normalize_optional_str(None) is None + + +def test_require_optional_non_empty_preserves_type_contract(): + assert TypeNormalizationHelper.require_optional_non_empty(" abc ", "err") == "abc" + assert TypeNormalizationHelper.require_optional_non_empty(None, "err") is None + with pytest.raises(ValueError, match="err"): + TypeNormalizationHelper.require_optional_non_empty(123, "err") + with pytest.raises(ValueError, match="err"): + TypeNormalizationHelper.require_optional_non_empty(" ", "err") + + +def test_normalize_flag_bool(): + assert TypeNormalizationHelper.normalize_flag_bool("true", default=False) is True + assert TypeNormalizationHelper.normalize_flag_bool("No", default=True) is False + assert TypeNormalizationHelper.normalize_flag_bool(None, default=True) is True + + +def test_get_file_size_bytes(tmp_path: Path): + target = tmp_path / "payload.txt" + target.write_text("hello", encoding="utf-8") + assert FileHelper.get_file_size_bytes(target) == 5 + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("ofs://adho/vol1/bucket1/dir/file.txt", "o3://adho/vol1/bucket1/dir/file.txt"), + ("o3fs://adho/vol1/bucket1/dir/file.txt", "o3://adho/vol1/bucket1/dir/file.txt"), + ("o3://adho/vol1/bucket1/dir/file.txt", "o3://adho/vol1/bucket1/dir/file.txt"), + ("s3a://bucket/key", "s3a://bucket/key"), + ("vol1/bucket1/dir/file.txt", "vol1/bucket1/dir/file.txt"), + ("/vol1/bucket1/dir/file.txt", "vol1/bucket1/dir/file.txt"), + ], +) +def test_to_key_uri_returns_ozone_sh_key_target(value: str, expected: str): + assert URIHelper.to_key_uri(value) == expected diff --git a/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_security.py b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_security.py new file mode 100644 index 0000000000000..0565901942bc4 --- /dev/null +++ b/providers/arenadata/ozone/tests/unit/arenadata/ozone/utils/test_security.py @@ -0,0 +1,501 @@ +# +# 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 pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.arenadata.ozone.utils.connection_schema import OzoneConnSnapshot +from airflow.providers.arenadata.ozone.utils.security import ( + KerberosConfig, + SecretResolver, + SSLConfig, +) + + +class TestGetSecretValue: + """Tests for get_secret_value.""" + + @patch("airflow.providers.arenadata.ozone.utils.security.ensure_secrets_loaded") + def test_resolves_secret_uri_from_backend(self, mock_ensure_loaded): + mock_backend = MagicMock() + mock_backend.get_config.return_value = "resolved_secret" + mock_ensure_loaded.return_value = [mock_backend] + + result = SecretResolver.get_secret_value("secret://vault/ozone/password", conn_id="ozone_default") + + assert result == "resolved_secret" + mock_backend.get_config.assert_called_once_with("secret://vault/ozone/password") + + @patch("airflow.providers.arenadata.ozone.utils.security.ensure_secrets_loaded") + def test_secret_not_found_raises(self, mock_ensure_loaded): + mock_backend = MagicMock() + mock_backend.get_config.return_value = None + mock_ensure_loaded.return_value = [mock_backend] + + with pytest.raises(ValueError, match="Secret not found"): + SecretResolver.get_secret_value("secret://missing/uri") + + def test_get_secret_value_masks_plain_value_by_default(self, monkeypatch): + masked: dict[str, object] = {} + + def _capture_masked(value: object) -> None: + masked["value"] = value + + monkeypatch.setattr("airflow.providers.arenadata.ozone.utils.security.mask_secret", _capture_masked) + result = SecretResolver.get_secret_value("plain_secret") + assert result == "plain_secret" + assert masked["value"] == "plain_secret" + + +class TestGetKerberosEnvVars: + """Tests for KerberosConfig.get_env_vars.""" + + def test_empty_snapshot_returns_empty(self): + snapshot = OzoneConnSnapshot(host="om", port=9862) + assert KerberosConfig.get_env_vars(snapshot, scope="ozone") == {} + + def test_ozone_kerberos_from_snapshot(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + kerberos_keytab="/etc/keytab/user.keytab", + ) + result = KerberosConfig.get_env_vars(snapshot, scope="ozone") + assert result["HADOOP_SECURITY_AUTHENTICATION"] == "kerberos" + assert result["KERBEROS_PRINCIPAL"] == "user@REALM" + assert result["KERBEROS_KEYTAB"] == "/etc/keytab/user.keytab" + + def test_ozone_kerberos_password_is_not_exported_to_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + kerberos_password="secret_password", + ) + result = KerberosConfig.get_env_vars(snapshot, scope="ozone") + assert result["HADOOP_SECURITY_AUTHENTICATION"] == "kerberos" + assert result["KERBEROS_PRINCIPAL"] == "user@REALM" + assert "KERBEROS_PASSWORD" not in result + + def test_ozone_scope_does_not_include_hdfs_kerberos_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + kerberos_keytab="/etc/keytab/user.keytab", + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_keytab="/etc/keytab/hdfs.keytab", + ) + result = KerberosConfig.get_env_vars(snapshot, scope="ozone") + assert "HDFS_KERBEROS_PRINCIPAL" not in result + assert "HDFS_KERBEROS_KEYTAB" not in result + + def test_hdfs_scope_does_not_include_ozone_kerberos_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + kerberos_keytab="/etc/keytab/user.keytab", + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_keytab="/etc/keytab/hdfs.keytab", + ) + result = KerberosConfig.get_env_vars(snapshot, scope="hdfs") + assert "HADOOP_SECURITY_AUTHENTICATION" not in result + assert "KERBEROS_PRINCIPAL" not in result + assert result["HDFS_KERBEROS_PRINCIPAL"] == "hdfs@REALM" + assert result["HDFS_KERBEROS_KEYTAB"] == "/etc/keytab/hdfs.keytab" + + def test_hdfs_kerberos_password_is_not_exported_to_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_password="secret_password", + ) + result = KerberosConfig.get_env_vars(snapshot, scope="hdfs") + assert result["HDFS_KERBEROS_PRINCIPAL"] == "hdfs@REALM" + assert "HDFS_KERBEROS_PASSWORD" not in result + + +class TestApplySslEnvVars: + """Tests for SSLConfig.apply_ssl_env_vars.""" + + def test_merges_into_existing(self): + overrides = {"A": "1"} + existing = {"B": "2"} + result = SSLConfig.apply_ssl_env_vars(overrides, existing) + assert result == {"A": "1", "B": "2"} + + def test_scope_ozone_does_not_include_hdfs_ssl_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + ozone_security_enabled=True, + ozone_om_https_port="9879", + hdfs_ssl_enabled=True, + hdfs_ssl_keystore_location="/etc/hdfs.ks", + ) + config = SSLConfig.from_snapshot(snapshot, scope="ozone") + env = config.as_env() + assert env["OZONE_SECURITY_ENABLED"] == "true" + assert "HDFS_SSL_ENABLED" not in env + + def test_scope_hdfs_does_not_include_ozone_ssl_env(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + ozone_security_enabled=True, + ozone_om_https_port="9879", + hdfs_ssl_enabled=True, + hdfs_ssl_keystore_location="/etc/hdfs.ks", + ) + config = SSLConfig.from_snapshot(snapshot, scope="hdfs") + env = config.as_env() + assert env["HDFS_SSL_ENABLED"] == "true" + assert "OZONE_SECURITY_ENABLED" not in env + + +class TestApplyKerberosEnvVars: + """Tests for KerberosConfig.apply_env_vars.""" + + def test_sets_hadoop_opts_when_kerberos_enabled(self): + env_vars = {"HADOOP_SECURITY_AUTHENTICATION": "kerberos"} + result = KerberosConfig.apply_env_vars(env_vars) + assert "-Dhadoop.security.authentication=kerberos" in result.get("HADOOP_OPTS", "") + assert "-Dhadoop.security.authentication=kerberos" in result.get("OZONE_OPTS", "") + + def test_reuses_existing_jvm_opts_when_kerberos_enabled(self): + env_vars = {"HADOOP_SECURITY_AUTHENTICATION": "kerberos"} + existing = { + "HADOOP_OPTS": "-Dexisting.hadoop=true", + "OZONE_OPTS": "-Dexisting.ozone=true", + } + result = KerberosConfig.apply_env_vars(env_vars, existing) + assert "-Dexisting.hadoop=true" in result["HADOOP_OPTS"] + assert "-Dhadoop.security.authentication=kerberos" in result["HADOOP_OPTS"] + assert "-Dexisting.ozone=true" in result["OZONE_OPTS"] + assert "-Dhadoop.security.authentication=kerberos" in result["OZONE_OPTS"] + assert "-Dozone.security.enabled=true" in result["OZONE_OPTS"] + + def test_reuses_explicit_existing_ozone_conf_dir_when_kerberos_enabled(self): + env_vars = {"HADOOP_SECURITY_AUTHENTICATION": "kerberos"} + result = KerberosConfig.apply_env_vars( + env_vars, + existing_env={"OZONE_CONF_DIR": "/opt/airflow/ozone-conf"}, + ) + assert result["OZONE_CONF_DIR"] == "/opt/airflow/ozone-conf" + assert result["HADOOP_CONF_DIR"] == "/opt/airflow/ozone-conf" + + def test_reuses_process_env_when_existing_env_is_not_provided(self, monkeypatch): + env_vars = {"HADOOP_SECURITY_AUTHENTICATION": "kerberos"} + monkeypatch.setenv("OZONE_CONF_DIR", "/env/ozone-conf") + monkeypatch.setenv("HADOOP_OPTS", "-Dexisting.hadoop=true") + monkeypatch.setenv("OZONE_OPTS", "-Dexisting.ozone=true") + + result = KerberosConfig.apply_env_vars(env_vars) + + assert result["OZONE_CONF_DIR"] == "/env/ozone-conf" + assert result["HADOOP_CONF_DIR"] == "/env/ozone-conf" + assert "-Dexisting.hadoop=true" in result["HADOOP_OPTS"] + assert "-Dhadoop.security.authentication=kerberos" in result["HADOOP_OPTS"] + assert "-Dexisting.ozone=true" in result["OZONE_OPTS"] + assert "-Dhadoop.security.authentication=kerberos" in result["OZONE_OPTS"] + + def test_kerberos_helpers_detect_enabled_state_and_config_dir(self): + kerberos_env = { + "HADOOP_SECURITY_AUTHENTICATION": "kerberos", + "OZONE_CONF_DIR": "/opt/airflow/ozone-conf", + } + + assert KerberosConfig.is_enabled(kerberos_env) + assert KerberosConfig.resolve_config_dir(kerberos_env) == "/opt/airflow/ozone-conf" + assert not KerberosConfig.is_enabled(None) + assert KerberosConfig.resolve_config_dir(None) is None + + +class TestSnapshotDrivenKerberosRuntime: + @patch("airflow.providers.arenadata.ozone.utils.security.FileHelper.is_readable_file", return_value=True) + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosCliRunner.run_kerberos") + def test_kinit_timeout_from_snapshot(self, mock_run_kerberos, _mock_readable): + mock_run_kerberos.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kinit_timeout_seconds=42, + ) + assert KerberosConfig.kinit_with_keytab( + "user@REALM", + "/tmp/user.keytab", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + assert mock_run_kerberos.call_args.kwargs["timeout"] == 42 + + @patch("airflow.providers.arenadata.ozone.utils.security.FileHelper.is_readable_file", return_value=True) + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosCliRunner.run_kerberos") + def test_kinit_with_password_uses_stdin(self, mock_run_kerberos, _mock_readable): + mock_run_kerberos.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kinit_timeout_seconds=42, + ) + assert KerberosConfig.kinit_with_password( + "user@REALM", + "secret_password", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + mock_run_kerberos.assert_called_once_with( + ["kinit", "user@REALM"], + env_overrides={"KRB5_CONFIG": "/tmp/krb5.conf"}, + timeout=42, + input_text="secret_password\n", + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_password") + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_keytab") + def test_kinit_from_snapshot_prefers_keytab_when_both_are_configured( + self, + mock_kinit_with_keytab: MagicMock, + mock_kinit_with_password: MagicMock, + ): + mock_kinit_with_keytab.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kerberos_principal="user@REALM", + kerberos_keytab="/tmp/user.keytab", + kerberos_password="secret_password", + ) + assert KerberosConfig.kinit_from_snapshot( + snapshot=snapshot, + ) + mock_kinit_with_keytab.assert_called_once_with( + "user@REALM", + "/tmp/user.keytab", + None, + snapshot=snapshot, + ) + mock_kinit_with_password.assert_not_called() + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_password") + def test_kinit_from_snapshot_uses_password_without_keytab( + self, + mock_kinit_with_password: MagicMock, + ): + mock_kinit_with_password.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kerberos_principal="user@REALM", + kerberos_password="secret_password", + ) + assert KerberosConfig.kinit_from_snapshot( + snapshot=snapshot, + ) + mock_kinit_with_password.assert_called_once_with( + "user@REALM", + "secret_password", + None, + snapshot=snapshot, + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_keytab") + def test_kinit_from_snapshot_uses_snapshot_keytab_and_krb5_conf( + self, + mock_kinit_with_keytab: MagicMock, + ): + mock_kinit_with_keytab.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kerberos_principal="user@REALM", + kerberos_keytab="/tmp/user.keytab", + krb5_conf="/tmp/krb5.conf", + ) + assert KerberosConfig.kinit_from_snapshot(snapshot=snapshot) + mock_kinit_with_keytab.assert_called_once_with( + "user@REALM", + "/tmp/user.keytab", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_keytab") + def test_kinit_hdfs_from_snapshot_uses_hdfs_snapshot_credentials( + self, + mock_kinit_with_keytab: MagicMock, + ): + mock_kinit_with_keytab.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_keytab="/tmp/hdfs.keytab", + krb5_conf="/tmp/krb5.conf", + ) + assert KerberosConfig.kinit_hdfs_from_snapshot(snapshot=snapshot, conn_id="hdfs_admin_default") + mock_kinit_with_keytab.assert_called_once_with( + "hdfs@REALM", + "/tmp/hdfs.keytab", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_password") + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_keytab") + def test_kinit_hdfs_from_snapshot_prefers_keytab_when_both_are_configured( + self, + mock_kinit_with_keytab: MagicMock, + mock_kinit_with_password: MagicMock, + ): + mock_kinit_with_keytab.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_keytab="/tmp/hdfs.keytab", + hdfs_kerberos_password="secret_password", + ) + assert KerberosConfig.kinit_hdfs_from_snapshot(snapshot=snapshot, conn_id="hdfs_admin_default") + mock_kinit_with_keytab.assert_called_once_with( + "hdfs@REALM", + "/tmp/hdfs.keytab", + None, + snapshot=snapshot, + ) + mock_kinit_with_password.assert_not_called() + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_password") + def test_kinit_hdfs_from_snapshot_uses_password_without_keytab( + self, + mock_kinit_with_password: MagicMock, + ): + mock_kinit_with_password.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hdfs_kerberos_enabled=True, + hdfs_kerberos_principal="hdfs@REALM", + hdfs_kerberos_password="secret_password", + krb5_conf="/tmp/krb5.conf", + ) + assert KerberosConfig.kinit_hdfs_from_snapshot(snapshot=snapshot, conn_id="hdfs_admin_default") + mock_kinit_with_password.assert_called_once_with( + "hdfs@REALM", + "secret_password", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_with_password") + def test_kinit_from_snapshot_uses_snapshot_password_and_krb5_conf( + self, + mock_kinit_with_password: MagicMock, + ): + mock_kinit_with_password.return_value = True + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + kerberos_principal="user@REALM", + kerberos_password="secret_password", + krb5_conf="/tmp/krb5.conf", + ) + assert KerberosConfig.kinit_from_snapshot(snapshot=snapshot) + mock_kinit_with_password.assert_called_once_with( + "user@REALM", + "secret_password", + "/tmp/krb5.conf", + snapshot=snapshot, + ) + + def test_check_config_files_exist_uses_snapshot_names(self, tmp_path: Path): + (tmp_path / "custom-core.xml").write_text("", encoding="utf-8") + (tmp_path / "custom-ozone.xml").write_text("", encoding="utf-8") + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + core_site_xml="custom-core.xml", + ozone_site_xml="custom-ozone.xml", + ) + assert KerberosConfig.check_config_files_exist(str(tmp_path), snapshot=snapshot) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosCliRunner.run_kerberos") + def test_has_valid_ticket_uses_klist(self, mock_run_kerberos: MagicMock): + mock_run_kerberos.return_value = True + snapshot = OzoneConnSnapshot(host="om", port=9862, kinit_timeout_seconds=44) + assert KerberosConfig.has_valid_ticket(snapshot=snapshot) + mock_run_kerberos.assert_called_once_with( + ["klist", "-s"], + timeout=44, + log_output=False, + ) + + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.has_valid_ticket") + @patch("airflow.providers.arenadata.ozone.utils.security.KerberosConfig.kinit_from_snapshot") + def test_ensure_ticket_rechecks_lifetime_when_cached( + self, + mock_kinit_from_snapshot: MagicMock, + mock_has_valid_ticket: MagicMock, + ): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + kerberos_keytab="/tmp/user.keytab", + ) + mock_has_valid_ticket.return_value = False + mock_kinit_from_snapshot.return_value = True + + assert KerberosConfig.ensure_ticket( + snapshot=snapshot, + conn_id="ozone_default", + kerberos_ticket_ready=True, + ) + mock_has_valid_ticket.assert_called_once_with(snapshot=snapshot) + mock_kinit_from_snapshot.assert_called_once_with(snapshot=snapshot, conn_id="ozone_default") + + def test_ensure_ticket_fails_fast_when_kerberos_credentials_missing(self): + snapshot = OzoneConnSnapshot( + host="om", + port=9862, + hadoop_security_authentication="kerberos", + kerberos_principal="user@REALM", + ) + with pytest.raises(AirflowException, match="neither 'kerberos_keytab' nor 'kerberos_password'"): + KerberosConfig.ensure_ticket( + snapshot=snapshot, + conn_id="ozone_default", + kerberos_ticket_ready=False, + ) diff --git a/pyproject.toml b/pyproject.toml index 34bbeab57b220..99a769e0eb4eb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -107,6 +107,9 @@ apache-airflow = "airflow.__main__:main" "amazon" = [ "apache-airflow-providers-amazon>=9.0.0" ] +"arenadata.ozone" = [ + "apache-airflow-providers-arenadata-ozone>=1.1.0" +] "apache.cassandra" = [ "apache-airflow-providers-apache-cassandra>=3.7.0; python_version !=\"3.14\"" ] @@ -398,6 +401,7 @@ apache-airflow = "airflow.__main__:main" "apache-airflow-providers-airbyte>=5.0.0", "apache-airflow-providers-alibaba>=3.0.0", "apache-airflow-providers-amazon>=9.0.0", + "apache-airflow-providers-arenadata-ozone>=1.1.0", "apache-airflow-providers-apache-cassandra>=3.7.0; python_version !=\"3.14\"", "apache-airflow-providers-apache-drill>=2.8.1", "apache-airflow-providers-apache-druid>=3.12.0", @@ -1056,6 +1060,8 @@ mypy_path = [ "$MYPY_CONFIG_FILE_DIR/providers/alibaba/tests", "$MYPY_CONFIG_FILE_DIR/providers/amazon/src", "$MYPY_CONFIG_FILE_DIR/providers/amazon/tests", + "$MYPY_CONFIG_FILE_DIR/providers/arenadata/ozone/src", + "$MYPY_CONFIG_FILE_DIR/providers/arenadata/ozone/tests", "$MYPY_CONFIG_FILE_DIR/providers/apache/cassandra/src", "$MYPY_CONFIG_FILE_DIR/providers/apache/cassandra/tests", "$MYPY_CONFIG_FILE_DIR/providers/apache/drill/src", @@ -1365,6 +1371,7 @@ apache-airflow-providers = false apache-airflow-providers-airbyte = false apache-airflow-providers-alibaba = false apache-airflow-providers-amazon = false +apache-airflow-providers-arenadata-ozone = false apache-airflow-providers-apache-cassandra = false apache-airflow-providers-apache-drill = false apache-airflow-providers-apache-druid = false @@ -1500,6 +1507,7 @@ apache-airflow-providers = false apache-airflow-providers-airbyte = false apache-airflow-providers-alibaba = false apache-airflow-providers-amazon = false +apache-airflow-providers-arenadata-ozone = false apache-airflow-providers-apache-cassandra = false apache-airflow-providers-apache-drill = false apache-airflow-providers-apache-druid = false @@ -1651,6 +1659,7 @@ apache-airflow-shared-timezones = { workspace = true } apache-airflow-providers-airbyte = { workspace = true } apache-airflow-providers-alibaba = { workspace = true } apache-airflow-providers-amazon = { workspace = true } +apache-airflow-providers-arenadata-ozone = { workspace = true } apache-airflow-providers-apache-cassandra = { workspace = true } apache-airflow-providers-apache-drill = { workspace = true } apache-airflow-providers-apache-druid = { workspace = true } @@ -1784,6 +1793,7 @@ members = [ "providers/airbyte", "providers/alibaba", "providers/amazon", + "providers/arenadata/ozone", "providers/apache/cassandra", "providers/apache/drill", "providers/apache/druid", diff --git a/scripts/ci/docker-compose/integration-ozone.yml b/scripts/ci/docker-compose/integration-ozone.yml new file mode 100644 index 0000000000000..7731d47d07542 --- /dev/null +++ b/scripts/ci/docker-compose/integration-ozone.yml @@ -0,0 +1,104 @@ +# 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. +--- +x-ozone-common: + &ozone-common + image: apache/ozone:2.0.0 + env_file: + - ./ozone/docker-config + restart: "on-failure" + +services: + ozone-scm: + <<: *ozone-common + container_name: ozone-scm + hostname: scm + labels: + breeze.description: "Ozone Storage Container Manager." + ports: + - "${OZONE_SCM_HOST_PORT}:9860" + environment: + ENSURE_SCM_INITIALIZED: /data/metadata/scm/current/VERSION + volumes: + - ozone-bin:/opt/hadoop + - ozone-java:/usr/local/jdk-21.0.2 + command: ["ozone", "scm"] + healthcheck: + test: ["CMD-SHELL", "/opt/hadoop/bin/ozone admin scm roles 2>/dev/null | grep -q LEADER"] + interval: 10s + timeout: 30s + retries: 30 + start_period: 20s + + ozone-om: + <<: *ozone-common + container_name: ozone-om + hostname: om + labels: + breeze.description: "Ozone Manager." + ports: + - "${OZONE_OM_HOST_PORT}:9862" + environment: + ENSURE_OM_INITIALIZED: /data/metadata/om/current/VERSION + command: ["ozone", "om"] + depends_on: + ozone-scm: + condition: service_healthy + healthcheck: + test: ["CMD-SHELL", "/opt/hadoop/bin/ozone sh volume list / 2>/dev/null | grep -q '\\['"] + interval: 10s + timeout: 30s + retries: 30 + start_period: 20s + + ozone-datanode: + <<: *ozone-common + container_name: ozone-datanode + hostname: datanode + labels: + breeze.description: "Ozone Datanode." + command: ["ozone", "datanode"] + depends_on: + ozone-scm: + condition: service_healthy + healthcheck: + test: ["CMD-SHELL", "bash -c '/dev/null"] + interval: 10s + timeout: 10s + retries: 40 + start_period: 30s + + airflow: + environment: + - INTEGRATION_OZONE=true + - PATH=/opt/ozone/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + - OZONE_HOME=/opt/ozone + - OZONE_CONF_DIR=/opt/ozone-conf + - JAVA_HOME=/opt/ozone-java + volumes: + - ozone-bin:/opt/ozone:ro + - ozone-java:/opt/ozone-java:ro + - ./ozone/ozone-site.xml:/opt/ozone-conf/ozone-site.xml:ro + depends_on: + ozone-om: + condition: service_healthy + ozone-datanode: + condition: service_healthy + +volumes: + ozone-bin: + ozone-java: diff --git a/scripts/ci/docker-compose/remove-sources.yml b/scripts/ci/docker-compose/remove-sources.yml index a4287f26e46d3..14c3252bb14fe 100644 --- a/scripts/ci/docker-compose/remove-sources.yml +++ b/scripts/ci/docker-compose/remove-sources.yml @@ -47,6 +47,7 @@ services: - ../../../empty:/opt/airflow/providers/apache/tinkerpop/src - ../../../empty:/opt/airflow/providers/apprise/src - ../../../empty:/opt/airflow/providers/arangodb/src + - ../../../empty:/opt/airflow/providers/arenadata/ozone/src - ../../../empty:/opt/airflow/providers/asana/src - ../../../empty:/opt/airflow/providers/atlassian/jira/src - ../../../empty:/opt/airflow/providers/celery/src diff --git a/scripts/ci/docker-compose/tests-sources.yml b/scripts/ci/docker-compose/tests-sources.yml index eb9da3cacebf5..119e4932ec2c1 100644 --- a/scripts/ci/docker-compose/tests-sources.yml +++ b/scripts/ci/docker-compose/tests-sources.yml @@ -60,6 +60,7 @@ services: - ../../../providers/apache/tinkerpop/tests:/opt/airflow/providers/apache/tinkerpop/tests - ../../../providers/apprise/tests:/opt/airflow/providers/apprise/tests - ../../../providers/arangodb/tests:/opt/airflow/providers/arangodb/tests + - ../../../providers/arenadata/ozone/tests:/opt/airflow/providers/arenadata/ozone/tests - ../../../providers/asana/tests:/opt/airflow/providers/asana/tests - ../../../providers/atlassian/jira/tests:/opt/airflow/providers/atlassian/jira/tests - ../../../providers/celery/tests:/opt/airflow/providers/celery/tests