Decouple hook executor from the operator, add Cloud Build execution and adjust testing

.gitignore

@@ -133,3 +133,11 @@ dmypy.json
 
 # pytype static type analyzer
 .pytype/
+
+.run/
+
+.terraform/
+
+.user.yaml
+
+.idea/

README.md

@@ -67,7 +67,7 @@ There are five operators currently implemented:
   * Calls [`dbt test`](https://docs.getdbt.com/docs/test)
 
 
-Each of the above operators accept the following arguments:
+Each of the above operators accept the arguments in [here (dbt_command_config)](airflow_dbt/dbt_command_config.py). The main ones being:
 
 * `profiles_dir`
   * If set, passed as the `--profiles-dir` argument to the `dbt` command
@@ -96,6 +96,68 @@ Typically you will want to use the `DbtRunOperator`, followed by the `DbtTestOpe
 
 You can also use the hook directly. Typically this can be used for when you need to combine the `dbt` command with another task in the same operators, for example running `dbt docs` and uploading the docs to somewhere they can be served from.
 
+## A more advanced example:
+
+If want to run your `dbt` project other tan in the airflow worker you can use
+the `DbtCloudBuildHook` and apply it to the `DbtBaseOperator` or simply use the
+provided `DbtCloudBuildOperator`:
+
+```python
+from airflow_dbt.hooks import DbtCloudBuildHook
+from airflow_dbt.operators import DbtBaseOperator, DbtCloudBuildOperator
+DbtBaseOperator(
+    task_id='provide_hook',
+    command='run',
+    use_colors=False,
+    config={
+        'profiles_dir': './jaffle-shop',
+        'project_dir': './jaffle-shop',
+    },
+    dbt_hook=DbtCloudBuildHook(
+        gcs_staging_location='gs://my-bucket/compressed-dbt-project.tar.gz'
+    )
+)
+
+DbtCloudBuildOperator(
+    task_id='default_hook_cloudbuild',
+    gcs_staging_location='gs://my-bucket/compressed-dbt-project.tar.gz',
+    command='run',
+    use_colors=False,
+    config={
+        'profiles_dir': './jaffle-shop',
+        'project_dir': './jaffle-shop',
+    },
+)
+```
+
+You can either define the dbt params/config/flags in the operator or you can 
+group them into a `config` param. They both have validation, but only the config
+has templating. The following two tasks are equivalent:
+
+```python
+from airflow_dbt.operators.dbt_operator import DbtBaseOperator
+
+DbtBaseOperator(
+    task_id='config_param',
+    command='run',
+    config={
+        'profiles_dir': './jaffle-shop',
+        'project_dir': './jaffle-shop',
+        'dbt_bin': '/usr/local/airflow/.local/bin/dbt',
+        'use_colors': False
+    }
+)
+
+DbtBaseOperator(
+    task_id='flat_config',
+    command='run',
+    profiles_dir='./jaffle-shop',
+    project_dir='./jaffle-shop',
+    dbt_bin='/usr/local/airflow/.local/bin/dbt',
+    use_colors=False
+)
+```
+
 ## Building Locally
 
 To install from the repository:
@@ -147,7 +209,9 @@ If you use MWAA, you just need to update the `requirements.txt` file and add `ai
 Then you can have your dbt code inside a folder `{DBT_FOLDER}` in the dags folder on S3 and configure the dbt task like below:
 
 ```python
-dbt_run = DbtRunOperator(
+from airflow_dbt.operators.dbt_operator import DbtRunOperator 
+
+dbt_run=DbtRunOperator(
   task_id='dbt_run',
   dbt_bin='/usr/local/airflow/.local/bin/dbt',
   profiles_dir='/usr/local/airflow/dags/{DBT_FOLDER}/',

airflow_dbt/__init__.py

@@ -1,9 +0,0 @@
-from .hooks import DbtCliHook
-from .operators import (
-    DbtSeedOperator,
-    DbtSnapshotOperator,
-    DbtRunOperator,
-    DbtTestOperator,
-    DbtDocsGenerateOperator,
-    DbtDepsOperator
-)

airflow_dbt/__version__.py

@@ -1,3 +1 @@
-VERSION = (0, 4, 0)
-
-__version__ = '.'.join(map(str, VERSION))
+__version__ = "0.5.10"

airflow_dbt/dbt_command_config.py

@@ -0,0 +1,69 @@
+import sys
+
+# Python versions older than 3.8 have the TypedDict in a different namespace.
+# In case we find ourselves in that situation, we use the `older` import
+if sys.version_info[0] == 3 and sys.version_info[1] >= 8:
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+
+class DbtCommandConfig(TypedDict, total=False):
+    """
+    Holds the structure of a dictionary containing dbt config. Provides the
+    types and names for each one, and also helps shortening the constructor
+    since we can nest it and reuse it
+    """
+    # global flags
+    version: bool
+    record_timing_info: bool
+    debug: bool
+    log_format: str  # either 'text', 'json' or 'default'
+    write_json: bool
+    strict: bool
+    warn_error: bool
+    partial_parse: bool
+    use_experimental_parser: bool
+    use_colors: bool
+    no_use_colors: bool
+
+    # per command flags
+    profiles_dir: str
+    project_dir: str
+    target: str
+    vars: dict
+    models: str
+    exclude: str
+
+    # run specific
+    full_refresh: bool
+    profile: str
+
+    # docs specific
+    no_compile: bool
+
+    # debug specific
+    config_dir: str
+
+    # ls specific
+    resource_type: str  # models, snapshots, seeds, tests, and sources.
+    select: str
+    models: str
+    exclude: str
+    selector: str
+    output: str
+    output_keys: str
+
+    # rpc specific
+    host: str
+    port: int
+
+    # run specific
+    fail_fast: bool
+
+    # run-operation specific
+    args: dict
+
+    # test specific
+    data: bool
+    schema: bool

airflow_dbt/hooks/__init__.py

@@ -1 +0,0 @@
-from .dbt_hook import DbtCliHook

airflow_dbt/hooks/base.py

@@ -0,0 +1,98 @@
+import json
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+# noinspection PyDeprecation
+from airflow.hooks.base_hook import BaseHook
+
+from airflow_dbt.dbt_command_config import DbtCommandConfig
+
+
+def render_config(config: Dict[str, Union[str, bool]]) -> List[str]:
+    """Renders a dictionary of options into a list of cli strings"""
+    dbt_command_config_annotations = DbtCommandConfig.__annotations__
+    command_params = []
+    for key, value in config.items():
+        if key not in dbt_command_config_annotations:
+            raise ValueError(f"{key} is not a valid key")
+        if value is not None:
+            param_value_type = type(value)
+            # check that the value has the correct type from dbt_command_config_annotations
+            if param_value_type != dbt_command_config_annotations[key]:
+                raise TypeError(f"{key} has to be of type {dbt_command_config_annotations[key]}")
+            # if the param is not bool it must have a non null value
+            flag_prefix = ''
+            if param_value_type is bool and not value:
+                flag_prefix = 'no-'
+            cli_param_from_kwarg = "--" + flag_prefix + key.replace("_", "-")
+            command_params.append(cli_param_from_kwarg)
+            if param_value_type is str:
+                command_params.append(value)
+            elif param_value_type is int:
+                command_params.append(str(value))
+            elif param_value_type is dict:
+                command_params.append(json.dumps(value))
+    return command_params
+
+
+def generate_dbt_cli_command(
+    dbt_bin: str,
+    command: str,
+    base_config: Dict[str, Union[str, bool]],
+    command_config: Dict[str, Union[str, bool]],
+) -> List[str]:
+    """
+    Creates a CLI string from the keys in the dictionary. If the key is none
+    it is ignored. If the key is of type boolean the name of the key is added.
+    If the key is of type string it adds the the key prefixed with tow dashes.
+    If the key is of type integer it adds the the key prefixed with three
+    dashes.
+    dbt_bin and command are mandatory.
+    Boolean flags must always be positive.
+
+    Available params are:
+    :param command_config: Specific params for the commands
+    :type command_config: dict
+    :param base_config: Params that apply to the `dbt` program regardless of
+        the command it is running
+    :type base_config: dict
+    :param command: The dbt sub-command to run
+    :type command: str
+    :param dbt_bin: Path to the dbt binary, defaults to `dbt` assumes it is
+        available in the PATH.
+    :type dbt_bin: str
+    :param command: The dbt sub command to run, for example for `dbt run`
+        the base_command will be `run`. If any other flag not contemplated
+        must be included it can also be added to this string
+    :type command: str
+    """
+    if not dbt_bin:
+        raise ValueError("dbt_bin is mandatory")
+    if not command:
+        raise ValueError("command mandatory")
+    base_params = render_config(base_config)
+    command_params = render_config(command_config)
+    # commands like 'dbt docs generate' need the command to be split in two
+    command_pieces = command.split(" ")
+    return [dbt_bin, *base_params, *command_pieces, *command_params]
+
+
+class DbtBaseHook(BaseHook, ABC):
+    """
+    Base abstract class for all DbtHooks to have a common interface and force
+    implement the mandatory `run_dbt()` function.
+    """
+
+    def __init__(self, env: Optional[Dict] = None):
+        """
+        :param env: If set will be passed over to cloud build to run in the
+            dbt step
+        :type env: dict
+        """
+        super().__init__()
+        self.env = env or {}
+
+    @abstractmethod
+    def run_dbt(self, dbt_cmd: Union[str, List[str]]):
+        """Run the dbt command"""
+        pass

airflow_dbt/hooks/cli.py

@@ -0,0 +1,51 @@
+from __future__ import print_function
+
+from typing import Any, Dict, List, Optional, Union
+
+from airflow import AirflowException
+from airflow.hooks.subprocess import SubprocessHook
+
+from airflow_dbt.hooks.base import DbtBaseHook
+
+
+class DbtCliHook(DbtBaseHook):
+    """
+    Run the dbt command in the same airflow worker the task is being run.
+    This requires the `dbt` python package to be installed in it first.
+    """
+
+    def __init__(self, env: Optional[Dict] = None):
+        """
+        :type env:
+        :param env: Environment variables that will be passed to the
+            subprocess. Must be a dictionary of key-values
+        """
+        self.sp = SubprocessHook()
+        super().__init__(env=env)
+
+    def get_conn(self) -> Any:
+        """
+        Return the subprocess connection, which isn't implemented, just for
+        conformity
+        """
+        return self.sp.get_conn()
+
+    def run_dbt(self, dbt_cmd: Union[str, List[str]]):
+        """
+         Run the dbt cli
+
+         :param dbt_cmd: The dbt whole command to run
+         :type dbt_cmd: List[str]
+         """
+        result = self.sp.run_command(
+            command=dbt_cmd,
+            env=self.env,
+        )
+
+        if result.exit_code != 0:
+            raise AirflowException(f'Error executing the DBT command: '
+                                   f'{result.output}')
+
+    def on_kill(self):
+        """Kill the open subprocess if the task gets killed by Airflow"""
+        self.sp.send_sigterm()