apache-airflow-providers-standard 1.0.0.dev1__py3-none-any.whl → 1.1.0rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of apache-airflow-providers-standard might be problematic. Click here for more details.

Files changed (50) hide show
  1. airflow/providers/standard/LICENSE +0 -52
  2. airflow/providers/standard/__init__.py +23 -1
  3. airflow/providers/standard/decorators/__init__.py +16 -0
  4. airflow/providers/standard/decorators/bash.py +121 -0
  5. airflow/providers/standard/decorators/branch_external_python.py +63 -0
  6. airflow/providers/standard/decorators/branch_python.py +62 -0
  7. airflow/providers/standard/decorators/branch_virtualenv.py +62 -0
  8. airflow/providers/standard/decorators/external_python.py +70 -0
  9. airflow/providers/standard/decorators/python.py +86 -0
  10. airflow/providers/standard/decorators/python_virtualenv.py +67 -0
  11. airflow/providers/standard/decorators/sensor.py +83 -0
  12. airflow/providers/standard/decorators/short_circuit.py +65 -0
  13. airflow/providers/standard/get_provider_info.py +89 -7
  14. airflow/providers/standard/hooks/__init__.py +16 -0
  15. airflow/providers/standard/hooks/filesystem.py +89 -0
  16. airflow/providers/standard/hooks/package_index.py +95 -0
  17. airflow/providers/standard/hooks/subprocess.py +119 -0
  18. airflow/providers/standard/operators/bash.py +73 -56
  19. airflow/providers/standard/operators/branch.py +105 -0
  20. airflow/providers/standard/operators/datetime.py +15 -5
  21. airflow/providers/standard/operators/empty.py +39 -0
  22. airflow/providers/standard/operators/latest_only.py +127 -0
  23. airflow/providers/standard/operators/python.py +1143 -0
  24. airflow/providers/standard/operators/smooth.py +38 -0
  25. airflow/providers/standard/operators/trigger_dagrun.py +391 -0
  26. airflow/providers/standard/operators/weekday.py +19 -9
  27. airflow/providers/standard/sensors/bash.py +15 -11
  28. airflow/providers/standard/sensors/date_time.py +32 -8
  29. airflow/providers/standard/sensors/external_task.py +593 -0
  30. airflow/providers/standard/sensors/filesystem.py +158 -0
  31. airflow/providers/standard/sensors/python.py +84 -0
  32. airflow/providers/standard/sensors/time.py +28 -5
  33. airflow/providers/standard/sensors/time_delta.py +68 -15
  34. airflow/providers/standard/sensors/weekday.py +25 -7
  35. airflow/providers/standard/triggers/__init__.py +16 -0
  36. airflow/providers/standard/triggers/external_task.py +288 -0
  37. airflow/providers/standard/triggers/file.py +131 -0
  38. airflow/providers/standard/triggers/temporal.py +113 -0
  39. airflow/providers/standard/utils/__init__.py +16 -0
  40. airflow/providers/standard/utils/python_virtualenv.py +209 -0
  41. airflow/providers/standard/utils/python_virtualenv_script.jinja2 +82 -0
  42. airflow/providers/standard/utils/sensor_helper.py +137 -0
  43. airflow/providers/standard/utils/skipmixin.py +192 -0
  44. airflow/providers/standard/utils/weekday.py +77 -0
  45. airflow/providers/standard/version_compat.py +36 -0
  46. {apache_airflow_providers_standard-1.0.0.dev1.dist-info → apache_airflow_providers_standard-1.1.0rc1.dist-info}/METADATA +16 -35
  47. apache_airflow_providers_standard-1.1.0rc1.dist-info/RECORD +51 -0
  48. {apache_airflow_providers_standard-1.0.0.dev1.dist-info → apache_airflow_providers_standard-1.1.0rc1.dist-info}/WHEEL +1 -1
  49. apache_airflow_providers_standard-1.0.0.dev1.dist-info/RECORD +0 -17
  50. {apache_airflow_providers_standard-1.0.0.dev1.dist-info → apache_airflow_providers_standard-1.1.0rc1.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,83 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
17
+
18
+ from __future__ import annotations
19
+
20
+ from collections.abc import Sequence
21
+ from typing import TYPE_CHECKING, Callable, ClassVar
22
+
23
+ from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
24
+
25
+ if AIRFLOW_V_3_0_PLUS:
26
+ from airflow.sdk.bases.decorator import get_unique_task_id, task_decorator_factory
27
+ else:
28
+ from airflow.decorators.base import get_unique_task_id, task_decorator_factory # type: ignore[no-redef]
29
+
30
+
31
+ from airflow.providers.standard.sensors.python import PythonSensor
32
+
33
+ if TYPE_CHECKING:
34
+ from airflow.sdk.bases.decorator import TaskDecorator
35
+
36
+
37
+ class DecoratedSensorOperator(PythonSensor):
38
+ """
39
+ Wraps a Python callable and captures args/kwargs when called for execution.
40
+
41
+ :param python_callable: A reference to an object that is callable
42
+ :param task_id: task Id
43
+ :param op_args: a list of positional arguments that will get unpacked when
44
+ calling your callable (templated)
45
+ :param op_kwargs: a dictionary of keyword arguments that will get unpacked
46
+ in your function (templated)
47
+ :param kwargs_to_upstream: For certain operators, we might need to upstream certain arguments
48
+ that would otherwise be absorbed by the DecoratedOperator (for example python_callable for the
49
+ PythonOperator). This gives a user the option to upstream kwargs as needed.
50
+ """
51
+
52
+ template_fields: Sequence[str] = ("op_args", "op_kwargs")
53
+ template_fields_renderers: ClassVar[dict[str, str]] = {"op_args": "py", "op_kwargs": "py"}
54
+
55
+ custom_operator_name = "@task.sensor"
56
+
57
+ # since we won't mutate the arguments, we should just do the shallow copy
58
+ # there are some cases we can't deepcopy the objects (e.g protobuf).
59
+ shallow_copy_attrs: Sequence[str] = ("python_callable",)
60
+
61
+ def __init__(
62
+ self,
63
+ *,
64
+ task_id: str,
65
+ **kwargs,
66
+ ) -> None:
67
+ kwargs["task_id"] = get_unique_task_id(task_id, kwargs.get("dag"), kwargs.get("task_group"))
68
+ super().__init__(**kwargs)
69
+
70
+
71
+ def sensor_task(python_callable: Callable | None = None, **kwargs) -> TaskDecorator:
72
+ """
73
+ Wrap a function into an Airflow operator.
74
+
75
+ Accepts kwargs for operator kwarg. Can be reused in a single DAG.
76
+ :param python_callable: Function to decorate
77
+ """
78
+ return task_decorator_factory(
79
+ python_callable=python_callable,
80
+ multiple_outputs=False,
81
+ decorated_operator_class=DecoratedSensorOperator,
82
+ **kwargs,
83
+ )
@@ -0,0 +1,65 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
17
+ from __future__ import annotations
18
+
19
+ from typing import TYPE_CHECKING, Callable
20
+
21
+ from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
22
+
23
+ if AIRFLOW_V_3_0_PLUS:
24
+ from airflow.sdk.bases.decorator import task_decorator_factory
25
+ else:
26
+ from airflow.decorators.base import task_decorator_factory # type: ignore[no-redef]
27
+
28
+ from airflow.providers.standard.decorators.python import _PythonDecoratedOperator
29
+ from airflow.providers.standard.operators.python import ShortCircuitOperator
30
+
31
+ if TYPE_CHECKING:
32
+ from airflow.sdk.bases.decorator import TaskDecorator
33
+
34
+
35
+ class _ShortCircuitDecoratedOperator(_PythonDecoratedOperator, ShortCircuitOperator):
36
+ """Wraps a Python callable and captures args/kwargs when called for execution."""
37
+
38
+ template_fields = ShortCircuitOperator.template_fields
39
+ custom_operator_name: str = "@task.short_circuit"
40
+
41
+
42
+ def short_circuit_task(
43
+ python_callable: Callable | None = None,
44
+ multiple_outputs: bool | None = None,
45
+ **kwargs,
46
+ ) -> TaskDecorator:
47
+ """
48
+ Wrap a function into an ShortCircuitOperator.
49
+
50
+ Accepts kwargs for operator kwarg. Can be reused in a single DAG.
51
+
52
+ This function is only used only used during type checking or auto-completion.
53
+
54
+ :param python_callable: Function to decorate
55
+ :param multiple_outputs: If set to True, the decorated function's return value will be unrolled to
56
+ multiple XCom values. Dict will unroll to XCom values with its keys as XCom keys. Defaults to False.
57
+
58
+ :meta private:
59
+ """
60
+ return task_decorator_factory(
61
+ python_callable=python_callable,
62
+ multiple_outputs=multiple_outputs,
63
+ decorated_operator_class=_ShortCircuitDecoratedOperator,
64
+ **kwargs,
65
+ )
@@ -15,8 +15,7 @@
15
15
  # specific language governing permissions and limitations
16
16
  # under the License.
17
17
 
18
- # NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
19
- # OVERWRITTEN WHEN PREPARING PACKAGES.
18
+ # NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
20
19
  #
21
20
  # IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
22
21
  # `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
@@ -27,16 +26,23 @@ def get_provider_info():
27
26
  "package-name": "apache-airflow-providers-standard",
28
27
  "name": "Standard",
29
28
  "description": "Airflow Standard Provider\n",
30
- "state": "not-ready",
31
- "source-date-epoch": 1718603992,
32
- "versions": ["1.0.0"],
33
- "dependencies": ["apache-airflow>=2.10.0"],
34
29
  "integrations": [
35
30
  {
36
31
  "integration-name": "Standard",
37
32
  "external-doc-url": "https://airflow.apache.org/",
38
33
  "tags": ["apache"],
39
- "how-to-guide": ["/docs/apache-airflow-providers-standard/operators.rst"],
34
+ "how-to-guide": [
35
+ "/docs/apache-airflow-providers-standard/operators/bash.rst",
36
+ "/docs/apache-airflow-providers-standard/operators/python.rst",
37
+ "/docs/apache-airflow-providers-standard/operators/datetime.rst",
38
+ "/docs/apache-airflow-providers-standard/operators/trigger_dag_run.rst",
39
+ "/docs/apache-airflow-providers-standard/operators/latest_only.rst",
40
+ "/docs/apache-airflow-providers-standard/sensors/bash.rst",
41
+ "/docs/apache-airflow-providers-standard/sensors/python.rst",
42
+ "/docs/apache-airflow-providers-standard/sensors/datetime.rst",
43
+ "/docs/apache-airflow-providers-standard/sensors/file.rst",
44
+ "/docs/apache-airflow-providers-standard/sensors/external_task_sensor.rst",
45
+ ],
40
46
  }
41
47
  ],
42
48
  "operators": [
@@ -46,6 +52,12 @@ def get_provider_info():
46
52
  "airflow.providers.standard.operators.datetime",
47
53
  "airflow.providers.standard.operators.weekday",
48
54
  "airflow.providers.standard.operators.bash",
55
+ "airflow.providers.standard.operators.python",
56
+ "airflow.providers.standard.operators.empty",
57
+ "airflow.providers.standard.operators.trigger_dagrun",
58
+ "airflow.providers.standard.operators.latest_only",
59
+ "airflow.providers.standard.operators.smooth",
60
+ "airflow.providers.standard.operators.branch",
49
61
  ],
50
62
  }
