_dbt

#compdef dbt dbtf

# OVERVIEW
#   Adds autocompletion to dbt CLI by:
#       1. Auto-detecting whether dbt is dbt-core (Click) or dbt-fusion (clap)
#       2. Using Click's built-in completion for dbt-core commands and flags
#          OR clap_complete (via `dbt completions zsh`) for dbt-fusion — cached
#          by binary mtime so it only regenerates when the binary changes
#       3. Enhancing with manifest.json parsing for intelligent model/selector
#          completions for both dbt-core and dbt-fusion
#       4. Finding the root of the repo (identified by dbt_project.yml)
#       5. Extracting valid model selectors from target/manifest.json for:
#           -m / --model[s]
#           -s / --select
#           --exclude
#           --selector
#
# NOTE: This script uses the manifest (assumed to be at target/manifest.json)
#       to _quickly_ provide a list of existing selectors. As such, a dbt
#       resource must be compiled before it will be available for tab completion.
#       Brand new models/sources/tags/packages will not be displayed in the tab
#       complete menu until they are compiled and appear in the manifest.
#
#
# CREDITS
#   Leveraging a lot of logic from dbt-completion.bash:
#   https://github.com/fishtown-analytics/dbt-completion.bash/blob/master/dbt-completion.bash
#
#  Inspired by zsh-completions
#  https://github.com/zsh-users/zsh-completions
#  and particularly https://github.com/zsh-users/zsh-completions/blob/40c6c768eabfa49d54a149149f338e23ee6dd83b/src/_ufw



# Inline a python script so we can deploy this as a single file
# the idea of doing this in bash natively is... daunting
_parse_manifest() {
manifest_path=$1
prefix=$2
prog=$(cat <<EOF
# Use a big try/catch so any errors (maybe from a corrupted or
# missing manifest?) are not printed on tab-complete

try:
    import fileinput, json, sys

    # If a prefix is given as an argument, include it in the
    # generated selector list. The bash completion logic below
    # will match these generated selectors against partially
    # written args when table completed. This helps the script
    # match selectors when a user does something like:
    #   dbt run --models +order<tab>

    prefix = sys.argv.pop() if len(sys.argv) == 2 else ""

    manifest = json.loads("\n".join([line for line in fileinput.input()]))

    models = set(
        "{}{}".format(prefix, node['name'])
        for node in manifest['nodes'].values()
        if node['resource_type'] in ['model', 'seed']
    )

    tags = set(
        "{}tag:{}".format(prefix, tag)
        for node in manifest['nodes'].values()
        for tag in node.get('tags', [])
        if node['resource_type'] == 'model'
    )

    # The + prefix for sources is not sensible, but allowed.
    # This script shouldn't be opinionated about these things
    sources = set(
        "{}source:{}".format(prefix, node['source_name'])
        for node in manifest['sources'].values()
        if node['resource_type'] == 'source'
    ) | set(
        "{}source:{}.{}".format(prefix, node['source_name'], node['name'])
        for node in manifest['sources'].values()
        if node['resource_type'] == 'source'
    )

    exposures = set(
        "{}exposure:{}".format(prefix, node['name'])
        for node in manifest['exposures'].values()
        if node['resource_type'] == 'exposure'
    )

    metrics = set(
        "{}metric:{}".format(prefix, node['name'])
        for node in manifest['metrics'].values()
        if node['resource_type'] == 'metric'
    )

    # Generate partial Fully Qualified Names with a wildcard
    # suffix. This matches things like directories and packag names
    fqns = set(
        "{}{}.*".format(prefix, ".".join(node['fqn'][:i-1]))
        for node in manifest['nodes'].values()
        for i in range(len(node.get('fqn', [])))
        if node['resource_type'] == 'model'
    )

    hard_coded = {
        "{}config.materialized:view".format(prefix),
        "{}config.materialized:table".format(prefix),
        "{}config.materialized:incremental".format(prefix),
        "{}config.materialized:snapshot".format(prefix),
        "{}state:new".format(prefix),
        "{}state:modified".format(prefix),
        "{}result:error".format(prefix),
        "{}result:fail".format(prefix)
    }

    selectors = [
        selector.replace(':', r'\:')
        for selector in (models | tags | sources | exposures | metrics | fqns | hard_coded)
        if selector != ''
    ]

    print(" ".join(selectors))
except Exception as e:
    print(e)
    # oops!
    pass
EOF
)

cat "$manifest_path" | python3 -c "$prog" $prefix
}


