Passing a SFTPHook to SFTPOperator results in a scheduler error if the connection doesn't exist

[camhom]

Apache Airflow version

2.9.2

If "Other Airflow 2 version" selected, which one?

No response

What happened?

This is an issue where DAG scheduler errors are showing up in our local development airflow containers due to missing connections.

To give some background: We had to work around and override the retrieve_file method for SFTPHook due to performance issues with large CSV file loads.

This is what the override looks like (setting max_concurrent_prefetch_requests to 64):

class BCSSFTPHook(SFTPHook):
    def retrieve_file(self, remote_full_path: str, local_full_path: str, prefetch: bool = True) -> None:
        # mainly copy/pasted from SFTPHook, passing a limit to get method
        conn = self.get_conn()
        conn.get(
            remote_full_path, 
            local_full_path, 
            prefetch=prefetch, 
            max_concurrent_prefetch_requests=64)

This then we use it in our DAG like so:

with DAG(...):
  t1 = SFTPOperator(
    task_id='data_download',
    sftp_hook=BCSSFTPHook(ssh_conn_id='my-connection-id'),
    local_filepath=Files.generate_dump_file_path(filename=dag_id, include_date=True),
    remote_filepath='some_path/file.csv',
    operation=SFTPOperation.GET,
  )

However, the DAG scheduler in a developer environment which doesn't have the connection ID defined in the connections will show an error:

Broken DAG: [/opt/airflow/dags/dags/my-dag.py]
Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.11/site-packages/airflow/hooks/base.py", line 83, in get_connection
    conn = Connection.get_connection_from_secrets(conn_id)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/airflow/.local/lib/python3.11/site-packages/airflow/models/connection.py", line 519, in get_connection_from_secrets
    raise AirflowNotFoundException(f"The conn_id `{conn_id}` isn't defined")
airflow.exceptions.AirflowNotFoundException: The conn_id `my-connection-id` isn't defined

I believe this is due to SSHHook calling get_connection() in __init__:

airflow/airflow/providers/ssh/hooks/ssh.py

Line 154 in 3583329

conn = self.get_connection(self.ssh_conn_id)

What you think should happen instead?

Is there a way SSHHook can handle fetching a connection differently so that it catches the missing conn_id on DAG runtime instead of scheduler check interval?

How to reproduce

1. Create a DAG in an airflow environment
2. Create a Task using the SFTPOperator
3. Set SFTPOperator.sftp_hook to SFTPHook(ssh_conn_id='undefined-conn-id')
4. Run airflow server
5. Observe error at the top of the airflow web GUI

Operating System

Windows 11

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

No response

Anything else?

This problem is consistent between many developers local machines

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[potiuk]

Converted to a discussion - it's not an airflow issue.

The easiest will likely be to override SFTPOperator with your own CustomSFTPOperator and override it's execute method where you will set the self.sftp_hook just before calling super.execute(). Smth like:

def execute(context):
   self.sftp_hook = ....
   super().execute(context=context)