dbt-selly/dbt-env/lib/python3.8/site-packages/dbt/task/list.py

import json

from dbt.contracts.graph.parsed import (
    ParsedExposure,
    ParsedSourceDefinition
)
from dbt.graph import ResourceTypeSelector
from dbt.task.runnable import GraphRunnableTask, ManifestTask
from dbt.task.test import TestSelector
from dbt.node_types import NodeType
from dbt.exceptions import RuntimeException, InternalException
from dbt.logger import log_manager, GLOBAL_LOGGER as logger


class ListTask(GraphRunnableTask):
    DEFAULT_RESOURCE_VALUES = frozenset((
        NodeType.Model,
        NodeType.Snapshot,
        NodeType.Seed,
        NodeType.Test,
        NodeType.Source,
        NodeType.Exposure,
    ))
    ALL_RESOURCE_VALUES = DEFAULT_RESOURCE_VALUES | frozenset((
        NodeType.Analysis,
    ))
    ALLOWED_KEYS = frozenset((
        'alias',
        'name',
        'package_name',
        'depends_on',
        'tags',
        'config',
        'resource_type',
        'source_name',
        'original_file_path',
        'unique_id'
    ))

    def __init__(self, args, config):
        super().__init__(args, config)
        if self.args.models:
            if self.args.select:
                raise RuntimeException(
                    '"models" and "select" are mutually exclusive arguments'
                )
            if self.args.resource_types:
                raise RuntimeException(
                    '"models" and "resource_type" are mutually exclusive '
                    'arguments'
                )

    @classmethod
    def pre_init_hook(cls, args):
        """A hook called before the task is initialized."""
        log_manager.stderr_console()
        super().pre_init_hook(args)

    def _iterate_selected_nodes(self):
        selector = self.get_node_selector()
        spec = self.get_selection_spec()
        nodes = sorted(selector.get_selected(spec))
        if not nodes:
            logger.warning('No nodes selected!')
            return
        if self.manifest is None:
            raise InternalException(
                'manifest is None in _iterate_selected_nodes'
            )
        for node in nodes:
            if node in self.manifest.nodes:
                yield self.manifest.nodes[node]
            elif node in self.manifest.sources:
                yield self.manifest.sources[node]
            elif node in self.manifest.exposures:
                yield self.manifest.exposures[node]
            else:
                raise RuntimeException(
                    f'Got an unexpected result from node selection: "{node}"'
                    f'Expected a source or a node!'
                )

    def generate_selectors(self):
        for node in self._iterate_selected_nodes():
            if node.resource_type == NodeType.Source:
                assert isinstance(node, ParsedSourceDefinition)
                # sources are searched for by pkg.source_name.table_name
                source_selector = '.'.join([
                    node.package_name, node.source_name, node.name
                ])
                yield f'source:{source_selector}'
            elif node.resource_type == NodeType.Exposure:
                assert isinstance(node, ParsedExposure)
                # exposures are searched for by pkg.exposure_name
                exposure_selector = '.'.join([node.package_name, node.name])
                yield f'exposure:{exposure_selector}'
            else:
                # everything else is from `fqn`
                yield '.'.join(node.fqn)

    def generate_names(self):
        for node in self._iterate_selected_nodes():
            yield node.search_name

    def generate_json(self):
        for node in self._iterate_selected_nodes():
            yield json.dumps({
                k: v
                for k, v in node.to_dict(omit_none=False).items()
                if (
                    k in self.args.output_keys
                    if self.args.output_keys is not None
                    else k in self.ALLOWED_KEYS
                )
            })

    def generate_paths(self):
        for node in self._iterate_selected_nodes():
            yield node.original_file_path

    def run(self):
        ManifestTask._runtime_initialize(self)
        output = self.args.output
        if output == 'selector':
            generator = self.generate_selectors
        elif output == 'name':
            generator = self.generate_names
        elif output == 'json':
            generator = self.generate_json
        elif output == 'path':
            generator = self.generate_paths
        else:
            raise InternalException(
                'Invalid output {}'.format(output)
            )

        return self.output_results(generator())

    def output_results(self, results):
        for result in results:
            self.node_results.append(result)
            print(result)
        return self.node_results

    @property
    def resource_types(self):
        if self.args.models:
            return [NodeType.Model]

        if not self.args.resource_types:
            return list(self.DEFAULT_RESOURCE_VALUES)

        values = set(self.args.resource_types)
        if 'default' in values:
            values.remove('default')
            values.update(self.DEFAULT_RESOURCE_VALUES)
        if 'all' in values:
            values.remove('all')
            values.update(self.ALL_RESOURCE_VALUES)
        return list(values)

    @property
    def selection_arg(self):
        # for backwards compatibility, list accepts both --models and --select,
        # with slightly different behavior: --models implies --resource-type model
        if self.args.models:
            return self.args.models
        else:
            return self.args.select

    def get_node_selector(self):
        if self.manifest is None or self.graph is None:
            raise InternalException(
                'manifest and graph must be set to get perform node selection'
            )
        if self.resource_types == [NodeType.Test]:
            return TestSelector(
                graph=self.graph,
                manifest=self.manifest,
                previous_state=self.previous_state,
            )
        else:
            return ResourceTypeSelector(
                graph=self.graph,
                manifest=self.manifest,
                previous_state=self.previous_state,
                resource_types=self.resource_types,
            )

    def interpret_results(self, results):
        # list command should always return 0 as exit code
        return True