from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from operators.dbt.dbt_base import BaseDBTOperator
[docs]
class DBTRunOperator(BaseDBTOperator):
"""Execute `dbt run` on the remote server."""
ui_color = "#9933FF"
ui_fgcolor = "#FFFFFF"
custom_operator_name = "DBTRunOperator"
template_fields = BaseDBTOperator.template_fields + ("dbt_dry_run", "full_refresh")
@apply_defaults
def __init__(
self,
dry_run: bool = False,
full_refresh: bool = False,
skip_deps: bool = False,
**kwargs,
):
"""
Initializes the DBTRunOperator.
:param dry_run: If True, runs `dbt compile` instead of `dbt run` for a dry run.
:type dry_run: bool
:param select: List of dbt model selectors to include in the run. if this list is None, it will run all models. if it is not empty, it will run only the models that match the selectors. if it has .sql extension, it will be removed.
:type select: Optional[List[str]]
:default select: None
:param exclude: List of dbt model selectors to exclude from the run. if this list is None, it will run all models. if it is not empty, it will exclude only the models that match the selectors. if it has .sql extension, it will be removed.
:type exclude: Optional[List[str]]
:default exclude: None
:param project_directory: Path to the dbt project directory. it is usually the folder where dbt_project.yml is located. if this is not set, it will be assigned based on the dag folder. `os.path.join(context["dag"].folder, "dbt_project")`
:type project_directory: str
:default project_directory: None
:param vars: Dictionary of dbt variables to pass to the command. for example, if you want to pass a variable to the dbt model, you can use this parameter. it will be passed to the dbt command as --vars '{"key": "value"}'. if this is not set, it will not pass any variables. we have a series of default variables which you can use, like `vars={"target_grapi_schema": "arkansas", "target_grapi_database": "postgres_grapi_dev"}`. For the complete list of variables, see the dbt project documentation on confluence or see `airflow/dbt/README.md`
:type vars: Optional[dict]
:default vars: None
:param target: dbt target environment. This is the target environment for dbt, which can be "dev", "prod", or "stage". It is used to determine which dbt profile to use. If you are in a development environment, you should use "dev" target. If you are in a production environment, you should use "prod" target. If you are in a staging environment, you should use "stage" target. This will be validated in the get_target method. in the airflow dev environment, only "dev" target is allowed. If you try to use "prod" or "stage" target, it will raise an AirflowSkipException.
:raises AirflowSkipException: If a non-dev target is used in a dev environment.
:type target: Literal["dev", "prod", "stage"]
:default target: dev
:param profile: dbt profile name. This is the profile name that dbt will use to connect to the database. It can be one of the following: "athena", "game", "grapi", or "default". If you are not sure which profile to use, you can use "athena" profile. The profile determines how dbt connects to the database and which credentials to use. To see all the profiles available, you can check the `airflow/dbt/shared/profiles_sample.yml` file or the dbt project documentation on confluence.
:type profile: Literal["athena", "game", "grapi", "default"]
:default profile: athena
:param profile_directory: Path to dbt profiles directory. You probably don't want to change this unless you have a custom profiles.yml file. By default, it will be set to `~/.dbt/profiles.yml` on the dbt vm.
:type profile_directory: str
:default profile_directory: None
:param verbose: Enables verbose logging and adds --debug and --log-path flags.
:type verbose: bool
:default verbose: True
"""
super().__init__(skip_deps=skip_deps, **kwargs)
# Store raw values — may be Jinja template strings at parse time.
# Conversion to bool happens in execute() after Airflow renders templates.
self.dbt_dry_run = dry_run
self.full_refresh = full_refresh
[docs]
def execute(self, context: Context):
"""Executes the `dbt run` command on the remote server.
Runs the parent class's execute method, builds and runs the `dbt run` command with
specified models and variables, logs compiled SQL on failure, and cleans up the target
directory.
:param context: Airflow task context.
:type context: Context
:raises Exception: If the `dbt run` command fails, after logging compiled SQL.
:return: Dictionary containing job details including job ID, target path, command run, and models executed.
"""
import time
super().execute(context)
# Convert templated boolean fields (strings after Jinja rendering) to actual bools
if isinstance(self.dbt_dry_run, str):
self.dbt_dry_run = self.dbt_dry_run.lower() == "true"
if isinstance(self.full_refresh, str):
self.full_refresh = self.full_refresh.lower() == "true"
if self.full_refresh:
self.log.info("Performing a full refresh of incremental models.")
super().run_dbt_refresh_sources()
dbt_command = "dbt compile" if self.dbt_dry_run else "dbt run"
if self.full_refresh:
dbt_command += " --full-refresh"
cmd = " ".join(self.build_command(dbt_command))
clock_start = time.time()
try:
self.log.info(f"Running: {cmd}")
if self.verbose:
self.log.info(f"dbt run job_id = {self.job_id}")
# When we create a new table, DBT doesn't get notified automatically.
# So we try to run the command, and if compilation fails, we refresh the source tables and try again.
# We don't refresh sources every time because it adds ~45 seconds to the run time.
try:
self.run_ssh_command(cmd)
except Exception as e:
print(e)
# If less than a minute has passed since we started, try refreshing sources and re-running
# If its been more than a minute, then it isn't a compilation issue, so just fail out
if time.time() - clock_start < 60:
# Before failing out, try to refresh sources and run again
self.log.info(
"dbt run failed, attempting to refresh sources and re-run..."
)
super().run_dbt_refresh_sources()
self.run_ssh_command(cmd)
else:
raise
except Exception:
# On failure, dump out all compiled SQL for debugging
self.log.error("dbt run failed—showing compiled SQL:")
dump = (
f"find {self.target_path}/compiled "
"-type f -name '*.sql' | sort | while read file; do "
'echo "=== $file ==="; cat "$file"; echo; done'
)
self.run_ssh_command(dump)
raise
finally:
# We want to get the list of models that were run, but the ssh command output only
# returns the last line of output, so we encode to base64 to get it all in one line.
import base64
models_run_output = self.run_ssh_command(
f"""source /home/ubuntu/dbt-venv/bin/activate && \
shopt -s globstar && \
if ls {self.target_path}/run/*/models/**/*.sql 2>/dev/null; then
ls {self.target_path}/run/*/models/**/*.sql -m | base64 -w 0
else
echo ""
fi"""
)
decoded_output = base64.b64decode(models_run_output).decode("utf-8")
models_run = [
model.strip().split("/")[-1].replace(".sql", "")
for model in decoded_output.split(",")
if ".yml" not in model
]
self.log.info(f"Models run: {models_run}")
# self.run_ssh_command(
# f'[ -e {self.target_path} ] && rm -rf {self.target_path}'
# )
pass
self.log.info("dbt run completed successfully.")
return {
"job_id": self.job_id,
"target_path": self.target_path,
"dbt_dry_run": self.dbt_dry_run,
"dbt_command": dbt_command,
"models_run": models_run,
}