import shutil
import os
import uuid
from config import Config
from typing import List, Literal, Optional
from functools import cached_property
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException, AirflowSkipException
from airflow.hooks.subprocess import SubprocessHook
from airflow.providers.slack.hooks.slack import SlackHook
[docs]
class BaseDBTOperator(BaseOperator):
"""Initializes the BaseDBTOperator
: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: name of target (database connection details; this is default entry in `outputs`). https://docs.getdbt.com/docs/core/connect-data-platform/profiles.yml and https://docs.getdbt.com/docs/core/connect-data-platform/connection-profiles: `target` is for default target, `outputs` is for additional targets. 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
:param args: Additional arguments for BaseOperator.
:type args: Any
:param kwargs: Additional keyword arguments for BaseOperator.
:type kwargs: Any
NOTE: super.execute() should be called in the execute method of the subclass. it will set the job_id and target_path and self.project_director attributes.
"""
ui_color = "#9933FF"
ui_fgcolor = "#FFFFFF"
custom_operator_name = "BaseDBTOperator"
template_fields = (
"project_directory",
"target",
"profile",
"profile_directory",
"vars",
"select",
"exclude",
)
target: str
"""The target environment for dbt, which can be "dev", "prod", or "stage". It is used to determine which dbt profile to use."""
profile: str
"""The dbt profile name that dbt will use to connect to the database. It can be one of the following: "athena", "game", "grapi" """
project_directory: str
"""Path to the dbt project directory. It is usually the folder where dbt_project.yml is located."""
profile_directory: str
"""Path to dbt profiles directory. You probably don't want to change this unless you have a custom profiles.yml file."""
vars: Optional[dict]
"""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."""
select: Optional[List[str]]
"""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."""
exclude: Optional[List[str]]
"""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."""
verbose: bool
"""Enables verbose logging and adds --debug and --log-path flags."""
job_id: str
"""Unique identifier for the dbt job run, used to create a unique target path."""
# def __init_subclass__(cls, **kwargs):
# super().__init_subclass__(**kwargs)
# if cls.__doc__ is not None:
# cls.__doc__ = f"{cls.__doc__} \n\n Subclass of BaseDBTOperator\n\n{BaseDBTOperator.__doc__}"
@apply_defaults
def __init__(
self,
select: Optional[List[str]] = None,
exclude: Optional[List[str]] = None,
project_directory: str = None,
vars: Optional[dict] = None,
target: Literal["dev", "prod", "stage"] = "dev",
profile: Literal[
"athena",
"game",
"grapi",
"default",
] = "athena",
profile_directory: str = None,
verbose: bool = True,
skip_deps: bool = False,
*args,
**kwargs,
):
"""Initializes the BaseDBTOperator
: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: name of target (database connection details; this is default entry in `outputs`). https://docs.getdbt.com/docs/core/connect-data-platform/profiles.yml and https://docs.getdbt.com/docs/core/connect-data-platform/connection-profiles: `target` is for default target, `outputs` is for additional targets. 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
:param skip_deps: If True, skip `dbt deps` and `dbt debug` during setup. Use when a
prior task in the same DAG run has already installed dependencies (e.g. a dedicated
setup task), allowing parallel model tasks to avoid concurrent writes to dbt_packages/.
:type skip_deps: bool
:default skip_deps: False
:param args: Additional arguments for BaseOperator.
:type args: Any
:param kwargs: Additional keyword arguments for BaseOperator.
:type kwargs: Any
NOTE: super.execute() should be called in the execute method of the subclass. it will set the job_id and target_path and self.project_director attributes.
"""
super().__init__(*args, **kwargs)
self.dag_root = Config.DAG_ROOT
self.dbt_server_folder = Config.DBT_SERVER_FOLDER
self.dbt_ssh_key = Config.DBT_SSH_KEY
self.dbt_ssh_user = Config.DBT_SSH_USER
self.dbt_ssh_host = Config.DBT_ALIAS
self.dbt_venv_path = Config.DBT_VENV_PATH
self.project_directory = project_directory
self.target = target
self.profile = profile
self.profile_directory = profile_directory # for example if you want to provide custom profile f"{self.dbt_server_folder}/dbt/shared" by default it will be '~/.dbt/profiles.yml'
self.job_id = None # we will set this in execute method
self.target_path = None # we will set this in execute method
self.select = select
self.exclude = exclude
self.vars = vars
self.verbose = verbose
self.skip_deps = skip_deps
def _validate_selection(self) -> None:
"""Validates the select and exclude lists to ensure they do not contain .sql extensions.
:raises AirflowException: If any selector contains a .sql extension.
"""
select = self.select
exclude = self.exclude
self.select = [select] if isinstance(select, str) else select
self.select = (
[m.replace(".sql", "") for m in (self.select or [])]
if self.select
else None
)
self.exclude = [exclude] if isinstance(exclude, str) else exclude
self.exclude = (
[m.replace(".sql", "") for m in (self.exclude or [])]
if self.exclude
else None
)
[docs]
def get_target(self) -> str:
"""Validates and returns the dbt target environment.
Ensures non-dev targets are blocked in development environments.
:param target: The dbt target environment.
:type target: str
:return: The validated target.
:rtype: str
:raises AirflowSkipException: If a non-dev target is used in a dev environment.
"""
if not Config.PRODUCTION and not self.target.endswith("dev"):
raise AirflowSkipException(
f"Target {self.target} is not allowed in dev environment. Please use dev target."
)
print(f"Using target: {self.target}")
return self.target
[docs]
def apply_default_vars(self) -> None:
"""Applies default variables to the operator.
This method is called to set default variables for the operator.
It can be overridden in subclasses to provide custom default variables.
:param context: Airflow task context.
:type context: Context
"""
if not self.vars:
self.vars = {}
default_vars = {"dest_bucket": Config.DATA_LAKE_BUCKET}
target = self.get_target()
if target == "dev":
default_vars = {
**default_vars,
**{
"target_grapi_database": (
"postgres_grapi_dev" if self.profile == "athena" else ""
), # we use postgres_grapi_dev for athena profile since it will convert it to postgres_grapi_dev.schema.table for cross db queries. if it is not athena profile, it will be empty string so it will be converted to grapidb.schema.table
"target_game_database": (
"postgres_game_dev" if self.profile == "athena" else ""
),
"target_lake_schema": "shoreline",
},
}
elif target == "stage":
default_vars = {
**default_vars,
**{
"target_grapi_database": (
"postgres_grapi_stage" if self.profile == "athena" else ""
),
"target_game_database": (
"postgres_game_stage" if self.profile == "athena" else ""
),
"target_lake_schema": "readydatalake",
},
}
else:
default_vars = {
**default_vars,
**{
"target_grapi_database": (
"postgres_grapi_prod" if self.profile == "athena" else ""
),
"target_game_database": (
"postgres_game_prod" if self.profile == "athena" else ""
),
"target_lake_schema": "readydatalake",
},
}
self.vars = {**default_vars, **self.vars}
@cached_property
def subprocess_hook(self):
"""Returns hook for running the bash command."""
return SubprocessHook()
[docs]
def run_ssh_command(
self,
command: str,
) -> str:
"""
ssh command to connect to the server
"""
ssh_target = f"{self.dbt_ssh_user}@{self.dbt_ssh_host}"
remote = f"source {self.dbt_venv_path} && {command}"
remote_escaped = remote.replace('"', '\\"')
ssh_cmd = f'ssh -i {self.dbt_ssh_key} {ssh_target} "{remote_escaped}"'
return self._run_bash_command(ssh_cmd)
def _run_bash_command(
self,
command: str,
) -> str:
"""
runs a command with subprocess hook
:param command: the command
:type command: str
:raises AirflowException: if the command fails
"""
system_env = os.environ.copy()
bash_path = shutil.which("bash") or "bash"
bash_result = self.subprocess_hook.run_command(
command=[bash_path, "-c", command], env=system_env
)
if bash_result.exit_code:
raise AirflowException(
f"Bash command failed. The command returned a non-zero exit code {bash_result.exit_code}. {bash_result.output}"
)
return bash_result.output
[docs]
def on_kill(self) -> None:
self.subprocess_hook.send_sigterm()
[docs]
def notify_slack(self, message) -> None:
slack_hook = SlackHook(slack_conn_id=Config.SLACK_DEFAULT_API)
slack_hook.call(
"chat.postMessage",
json={
"channel": "#dbt-notification",
"text": message,
},
)
[docs]
def assign_project_folder(self, context: Context) -> None:
"""Assigns the project folder based on the DAG folder if not set.
:param context: Airflow task context containing DAG information.
:type context: Context
"""
if not self.project_directory:
self.project_directory = os.path.join(context["dag"].folder, "dbt_project")
[docs]
def get_target_path(self) -> str:
"""Generates a unique temporary path for dbt run output.
:return: Path in the format /tmp/dbt-target/<job_id>.
:rtype: str
"""
job_id = self.job_id
return os.path.join("/tmp/dbt-target", job_id)
[docs]
def build_command(self, dbt_base_command) -> List[str]:
"""Builds the dbt command to be executed on the server.
:param dbt_base_command: The base dbt command (e.g., "dbt run", "dbt test").
:type dbt_base_command: str
:return: The complete dbt command with arguments.
:rtype: List[str]
"""
import yaml
dbt_command = [dbt_base_command]
dbt_command.extend(["--project-dir", self.project_directory])
if self.profile_directory:
dbt_command.extend(["--profiles-dir", self.profile_directory])
target = self.get_target()
dbt_command.extend(["--target-path", self.target_path])
dbt_command.extend(["--target", target])
dbt_command.extend(["--profile", self.profile])
if self.select:
dbt_command.append("--select")
dbt_command.extend(self.select)
if self.exclude:
dbt_command.append("--exclude")
dbt_command.extend(self.exclude)
if self.vars:
import json, shlex
vars_json = json.dumps(self.vars, separators=(",", ":"))
vars_arg = shlex.quote(vars_json)
dbt_command.extend(["--vars", vars_arg])
if self.verbose:
dbt_command.append("--debug")
dbt_command.append(f"--log-path {self.target_path}/logs")
print(f"DBT command: {dbt_command}")
return dbt_command
[docs]
def run_dbt_deps(self) -> None:
"""
Runs dbt deps command to install the dependencies.
"""
dbt_command = ["dbt", "deps"]
dbt_command.extend(["--project-dir", self.project_directory])
if self.profile_directory:
dbt_command.extend(["--profiles-dir", self.profile_directory])
dbt_command.extend(["--profile", self.profile])
if self.vars:
import json, shlex
vars_json = json.dumps(self.vars, separators=(",", ":"))
vars_arg = shlex.quote(vars_json)
dbt_command.extend(["--vars", vars_arg])
print(f"DBT command: {dbt_command}")
self.run_ssh_command(command=" ".join(dbt_command))
[docs]
def run_dbt_debug(self) -> None:
"""
Runs dbt debug command to check if the connection is working.
"""
dbt_command = ["dbt", "debug"]
dbt_command.extend(["--project-dir", self.project_directory])
if self.profile_directory:
dbt_command.extend(["--profiles-dir", self.profile_directory])
dbt_command.extend(["--profile", self.profile])
if self.vars:
import json, shlex
vars_json = json.dumps(self.vars, separators=(",", ":"))
vars_arg = shlex.quote(vars_json)
dbt_command.extend(["--vars", vars_arg])
print(f"DBT command: {dbt_command}")
self.run_ssh_command(command=" ".join(dbt_command))
[docs]
def run_dbt_refresh_sources(self) -> None:
"""
Runs the script that gets all tables and write the sources.yml file.
"""
command = "python airflow/dbt/generate_sources.py"
self.run_ssh_command(command=command)
[docs]
def get_job_id(self, context: Context) -> str:
"""Retrieves the job ID from the context or generates a new one.
Includes the task_id in the seed so that parallel tasks within the
same DAG run get unique target paths and don't collide on artifacts.
:param context: Airflow task context.
:type context: Context
:return: The job ID.
:rtype: str
"""
execution_time = context["execution_date"]
execution_date_str = execution_time.isoformat()
task_id = context["task"].task_id
# Generate UUID using UUID5 (deterministic, requires a namespace)
namespace = uuid.NAMESPACE_DNS
uuid_value = uuid.uuid5(namespace, f"{execution_date_str}:{task_id}")
return str(uuid_value)
[docs]
def execute(self, context):
"""Executes the dbt operator logic.
Performs setup, runs `dbt deps` and `dbt debug`, and logs relevant information.
NOTE: this method should be called in the execute method of the subclass. it will set the job_id and target_path and self.project_director attributes.
:param context: Airflow task context.
:type context: Context
:raises AirflowSkipException: If target or project_directory is not set.
"""
self._validate_selection()
self.apply_default_vars()
self.job_id = self.get_job_id(context)
self.target_path = self.get_target_path()
if self.verbose:
self.log.info(f"Job ID: {self.job_id}")
self.log.info(f"Target path: {self.target_path}")
self.log.info("Executing BaseDBTOperator")
self.log.info(f"Server folder: {self.dbt_server_folder}")
self.log.info(f"Dag root: {self.dag_root}")
self.log.info(f"Dag id: {context['dag'].dag_id}")
self.log.info(f"Task id: {context['task'].task_id}")
dag_folder = context["dag"].folder
self.log.info(f"Current DAG folder: {dag_folder}")
if not self.dag_id:
self.dag_id = context["dag"].dag_id
# we assign the project folder based on the dag folder
self.assign_project_folder(context)
if self.verbose:
self.log.info(f"Project folder: {self.project_directory}")
if self.profile_directory:
self.log.info(f"Profile folder: {self.profile_directory}")
self.log.info(f"Target path: {self.target_path}")
self.run_ssh_command(
command="dbt --version",
)
if not self.get_target():
raise AirflowSkipException("DBT profile is not set")
if not self.project_directory:
raise AirflowSkipException("DBT project directory is not set")
if not self.skip_deps:
self.run_dbt_deps()
self.run_dbt_debug()
else:
self.log.info("Skipping dbt deps/debug (skip_deps=True)")