51
63
  ],
@@ -58,7 +70,77 @@ def get_provider_info():
58
70
  "airflow.providers.standard.sensors.time",
59
71
  "airflow.providers.standard.sensors.weekday",
60
72
  "airflow.providers.standard.sensors.bash",
73
+ "airflow.providers.standard.sensors.python",
74
+ "airflow.providers.standard.sensors.filesystem",
75
+ "airflow.providers.standard.sensors.external_task",
76
+ ],
77
+ }
78
+ ],
79
+ "hooks": [
80
+ {
81
+ "integration-name": "Standard",
82
+ "python-modules": [
83
+ "airflow.providers.standard.hooks.filesystem",
84
+ "airflow.providers.standard.hooks.package_index",
85
+ "airflow.providers.standard.hooks.subprocess",
86
+ ],
87
+ }
88
+ ],
89
+ "triggers": [
90
+ {
91
+ "integration-name": "Standard",
92
+ "python-modules": [
93
+ "airflow.providers.standard.triggers.external_task",
94
+ "airflow.providers.standard.triggers.file",
95
+ "airflow.providers.standard.triggers.temporal",
61
96
  ],
62
97
  }
63
98
  ],
99
+ "extra-links": [
100
+ "airflow.providers.standard.operators.trigger_dagrun.TriggerDagRunLink",
101
+ "airflow.providers.standard.sensors.external_task.ExternalDagLink",
102
+ ],
103
+ "config": {
104
+ "standard": {
105
+ "description": "Options for the standard provider operators.",
106
+ "options": {
107
+ "venv_install_method": {
108
+ "description": "Which python tooling should be used to install the virtual environment.\n\nThe following options are available:\n- ``auto``: Automatically select, use ``uv`` if available, otherwise use ``pip``.\n- ``pip``: Use pip to install the virtual environment.\n- ``uv``: Use uv to install the virtual environment. Must be available in environment PATH.\n",
109
+ "version_added": None,
110
+ "type": "string",
111
+ "example": "uv",
112
+ "default": "auto",
113
+ }
114
+ },
115
+ }
116
+ },
117
+ "task-decorators": [
118
+ {"class-name": "airflow.providers.standard.decorators.python.python_task", "name": "python"},
119
+ {"class-name": "airflow.providers.standard.decorators.bash.bash_task", "name": "bash"},
120
+ {
121
+ "class-name": "airflow.providers.standard.decorators.branch_external_python.branch_external_python_task",
122
+ "name": "branch_external_python",
123
+ },
124
+ {
125
+ "class-name": "airflow.providers.standard.decorators.branch_python.branch_task",
126
+ "name": "branch",
127
+ },
128
+ {
129
+ "class-name": "airflow.providers.standard.decorators.branch_virtualenv.branch_virtualenv_task",
130
+ "name": "branch_virtualenv",
131
+ },
132
+ {
133
+ "class-name": "airflow.providers.standard.decorators.external_python.external_python_task",
134
+ "name": "external_python",
135
+ },
136
+ {
137
+ "class-name": "airflow.providers.standard.decorators.python_virtualenv.virtualenv_task",
138
+ "name": "virtualenv",
139
+ },
140
+ {"class-name": "airflow.providers.standard.decorators.sensor.sensor_task", "name": "sensor"},
141
+ {
142
+ "class-name": "airflow.providers.standard.decorators.short_circuit.short_circuit_task",
143
+ "name": "short_circuit",
144
+ },
145
+ ],
64
146
  }
