relation.py

from collections.abc import Hashable
from dataclasses import dataclass
from typing import (
    Optional, TypeVar, Any, Type, Dict, Union, Iterator, Tuple, Set
)

from dbt.contracts.graph.compiled import CompiledNode
from dbt.contracts.graph.parsed import ParsedSourceDefinition, ParsedNode
from dbt.contracts.relation import (
    RelationType, ComponentName, HasQuoting, FakeAPIObject, Policy, Path
)
from dbt.exceptions import InternalException
from dbt.node_types import NodeType
from dbt.utils import filter_null_values, deep_merge, classproperty

import dbt.exceptions


Self = TypeVar('Self', bound='BaseRelation')


@dataclass(frozen=True, eq=False, repr=False)
class BaseRelation(FakeAPIObject, Hashable):
    path: Path
    type: Optional[RelationType] = None
    quote_character: str = '"'
    include_policy: Policy = Policy()
    quote_policy: Policy = Policy()
    dbt_created: bool = False

    def _is_exactish_match(self, field: ComponentName, value: str) -> bool:
        if self.dbt_created and self.quote_policy.get_part(field) is False:
            return self.path.get_lowered_part(field) == value.lower()
        else:
            return self.path.get_part(field) == value

    @classmethod
    def _get_field_named(cls, field_name):
        for field, _ in cls._get_fields():
            if field.name == field_name:
                return field
        # this should be unreachable
        raise ValueError(f'BaseRelation has no {field_name} field!')

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False
        return self.to_dict(omit_none=True) == other.to_dict(omit_none=True)

    @classmethod
    def get_default_quote_policy(cls) -> Policy:
        return cls._get_field_named('quote_policy').default

    @classmethod
    def get_default_include_policy(cls) -> Policy:
        return cls._get_field_named('include_policy').default

    def get(self, key, default=None):
        """Override `.get` to return a metadata object so we don't break
        dbt_utils.
        """
        if key == 'metadata':
            return {
                'type': self.__class__.__name__
            }
        return super().get(key, default)

    def matches(
        self,
        database: Optional[str] = None,
        schema: Optional[str] = None,
        identifier: Optional[str] = None,
    ) -> bool:
        search = filter_null_values({
            ComponentName.Database: database,
            ComponentName.Schema: schema,
            ComponentName.Identifier: identifier
        })

        if not search:
            # nothing was passed in
            raise dbt.exceptions.RuntimeException(
                "Tried to match relation, but no search path was passed!")

        exact_match = True
        approximate_match = True

        for k, v in search.items():
            if not self._is_exactish_match(k, v):
                exact_match = False

            if self.path.get_lowered_part(k) != v.lower():
                approximate_match = False

        if approximate_match and not exact_match:
            target = self.create(
                database=database, schema=schema, identifier=identifier
            )
            dbt.exceptions.approximate_relation_match(target, self)

        return exact_match

    def replace_path(self, **kwargs):
        return self.replace(path=self.path.replace(**kwargs))

    def quote(
        self: Self,
        database: Optional[bool] = None,
        schema: Optional[bool] = None,
        identifier: Optional[bool] = None,
    ) -> Self:
        policy = filter_null_values({
            ComponentName.Database: database,
            ComponentName.Schema: schema,
            ComponentName.Identifier: identifier
        })

        new_quote_policy = self.quote_policy.replace_dict(policy)
        return self.replace(quote_policy=new_quote_policy)

    def include(
        self: Self,
        database: Optional[bool] = None,
        schema: Optional[bool] = None,
        identifier: Optional[bool] = None,
    ) -> Self:
        policy = filter_null_values({
            ComponentName.Database: database,
            ComponentName.Schema: schema,
            ComponentName.Identifier: identifier
        })

        new_include_policy = self.include_policy.replace_dict(policy)
        return self.replace(include_policy=new_include_policy)

    def information_schema(self, view_name=None) -> 'InformationSchema':
        # some of our data comes from jinja, where things can be `Undefined`.
        if not isinstance(view_name, str):
            view_name = None

        # Kick the user-supplied schema out of the information schema relation
        # Instead address this as <database>.information_schema by default
        info_schema = InformationSchema.from_relation(self, view_name)
        return info_schema.incorporate(path={"schema": None})

    def information_schema_only(self) -> 'InformationSchema':
        return self.information_schema()

    def without_identifier(self) -> 'BaseRelation':
        """Return a form of this relation that only has the database and schema
        set to included. To get the appropriately-quoted form the schema out of
        the result (for use as part of a query), use `.render()`. To get the
        raw database or schema name, use `.database` or `.schema`.

        The hash of the returned object is the result of render().
        """
        return self.include(identifier=False).replace_path(identifier=None)

    def _render_iterator(
        self
    ) -> Iterator[Tuple[Optional[ComponentName], Optional[str]]]:

        for key in ComponentName:
            path_part: Optional[str] = None
            if self.include_policy.get_part(key):
                path_part = self.path.get_part(key)
                if path_part is not None and self.quote_policy.get_part(key):
                    path_part = self.quoted(path_part)
            yield key, path_part

    def render(self) -> str:
        # if there is nothing set, this will return the empty string.
        return '.'.join(
            part for _, part in self._render_iterator()
            if part is not None
        )

    def quoted(self, identifier):
        return '{quote_char}{identifier}{quote_char}'.format(
            quote_char=self.quote_character,
            identifier=identifier,
        )

    @classmethod
    def create_from_source(
        cls: Type[Self], source: ParsedSourceDefinition, **kwargs: Any
    ) -> Self:
        source_quoting = source.quoting.to_dict(omit_none=True)
        source_quoting.pop('column', None)
        quote_policy = deep_merge(
            cls.get_default_quote_policy().to_dict(omit_none=True),
            source_quoting,
            kwargs.get('quote_policy', {}),
        )

        return cls.create(
            database=source.database,
            schema=source.schema,
            identifier=source.identifier,
            quote_policy=quote_policy,
            **kwargs
        )

    @staticmethod
    def add_ephemeral_prefix(name: str):
        return f'__dbt__cte__{name}'

    @classmethod
    def create_ephemeral_from_node(
        cls: Type[Self],
        config: HasQuoting,
        node: Union[ParsedNode, CompiledNode],
    ) -> Self:
        # Note that ephemeral models are based on the name.
        identifier = cls.add_ephemeral_prefix(node.name)
        return cls.create(
            type=cls.CTE,
            identifier=identifier,
        ).quote(identifier=False)

    @classmethod
    def create_from_node(
        cls: Type[Self],
        config: HasQuoting,
        node: Union[ParsedNode, CompiledNode],
        quote_policy: Optional[Dict[str, bool]] = None,
        **kwargs: Any,
    ) -> Self:
        if quote_policy is None:
            quote_policy = {}

        quote_policy = dbt.utils.merge(config.quoting, quote_policy)

        return cls.create(
            database=node.database,
            schema=node.schema,
            identifier=node.alias,
            quote_policy=quote_policy,
            **kwargs)

    @classmethod
    def create_from(
        cls: Type[Self],
        config: HasQuoting,
        node: Union[CompiledNode, ParsedNode, ParsedSourceDefinition],
        **kwargs: Any,
    ) -> Self:
        if node.resource_type == NodeType.Source:
            if not isinstance(node, ParsedSourceDefinition):
                raise InternalException(
                    'type mismatch, expected ParsedSourceDefinition but got {}'
                    .format(type(node))
                )
            return cls.create_from_source(node, **kwargs)
        else:
            if not isinstance(node, (ParsedNode, CompiledNode)):
                raise InternalException(
                    'type mismatch, expected ParsedNode or CompiledNode but '
                    'got {}'.format(type(node))
                )
            return cls.create_from_node(config, node, **kwargs)

    @classmethod
    def create(
        cls: Type[Self],
        database: Optional[str] = None,
        schema: Optional[str] = None,
        identifier: Optional[str] = None,
        type: Optional[RelationType] = None,
        **kwargs,
    ) -> Self:
        kwargs.update({
            'path': {
                'database': database,
                'schema': schema,
                'identifier': identifier,
            },
            'type': type,
        })
        return cls.from_dict(kwargs)

    def __repr__(self) -> str:
        return "<{} {}>".format(self.__class__.__name__, self.render())

    def __hash__(self) -> int:
        return hash(self.render())

    def __str__(self) -> str:
        return self.render()

    @property
    def database(self) -> Optional[str]:
        return self.path.database

    @property
    def schema(self) -> Optional[str]:
        return self.path.schema

    @property
    def identifier(self) -> Optional[str]:
        return self.path.identifier

    @property
    def table(self) -> Optional[str]:
        return self.path.identifier

    # Here for compatibility with old Relation interface
    @property
    def name(self) -> Optional[str]:
        return self.identifier

    @property
    def is_table(self) -> bool:
        return self.type == RelationType.Table

    @property
    def is_cte(self) -> bool:
        return self.type == RelationType.CTE

    @property
    def is_view(self) -> bool:
        return self.type == RelationType.View

    @classproperty
    def Table(cls) -> str:
        return str(RelationType.Table)

    @classproperty
    def CTE(cls) -> str:
        return str(RelationType.CTE)

    @classproperty
    def View(cls) -> str:
        return str(RelationType.View)

    @classproperty
    def External(cls) -> str:
        return str(RelationType.External)

    @classproperty
    def get_relation_type(cls) -> Type[RelationType]:
        return RelationType


