Skip to content

Commit

Permalink
fix: Avoid heavy imports in operator constructor
Browse files Browse the repository at this point in the history
After a previous commit [1] improving import time on operator import,
this commit fixes the import time issues when instantiating any of the
dbt operators.

The main issue still present was the enum imports in the operator
constructors, which make DAGs slow to import any time they instantiated
one of them.

Profiling can be easily run locally, with the following command:
```shell
python -X importtime -c "from airflow_dbt_python.operators.dbt import DbtRunOperator; DbtRunOperator(task_id='test')" 2>import-times.log
```

[1] 67c8d78
  • Loading branch information
adamantike authored and tomasfarias committed Nov 25, 2022
1 parent 3ca0fa2 commit ed205ca
Show file tree
Hide file tree
Showing 4 changed files with 39 additions and 37 deletions.
35 changes: 2 additions & 33 deletions airflow_dbt_python/hooks/dbt.py
Expand Up @@ -6,7 +6,6 @@
import os
import pickle
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Any, Optional, Type, Union
from urllib.parse import urlparse
Expand Down Expand Up @@ -50,39 +49,9 @@
except ImportError:
from airflow.hooks.base_hook import BaseHook # type: ignore

from .backends import DbtBackend, StrPath, build_backend


class FromStrEnum(Enum):
"""Access enum variants with strings ensuring uppercase."""

@classmethod
def from_str(cls, s: str):
"""Instantiate an Enum from a string."""
return cls[s.replace("-", "_").upper()]


class LogFormat(FromStrEnum):
"""Allowed dbt log formats."""
from airflow_dbt_python.utils.enums import FromStrEnum, LogFormat, Output

DEFAULT = "default"
JSON = "json"
TEXT = "text"


class Output(FromStrEnum):
"""Allowed output arguments."""

JSON = "json"
NAME = "name"
PATH = "path"
SELECTOR = "selector"

def __eq__(self, other):
"""Override equality for string comparison."""
if isinstance(other, str):
return other.upper() == self.name
return Enum.__eq__(self, other)
from .backends import DbtBackend, StrPath, build_backend


def parse_yaml_args(args: Optional[Union[str, dict[str, Any]]]) -> dict[str, Any]:
Expand Down
6 changes: 2 additions & 4 deletions airflow_dbt_python/operators/dbt.py
Expand Up @@ -15,6 +15,8 @@
from airflow.models.xcom import XCOM_RETURN_KEY
from airflow.version import version

from airflow_dbt_python.utils.enums import LogFormat, Output

# apply_defaults is deprecated in version 2 and beyond. This allows us to
# support version 1 and deal with the deprecation warning.
if int(version[0]) == 1:
Expand Down Expand Up @@ -110,8 +112,6 @@ def __init__(
replace_on_push: bool = False,
**kwargs,
) -> None:
from airflow_dbt_python.hooks.dbt import LogFormat

super().__init__(**kwargs)
self.project_dir = project_dir
self.profiles_dir = profiles_dir
Expand Down Expand Up @@ -618,8 +618,6 @@ def __init__(
indirect_selection: Optional[str] = None,
**kwargs,
) -> None:
from airflow_dbt_python.hooks.dbt import Output

super().__init__(**kwargs)
self.resource_types = resource_types
self.select = select
Expand Down
1 change: 1 addition & 0 deletions airflow_dbt_python/utils/__init__.py
@@ -0,0 +1 @@
"""Utils for project operators and hooks."""
34 changes: 34 additions & 0 deletions airflow_dbt_python/utils/enums.py
@@ -0,0 +1,34 @@
"""Enumerations with allowed set of values."""
from enum import Enum


class FromStrEnum(Enum):
"""Access enum variants with strings ensuring uppercase."""

@classmethod
def from_str(cls, s: str):
"""Instantiate an Enum from a string."""
return cls[s.replace("-", "_").upper()]


class LogFormat(FromStrEnum):
"""Allowed dbt log formats."""

DEFAULT = "default"
JSON = "json"
TEXT = "text"


class Output(FromStrEnum):
"""Allowed output arguments."""

JSON = "json"
NAME = "name"
PATH = "path"
SELECTOR = "selector"

def __eq__(self, other):
"""Override equality for string comparison."""
if isinstance(other, str):
return other.upper() == self.name
return Enum.__eq__(self, other)

0 comments on commit ed205ca

Please sign in to comment.