@@ -0,0 +1,16 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
@@ -0,0 +1,89 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ from __future__ import annotations
19
+
20
+ from pathlib import Path
21
+ from typing import Any
22
+
23
+ from airflow.hooks.base import BaseHook
24
+
25
+
26
+ class FSHook(BaseHook):
27
+ """
28
+ Allows for interaction with an file server.
29
+
30
+ Connection should have a name and a path specified under extra:
31
+
32
+ example:
33
+ Connection Id: fs_test
34
+ Connection Type: File (path)
35
+ Host, Schema, Login, Password, Port: empty
36
+ Extra: {"path": "/tmp"}
37
+ """
38
+
39
+ conn_name_attr = "fs_conn_id"
40
+ default_conn_name = "fs_default"
41
+ conn_type = "fs"
42
+ hook_name = "File (path)"
43
+
44
+ @classmethod
45
+ def get_connection_form_widgets(cls) -> dict[str, Any]:
46
+ """Return connection widgets to add to connection form."""
47
+ from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
48
+ from flask_babel import lazy_gettext
49
+ from wtforms import StringField
50
+
51
+ return {"path": StringField(lazy_gettext("Path"), widget=BS3TextFieldWidget())}
52
+
53
+ @classmethod
54
+ def get_ui_field_behaviour(cls) -> dict[str, Any]:
55
+ """Return custom field behaviour."""
56
+ return {
57
+ "hidden_fields": ["host", "schema", "port", "login", "password", "extra"],
58
+ "relabeling": {},
59
+ "placeholders": {},
60
+ }
61
+
62
+ def __init__(self, fs_conn_id: str = default_conn_name, **kwargs):
63
+ super().__init__(**kwargs)
64
+ conn = self.get_connection(fs_conn_id)
65
+ self.basepath = conn.extra_dejson.get("path", "")
66
+ self.conn = conn
67
+
68
+ def get_conn(self) -> None:
69
+ pass
70
+
71
+ def get_path(self) -> str:
72
+ """
73
+ Get the path to the filesystem location.
74
+
75
+ :return: the path.
76
+ """
77
+ return self.basepath
78
+
79
+ def test_connection(self):
80
+ """Test File connection."""
81
+ try:
82
+ p = self.get_path()
83
+ if not p:
84
+ return False, "File Path is undefined."
85
+ if not Path(p).exists():
86
+ return False, f"Path {p} does not exist."
87
+ return True, f"Path {p} is existing."
88
+ except Exception as e:
89
+ return False, str(e)
@@ -0,0 +1,95 @@
1
+ #
2
+ # Licensed to the Apache Software Foundation (ASF) under one
3
+ # or more contributor license agreements. See the NOTICE file
4
+ # distributed with this work for additional information
5
+ # regarding copyright ownership. The ASF licenses this file
6
+ # to you under the Apache License, Version 2.0 (the
7
+ # "License"); you may not use this file except in compliance
8
+ # with the License. You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing,
13
+ # software distributed under the License is distributed on an
14
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15
+ # KIND, either express or implied. See the License for the
16
+ # specific language governing permissions and limitations
17
+ # under the License.
18
+ """Hook for additional Package Indexes (Python)."""
19
+
20
+ from __future__ import annotations
21
+
22
+ import subprocess
23
+ from typing import Any
24
+ from urllib.parse import quote, urlparse
25
+
26
+ from airflow.hooks.base import BaseHook
27
+
28
+
29
+ class PackageIndexHook(BaseHook):
30
+ """Specify package indexes/Python package sources using Airflow connections."""
31
+
32
+ conn_name_attr = "pi_conn_id"
33
+ default_conn_name = "package_index_default"
34
+ conn_type = "package_index"
35
+ hook_name = "Package Index (Python)"
36
+
37
+ def __init__(self, pi_conn_id: str = default_conn_name, **kwargs) -> None:
38
+ super().__init__(**kwargs)
39
+ self.pi_conn_id = pi_conn_id
40
+ self.conn = None
41
+
42
+ @staticmethod
43
+ def get_ui_field_behaviour() -> dict[str, Any]:
44
+ """Return custom field behaviour."""
45
+ return {
46
+ "hidden_fields": ["schema", "port", "extra"],
47
+ "relabeling": {"host": "Package Index URL"},
48
+ "placeholders": {
49
+ "host": "Example: https://my-package-mirror.net/pypi/repo-name/simple",
50
+ "login": "Username for package index",
51
+ "password": "Password for package index (will be masked)",
52
+ },
53
+ }
54
+
55
+ @staticmethod
56
+ def _get_basic_auth_conn_url(index_url: str, user: str | None, password: str | None) -> str:
57
+ """Return a connection URL with basic auth credentials based on connection config."""
58
+ url = urlparse(index_url)
59
+ host = url.netloc.split("@")[-1]
60
+ if user:
61
+ if password:
62
+ host = f"{quote(user)}:{quote(password)}@{host}"
63
+ else:
64
+ host = f"{quote(user)}@{host}"
65
+ return url._replace(netloc=host).geturl()
66
+
67
+ def get_conn(self) -> Any:
68
+ """Return connection for the hook."""
69
+ return self.get_connection_url()
70
+
71
+ def get_connection_url(self) -> Any:
72
+ """Return a connection URL with embedded credentials."""
73
+ conn = self.get_connection(self.pi_conn_id)
74
+ index_url = conn.host
75
+ if not index_url:
76
+ raise ValueError("Please provide an index URL.")
77
+ return self._get_basic_auth_conn_url(index_url, conn.login, conn.password)
78
+
79
+ def test_connection(self) -> tuple[bool, str]:
80
+ """Test connection to package index url."""
81
+ conn_url = self.get_connection_url()
82
+ proc = subprocess.run(
83
+ ["pip", "search", "not-existing-test-package", "--no-input", "--index", conn_url],
84
+ check=False,
85
+ capture_output=True,
86
+ )
87
+ conn = self.get_connection(self.pi_conn_id)
88
+ if proc.returncode not in [
89
+ 0, # executed successfully, found package
90
+ 23, # executed successfully, didn't find any packages
91
+ # (but we do not expect it to find 'not-existing-test-package')
92
+ ]:
93
+ return False, f"Connection test to {conn.host} failed. Error: {str(proc.stderr)}"
94
+
95
+ return True, f"Connection to {conn.host} tested successfully!"
@@ -0,0 +1,119 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
17
+ from __future__ import annotations
18
+
19
+ import contextlib
20
+ import os
21
+ import signal
22
+ from collections import namedtuple
23
+ from collections.abc import Iterator
24
+ from subprocess import PIPE, STDOUT, Popen
25
+ from tempfile import TemporaryDirectory, gettempdir
26
+
27
+ from airflow.hooks.base import BaseHook
28
+
29
+ SubprocessResult = namedtuple("SubprocessResult", ["exit_code", "output"])
30
+
31
+
32
+ @contextlib.contextmanager
33
+ def working_directory(cwd: str | None = None) -> Iterator[str]:
34
+ """
35
+ Context manager for handling (temporary) working directory.
36
+
37
+ Use the given cwd as working directory, if provided.
38
+ Otherwise, create a temporary directory.
39
+ """
40
+ with contextlib.ExitStack() as stack:
41
+ if cwd is None:
42
+ cwd = stack.enter_context(TemporaryDirectory(prefix="airflowtmp"))
43
+ yield cwd
44
+
45
+
46
+ class SubprocessHook(BaseHook):
47
+ """Hook for running processes with the ``subprocess`` module."""
48
+
49
+ def __init__(self, **kwargs) -> None:
50
+ self.sub_process: Popen[bytes] | None = None
51
+ super().__init__(**kwargs)
52
+
53
+ def run_command(
54
+ self,
55
+ command: list[str],
56
+ env: dict[str, str] | None = None,
57
+ output_encoding: str = "utf-8",
58
+ cwd: str | None = None,
59
+ ) -> SubprocessResult:
60
+ """
61
+ Execute the command.
62
+
63
+ If ``cwd`` is None, execute the command in a temporary directory which will be cleaned afterwards.
64
+ If ``env`` is not supplied, ``os.environ`` is passed
65
+
66
+ :param command: the command to run
67
+ :param env: Optional dict containing environment variables to be made available to the shell
68
+ environment in which ``command`` will be executed. If omitted, ``os.environ`` will be used.
69
+ Note, that in case you have Sentry configured, original variables from the environment
70
+ will also be passed to the subprocess with ``SUBPROCESS_`` prefix. See:
71
+ https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/logging-monitoring/errors.html for details.
72
+ :param output_encoding: encoding to use for decoding stdout
73
+ :param cwd: Working directory to run the command in.
74
+ If None (default), the command is run in a temporary directory.
75
+ :return: :class:`namedtuple` containing ``exit_code`` and ``output``, the last line from stderr
76
+ or stdout
77
+ """
78
+ self.log.info("Tmp dir root location: %s", gettempdir())
79
+ with working_directory(cwd=cwd) as cwd:
80
+
81
+ def pre_exec():
82
+ # Restore default signal disposition and invoke setsid
83
+ for sig in ("SIGPIPE", "SIGXFZ", "SIGXFSZ"):
84
+ if hasattr(signal, sig):
85
+ signal.signal(getattr(signal, sig), signal.SIG_DFL)
86
+ os.setsid()
87
+
88
+ self.log.info("Running command: %s", command)
89
+
90
+ self.sub_process = Popen(
91
+ command,
92
+ stdout=PIPE,
93
+ stderr=STDOUT,
94
+ cwd=cwd,
95
+ env=env if env or env == {} else os.environ,
96
+ preexec_fn=pre_exec,
97
+ )
98
+
99
+ self.log.info("Output:")
100
+ line = ""
101
+ if self.sub_process is None:
102
+ raise RuntimeError("The subprocess should be created here and is None!")
103
+ if self.sub_process.stdout is not None:
104
+ for raw_line in iter(self.sub_process.stdout.readline, b""):
105
+ line = raw_line.decode(output_encoding, errors="backslashreplace").rstrip()
106
+ self.log.info("%s", line)
107
+
108
+ self.sub_process.wait()
109
+
110
+ self.log.info("Command exited with return code %s", self.sub_process.returncode)
111
+ return_code: int = self.sub_process.returncode
112
+
113
+ return SubprocessResult(exit_code=return_code, output=line)
114
+
115
+ def send_sigterm(self):
116
+ """Send SIGTERM signal to ``self.sub_process`` if one exists."""
117
+ self.log.info("Sending SIGTERM signal to process group")
118
+ if self.sub_process and hasattr(self.sub_process, "pid"):
119
+ os.killpg(os.getpgid(self.sub_process.pid), signal.SIGTERM)