Info = TypeVar('Info', bound='InformationSchema')


@dataclass(frozen=True, eq=False, repr=False)
class InformationSchema(BaseRelation):
    information_schema_view: Optional[str] = None

    def __post_init__(self):
        if not isinstance(self.information_schema_view, (type(None), str)):
            raise dbt.exceptions.CompilationException(
                'Got an invalid name: {}'.format(self.information_schema_view)
            )

    @classmethod
    def get_path(
        cls, relation: BaseRelation, information_schema_view: Optional[str]
    ) -> Path:
        return Path(
            database=relation.database,
            schema=relation.schema,
            identifier='INFORMATION_SCHEMA',
        )

    @classmethod
    def get_include_policy(
        cls,
        relation,
        information_schema_view: Optional[str],
    ) -> Policy:
        return relation.include_policy.replace(
            database=relation.database is not None,
            schema=False,
            identifier=True,
        )

    @classmethod
    def get_quote_policy(
        cls,
        relation,
        information_schema_view: Optional[str],
    ) -> Policy:
        return relation.quote_policy.replace(
            identifier=False,
        )

    @classmethod
    def from_relation(
        cls: Type[Info],
        relation: BaseRelation,
        information_schema_view: Optional[str],
    ) -> Info:
        include_policy = cls.get_include_policy(
            relation, information_schema_view
        )
        quote_policy = cls.get_quote_policy(relation, information_schema_view)
        path = cls.get_path(relation, information_schema_view)
        return cls(
            type=RelationType.View,
            path=path,
            include_policy=include_policy,
            quote_policy=quote_policy,
            information_schema_view=information_schema_view,
        )

    def _render_iterator(self):
        for k, v in super()._render_iterator():
            yield k, v
        yield None, self.information_schema_view


class SchemaSearchMap(Dict[InformationSchema, Set[Optional[str]]]):
    """A utility class to keep track of what information_schema tables to
    search for what schemas. The schema values are all lowercased to avoid
    duplication.
    """
    def add(self, relation: BaseRelation):
        key = relation.information_schema_only()
        if key not in self:
            self[key] = set()
        schema: Optional[str] = None
        if relation.schema is not None:
            schema = relation.schema.lower()
        self[key].add(schema)

    def search(
        self
    ) -> Iterator[Tuple[InformationSchema, Optional[str]]]:
        for information_schema_name, schemas in self.items():
            for schema in schemas:
                yield information_schema_name, schema

    def flatten(self, allow_multiple_databases: bool = False):
        new = self.__class__()

        # make sure we don't have multiple databases if allow_multiple_databases is set to False
        if not allow_multiple_databases:
            seen = {r.database.lower() for r in self if r.database}
            if len(seen) > 1:
                dbt.exceptions.raise_compiler_error(str(seen))

        for information_schema_name, schema in self.search():
            path = {
                'database': information_schema_name.database,
                'schema': schema
            }
            new.add(information_schema_name.incorporate(
                path=path,
                quote_policy={'database': False},
                include_policy={'database': False},
            ))

        return new