Skip to content

Commit

Permalink
fix: Conditionally import required hook
Browse files Browse the repository at this point in the history
This change makes the library to only import the hook it's going to need
based on the remote scheme.

As it is currently implemented, an environment that uses a S3 remote
fails if it doesn't have the `ssh` Airflow provider installer:

```
Traceback (most recent call last):
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/dbt.py", line 126, in get_remote
    return self.remotes[(scheme, conn_id)]
KeyError: ('s3', None)

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/dbt.py", line 300, in dbt_directory
    project_dir, profiles_dir = self.prepare_directory(
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/dbt.py", line 345, in prepare_directory
    project_dir_path = self.download_dbt_project(
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/dbt.py", line 158, in download_dbt_project
    remote = self.get_remote(scheme, self.project_conn_id)
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/dbt.py", line 128, in get_remote
    remote = get_remote(scheme, conn_id)
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/remote.py", line 147, in get_remote
    from .git import DbtGitRemoteHook
  File "/usr/local/airflow/.local/lib/python3.10/site-packages/airflow_dbt_python/hooks/git.py", line 5, in <module>
    from airflow.providers.ssh.hooks.ssh import SSHHook
ModuleNotFoundError: No module named 'airflow.providers.ssh'
```

The same happens when the environment uses a remote that requires the
SSH hook, and doesn't install the `aws` Airflow provider.
  • Loading branch information
adamantike authored and tomasfarias committed Mar 16, 2023
1 parent ef1449d commit abd8149
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions airflow_dbt_python/hooks/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,15 +144,17 @@ def get_remote(scheme: str, conn_id: Optional[str] = None) -> DbtRemoteHook:
In the future we should make our hooks discoverable and package ourselves as a
proper Airflow providers package.
"""
from .git import DbtGitRemoteHook
from .localfs import DbtLocalFsRemoteHook
from .s3 import DbtS3RemoteHook

if scheme == "s3":
from .s3 import DbtS3RemoteHook

remote_cls: Type[DbtRemoteHook] = DbtS3RemoteHook
elif scheme in ("https", "git", "git+ssh", "ssh", "http"):
from .git import DbtGitRemoteHook

remote_cls = DbtGitRemoteHook
elif scheme == "":
from .localfs import DbtLocalFsRemoteHook

remote_cls = DbtLocalFsRemoteHook
else:
raise NotImplementedError(f"Backend {scheme} is not supported")
Expand Down

0 comments on commit abd8149

Please sign in to comment.