_parse_selectors() {
manifest_path=$1
prog=$(cat <<EOF
# Use a big try/catch so any errors (maybe from a corrupted or
# missing manifest?) are not printed on tab-complete

try:
    import fileinput, json, sys

    # No prefix is allowed for YAML selectors

    manifest = json.loads("\n".join([line for line in fileinput.input()]))

    selectors = set(
        "{}".format(node['name'])
        for node in manifest['selectors'].values()
    )

    if selectors:
        print(" ".join(selectors))
    else:
        print("")
except Exception as e:
    print(e)
    # oops!
    pass
EOF
)

cat "$manifest_path" | python3 -c "$prog"
}


# Check if DBT_PROJECT_DIR is set and not empty
# Otherwise, walk up the filesystem until we find a dbt_project.yml file,
# then return the path which contains it (if found)
_get_project_root() {
  if [ -n "$DBT_PROJECT_DIR" ]; then
    echo "$DBT_PROJECT_DIR" && return
  fi

  slashes=${PWD//[^\/]/}
  directory="$PWD"
  for (( n=${#slashes}; n>0; --n ))
  do
    test -e "$directory/dbt_project.yml" && echo "$directory" && return
    directory="$directory/.."
  done
}


# Lists the difference models, extracted from manifest.json
_dbt_list_models() {

    # Wanted to look at $words to know when or not to suggest models,
    # but the issue is that $words seems to reset after the first argument has been typed


    # Option 1 to try to capture if the last flag is for a model (does not work because $words resets)
    # if [ $(($words[(I)--model] + $words[(I)-m])) != $words[(I)-*] ] || [ $words[(I)-*] = 0 ]
    # then
    #     return
    # fi

    # Option 2 to try to capture if the last flag is for a model  (does not work because $words resets)
    # last_flag=$(_get_last_flag ${#words} $words)
    # is_selector=$(_flag_is_selector $last_flag)
    #
    # if [[ $is_selector == 0 ]] ; then
    #     return
    # fi

    project_dir="$(_get_project_root)"

    # Attempt to fetch the manifest path from the environment variable
    if [ -z "$DBT_MANIFEST_PATH" ] ; then
        manifest_path="${project_dir}/target/manifest.json"
    else
        manifest_path="$DBT_MANIFEST_PATH"
    fi

    if [ ! -f "$manifest_path" ] ; then
        return
    fi

    local first_letter
    first_letter=${words[-1]:0:1}

    if [ "$first_letter" = "+" ] || [ "$first_letter" = "@" ]; then
        local models_list=( $(_parse_manifest "$manifest_path" "$first_letter") )
    else
        local models_list=( $(_parse_manifest "$manifest_path" "") )
    fi
    _values -s , 'models' $models_list
}


# Lists the different selectors, extracted from manifest.json
_dbt_list_selectors() {

    project_dir="$(_get_project_root)"

    # Attempt to fetch the manifest path from the environment variable
    if [ -z "$DBT_MANIFEST_PATH" ] ; then
        manifest_path="${project_dir}/target/manifest.json"
    else
        manifest_path="$DBT_MANIFEST_PATH"
    fi

    if [ ! -f "$manifest_path" ] ; then
        return
    fi

    local selectors_list=( $(_parse_selectors "$manifest_path") )
    if [ "$selectors_list" != "" ]; then
        _values -s , 'selectors' $selectors_list
    fi
}


# Detect whether the dbt binary is dbt-fusion or dbt-core.
# Result is cached (keyed by binary path + mtime) to avoid running --help on every tab press.
_dbt_detect_type() {
    local bin_path
    bin_path=$(command -v dbt) || return

    local cache_dir="${XDG_CACHE_HOME:-$HOME/.cache}"
    local type_cache="$cache_dir/dbt_binary_type"
    local key_cache="$cache_dir/dbt_binary_key"
    local current_key="${bin_path}:$(stat -f %m "$bin_path" 2>/dev/null)"

    if [[ -f "$type_cache" ]] && [[ "$(cat "$key_cache" 2>/dev/null)" == "$current_key" ]]; then
        cat "$type_cache"
        return
    fi

    local first_line
    first_line=$(dbt --help 2>/dev/null | head -1)

    local type="core"
    if [[ "$first_line" == *"dbt-fusion"* ]]; then
        type="fusion"
    fi

    mkdir -p "$cache_dir"
    printf '%s' "$type" > "$type_cache"
    printf '%s' "$current_key" > "$key_cache"
    printf '%s' "$type"
}


# For dbt-fusion: source the clap_complete-generated zsh script (cached by binary mtime),
# then delegate to the _dbt function it defines.
# Accepts an optional explicit binary path (used when completing the dbtf alias).
_dbt_fusion_complete() {
    local bin_path="${1:-$(command -v dbt)}"
    [[ -z "$bin_path" ]] && return 1

    local cache_dir="${XDG_CACHE_HOME:-$HOME/.cache}"
    local script_cache="$cache_dir/dbt_fusion_completions.zsh"
    local key_cache="$cache_dir/dbt_fusion_completions.key"
    local current_key="${bin_path}:$(stat -f %m "$bin_path" 2>/dev/null)"

    if [[ ! -f "$script_cache" ]] || [[ "$(cat "$key_cache" 2>/dev/null)" != "$current_key" ]]; then
        mkdir -p "$cache_dir"
        "$bin_path" completions zsh > "$script_cache" 2>/dev/null
        printf '%s' "$current_key" > "$key_cache"
    fi

    # Source the clap-generated script; it redefines _dbt with clap's completion logic.
    # Rename it to _dbt_clap so it doesn't clobber this function, then call it.
    local script_content
    script_content=$(sed 's/^_dbt()/_dbt_clap()/' "$script_cache")
    eval "$script_content"
    _dbt_clap "$@"
}


# For dbt-core: Click runtime completion.
# Note: spawns a Python process per tab press, so inherently slower than dbt-fusion.
_dbt_core_complete() {
    local ret=1
    local IFS=$'\n'
    local response

    # COMP_CWORD in Click is 0-indexed for the position we're completing
    response=$(env COMP_WORDS="${words[*]}" COMP_CWORD=$((CURRENT - 1)) _DBT_COMPLETE=zsh_complete dbt 2>/dev/null)

    if [[ -n "$response" ]]; then
        local lines=("${(@f)response}")
        local i=1
        local -a completions descriptions

        while (( i <= ${#lines[@]} )); do
            local comp_type="${lines[i]}"
            local comp_value="${lines[i+1]}"
            local comp_desc="${lines[i+2]}"

            if [[ $comp_type == 'dir' ]]; then
                _files -/ && ret=0
            elif [[ $comp_type == 'file' ]]; then
                _files && ret=0
            elif [[ $comp_type == 'plain' ]]; then
                if [[ -n "$comp_desc" ]]; then
                    descriptions+=("$comp_value:$comp_desc")
                else
                    completions+=("$comp_value")
                fi
                ret=0
            fi

            i=$((i + 3))
        done

        if (( ${#descriptions[@]} > 0 )); then
            _describe 'dbt commands' descriptions
        elif (( ${#completions[@]} > 0 )); then
            compadd -U -a completions
        fi
    fi

    return ret
}


# Main function
# Auto-detects dbt-core vs dbt-fusion and routes to the appropriate completion backend.
# For the dbtf alias (always Fusion), the binary path is resolved from the alias value.
# Manifest-based model/selector completions are applied for both.
_dbt() {
    local ret=1

    # Manifest-based model/selector completions apply regardless of binary type
    local prev_word=""
    if (( CURRENT > 1 )); then
        prev_word="${words[CURRENT-1]}"
    fi

    case "$prev_word" in
        -s|--select|-m|--model|--models|--exclude)
            _dbt_list_models
            return $?
            ;;
        --selector)
            _dbt_list_selectors
            return $?
            ;;
    esac

    # dbtf is always Fusion; resolve the binary from the alias value directly
    if [[ "$service" == "dbtf" ]]; then
        local dbtf_bin="${aliases[dbtf]}"
        _dbt_fusion_complete "$dbtf_bin"
        return $?
    fi

    if ! command -v dbt &> /dev/null; then
        return 1
    fi

    local dbt_type
    dbt_type=$(_dbt_detect_type)

    if [[ "$dbt_type" == "fusion" ]]; then
        _dbt_fusion_complete
    else
        _dbt_core_complete
    fi
    return $?
}

_dbt