Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add port management for supervisor #87

Merged
merged 1 commit into from
Feb 16, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion airflow_supervisor/airflow/ssh.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
from airflow.models.dag import DAG
from airflow.models.operator import Operator
from airflow.providers.ssh.operators.ssh import SSHOperator
from airflow_balancer import Host
from airflow_balancer import Host, Port
from supervisor_pydantic.convenience import SupervisorTaskStep

from airflow_supervisor.config import SupervisorSSHAirflowConfiguration
Expand All @@ -21,6 +21,7 @@
dag: DAG,
cfg: SupervisorSSHAirflowConfiguration,
host: "Host" = None,
port: "Port" = None,
**kwargs,
):
for attr in ("command_prefix",):
Expand Down Expand Up @@ -56,7 +57,15 @@

# Ensure host matches the configuration
cfg.convenience.host = host.name

if port:
# Ensure port matches the configuration
cfg.convenience.port = f"*:{port.port}"

Check warning on line 63 in airflow_supervisor/airflow/ssh.py

View check run for this annotation

Codecov / codecov/patch

airflow_supervisor/airflow/ssh.py#L63

Added line #L63 was not covered by tests

if host or port:
# revalidate
cfg._setup_convenience_defaults()

super().__init__(dag=dag, cfg=cfg, **kwargs)

def get_step_kwargs(self, step: SupervisorTaskStep) -> Dict:
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ develop = [
# airflow-config integration
"airflow-config>=0.1.2",
# airflow-balancer integration
"airflow-balancer>=0.3.0",
"airflow-balancer>=0.3.1",
]

[project.scripts]
Expand Down