Databases¶

class Query(object):
    """
    This class accepts a single `dict` argument which is necessary to run the query. The representation
    is based on our `custom_queries` format originally designed and implemented in !1528.

    It is now part of all our database integrations and
    [other](https://cloud.google.com/solutions/sap/docs/sap-hana-monitoring-agent-planning-guide#defining_custom_queries)
    products have since adopted this format.
    """

    def __init__(self, query_data):
        '''
        Parameters:
            query_data (Dict[str, Any]): The query data to run the query. It should contain the following fields:
                - name (str): The name of the query.
                - query (str): The query to run.
                - columns (List[Dict[str, Any]]): Each column should contain the following fields:
                    - name (str): The name of the column.
                    - type (str): The type of the column.
                    - (Optional) Any other field that the column transformer for the type requires.
                - (Optional) extras (List[Dict[str, Any]]): Each extra transformer should contain the following fields:
                    - name (str): The name of the extra transformer.
                    - type (str): The type of the extra transformer.
                    - (Optional) Any other field that the extra transformer for the type requires.
                - (Optional) tags (List[str]): The tags to add to the query result.
                - (Optional) collection_interval (int): The collection interval (in seconds) of the query.
                    Note:
                        If collection_interval is None, the query will be run every check run.
                        If the collection interval is less than check collection interval,
                        the query will be run every check run.
                        If the collection interval is greater than check collection interval,
                        the query will NOT BE RUN exactly at the collection interval.
                        The query will be run at the next check run after the collection interval has passed.
                - (Optional) metric_prefix (str): The prefix to add to the metric name.
                    Note: If the metric prefix is None, the default metric prefix `<INTEGRATION>.` will be used.
        '''
        # Contains the data to fill the rest of the attributes
        self.query_data = deepcopy(query_data or {})  # type: Dict[str, Any]
        self.name = None  # type: str
        # The actual query
        self.query = None  # type: str
        # Contains a mapping of column_name -> column_type, transformer
        self.column_transformers = None  # type: Tuple[Tuple[str, Tuple[str, Transformer]]]
        # These transformers are used to collect extra metrics calculated from the query result
        self.extra_transformers = None  # type: List[Tuple[str, Transformer]]
        # Contains the tags defined in query_data, more tags can be added later from the query result
        self.base_tags = None  # type: List[str]
        # The collecton interval (in seconds) of the query. If None, the query will be run every check run.
        self.collection_interval = None  # type: int
        # The last time the query was executed. If None, the query has never been executed.
        # This is only used when the collection_interval is not None.
        self.__last_execution_time = None  # type: float
        # whether to ignore any defined namespace prefix. True when `metric_prefix` is defined.
        self.metric_name_raw = False  # type: bool

    def compile(
        self,
        column_transformers,  # type: Dict[str, TransformerFactory]
        extra_transformers,  # type: Dict[str, TransformerFactory]
    ):
        # type: (...) -> None

        """
        This idempotent method will be called by `QueryManager.compile_queries` so you
        should never need to call it directly.
        """
        # Check for previous compilation
        if self.name is not None:
            return

        query_name = self.query_data.get('name')
        if not query_name:
            raise ValueError('query field `name` is required')
        elif not isinstance(query_name, str):
            raise ValueError('query field `name` must be a string')

        metric_prefix = self.query_data.get('metric_prefix')
        if metric_prefix is not None:
            if not isinstance(metric_prefix, str):
                raise ValueError('field `metric_prefix` for {} must be a string'.format(query_name))
            elif not metric_prefix:
                raise ValueError('field `metric_prefix` for {} must not be empty'.format(query_name))

        query = self.query_data.get('query')
        if not query:
            raise ValueError('field `query` for {} is required'.format(query_name))
        elif query_name.startswith('custom query #') and not isinstance(query, str):
            raise ValueError('field `query` for {} must be a string'.format(query_name))

        columns = self.query_data.get('columns')
        if not columns:
            raise ValueError('field `columns` for {} is required'.format(query_name))
        elif not isinstance(columns, list):
            raise ValueError('field `columns` for {} must be a list'.format(query_name))

        tags = self.query_data.get('tags', [])
        if tags is not None and not isinstance(tags, list):
            raise ValueError('field `tags` for {} must be a list'.format(query_name))

        # Keep track of all defined names
        sources = {}

        column_data = []
        for i, column in enumerate(columns, 1):
            # Columns can be ignored via configuration.
            if not column:
                column_data.append((None, None))
                continue
            elif not isinstance(column, dict):
                raise ValueError('column #{} of {} is not a mapping'.format(i, query_name))

            column_name = column.get('name')
            if not column_name:
                raise ValueError('field `name` for column #{} of {} is required'.format(i, query_name))
            elif not isinstance(column_name, str):
                raise ValueError('field `name` for column #{} of {} must be a string'.format(i, query_name))
            elif column_name in sources:
                raise ValueError(
                    'the name {} of {} was already defined in {} #{}'.format(
                        column_name, query_name, sources[column_name]['type'], sources[column_name]['index']
                    )
                )

            sources[column_name] = {'type': 'column', 'index': i}

            column_type = column.get('type')
            if not column_type:
                raise ValueError('field `type` for column {} of {} is required'.format(column_name, query_name))
            elif not isinstance(column_type, str):
                raise ValueError('field `type` for column {} of {} must be a string'.format(column_name, query_name))
            elif column_type == 'source':
                column_data.append((column_name, (None, None)))
                continue
            elif column_type not in column_transformers:
                raise ValueError('unknown type `{}` for column {} of {}'.format(column_type, column_name, query_name))

            __column_type_is_tag = column_type in ('tag', 'tag_list', 'tag_not_null')
            modifiers = {key: value for key, value in column.items() if key not in ('name', 'type')}

            try:
                if not __column_type_is_tag and metric_prefix:
                    # if metric_prefix is defined, we prepend it to the column name
                    column_name = "{}.{}".format(metric_prefix, column_name)
                transformer = column_transformers[column_type](column_transformers, column_name, **modifiers)
            except Exception as e:
                error = 'error compiling type `{}` for column {} of {}: {}'.format(
                    column_type, column_name, query_name, e
                )

                # Prepend helpful error text.
                #
                # When an exception is raised in the context of another one, both will be printed. To avoid
                # this we set the context to None. https://www.python.org/dev/peps/pep-0409/
                raise type(e)(error) from None
            else:
                if __column_type_is_tag:
                    column_data.append((column_name, (column_type, transformer)))
                else:
                    # All these would actually submit data. As that is the default case, we represent it as
                    # a reference to None since if we use e.g. `value` it would never be checked anyway.
                    column_data.append((column_name, (None, transformer)))

        submission_transformers = column_transformers.copy()  # type: Dict[str, Transformer]
        submission_transformers.pop('tag')
        submission_transformers.pop('tag_list')
        submission_transformers.pop('tag_not_null')

        extras = self.query_data.get('extras', [])  # type: List[Dict[str, Any]]
        if not isinstance(extras, list):
            raise ValueError('field `extras` for {} must be a list'.format(query_name))

        extra_data = []  # type: List[Tuple[str, Transformer]]
        for i, extra in enumerate(extras, 1):
            if not isinstance(extra, dict):
                raise ValueError('extra #{} of {} is not a mapping'.format(i, query_name))

            extra_type = extra.get('type')  # type: str
            extra_name = extra.get('name')  # type: str
            if extra_type == 'log':
                # The name is unused
                extra_name = 'log'
            elif not extra_name:
                raise ValueError('field `name` for extra #{} of {} is required'.format(i, query_name))
            elif not isinstance(extra_name, str):
                raise ValueError('field `name` for extra #{} of {} must be a string'.format(i, query_name))
            elif extra_name in sources:
                raise ValueError(
                    'the name {} of {} was already defined in {} #{}'.format(
                        extra_name, query_name, sources[extra_name]['type'], sources[extra_name]['index']
                    )
                )

            sources[extra_name] = {'type': 'extra', 'index': i}

            if not extra_type:
                if 'expression' in extra:
                    extra_type = 'expression'
                else:
                    raise ValueError('field `type` for extra {} of {} is required'.format(extra_name, query_name))
            elif not isinstance(extra_type, str):
                raise ValueError('field `type` for extra {} of {} must be a string'.format(extra_name, query_name))
            elif extra_type not in extra_transformers and extra_type not in submission_transformers:
                raise ValueError('unknown type `{}` for extra {} of {}'.format(extra_type, extra_name, query_name))

            transformer_factory = extra_transformers.get(
                extra_type, submission_transformers.get(extra_type)
            )  # type: TransformerFactory

            extra_source = extra.get('source')
            if extra_type in submission_transformers:
                if not extra_source:
                    raise ValueError('field `source` for extra {} of {} is required'.format(extra_name, query_name))

                modifiers = {key: value for key, value in extra.items() if key not in ('name', 'type', 'source')}
            else:
                modifiers = {key: value for key, value in extra.items() if key not in ('name', 'type')}
                modifiers['sources'] = sources

            try:
                transformer = transformer_factory(submission_transformers, extra_name, **modifiers)
            except Exception as e:
                error = 'error compiling type `{}` for extra {} of {}: {}'.format(extra_type, extra_name, query_name, e)

                raise type(e)(error) from None
            else:
                if extra_type in submission_transformers:
                    transformer = create_extra_transformer(transformer, extra_source)

                extra_data.append((extra_name, transformer))

        collection_interval = self.query_data.get('collection_interval')
        if collection_interval is not None:
            if not isinstance(collection_interval, (int, float)):
                raise ValueError('field `collection_interval` for {} must be a number'.format(query_name))
            elif int(collection_interval) <= 0:
                raise ValueError(
                    'field `collection_interval` for {} must be a positive number after rounding'.format(query_name)
                )
            collection_interval = int(collection_interval)

        self.name = query_name
        self.query = query
        self.column_transformers = tuple(column_data)
        self.extra_transformers = tuple(extra_data)
        self.base_tags = tags
        self.collection_interval = collection_interval
        self.metric_name_raw = metric_prefix is not None
        del self.query_data

    def should_execute(self):
        '''
        Check if the query should be executed based on the collection interval.

        :return: True if the query should be executed, False otherwise.
        '''
        if self.collection_interval is None:
            # if the collection interval is None, the query should always be executed.
            return True

        now = get_timestamp()
        if self.__last_execution_time is None or now - self.__last_execution_time >= self.collection_interval:
            # if the last execution time is None (the query has never been executed),
            # if the time since the last execution is greater than or equal to the collection interval,
            # the query should be executed.
            self.__last_execution_time = now
            return True

        return False

class QueryManager(QueryExecutor):
    """
    This class is in charge of running any number of `Query` instances for a single Check instance.

    You will most often see it created during Check initialization like this:

    ```python
    self._query_manager = QueryManager(
        self,
        self.execute_query,
        queries=[
            queries.SomeQuery1,
            queries.SomeQuery2,
            queries.SomeQuery3,
            queries.SomeQuery4,
            queries.SomeQuery5,
        ],
        tags=self.instance.get('tags', []),
        error_handler=self._error_sanitizer,
    )
    self.check_initializations.append(self._query_manager.compile_queries)
    ```

    Note: This class is not in charge of opening or closing connections, just running queries.
    """

    def __init__(
        self,
        check,  # type: AgentCheck
        executor,  # type:  QueriesExecutor
        queries=None,  # type: List[Dict[str, Any]]
        tags=None,  # type: List[str]
        error_handler=None,  # type: Callable[[str], str]
        hostname=None,  # type: str
    ):  # type: (...) -> QueryManager
        """
        - **check** (_AgentCheck_) - an instance of a Check
        - **executor** (_callable_) - a callable accepting a `str` query as its sole argument and returning
          a sequence representing either the full result set or an iterator over the result set
        - **queries** (_List[Dict]_) - a list of queries in dict format
        - **tags** (_List[str]_) - a list of tags to associate with every submission
        - **error_handler** (_callable_) - a callable accepting a `str` error as its sole argument and returning
          a sanitized string, useful for scrubbing potentially sensitive information libraries emit
        """
        super(QueryManager, self).__init__(
            executor=executor,
            submitter=check,
            queries=queries,
            tags=tags,
            error_handler=error_handler,
            hostname=hostname,
            logger=check.log,
        )
        self.check = check  # type: AgentCheck

        only_custom_queries = is_affirmative(self.check.instance.get('only_custom_queries', False))  # type: bool
        custom_queries = list(self.check.instance.get('custom_queries', []))  # type: List[str]
        use_global_custom_queries = self.check.instance.get('use_global_custom_queries', True)  # type: str

        # Handle overrides
        if use_global_custom_queries == 'extend':
            custom_queries.extend(self.check.init_config.get('global_custom_queries', []))
        elif (
            not custom_queries
            and 'global_custom_queries' in self.check.init_config
            and is_affirmative(use_global_custom_queries)
        ):
            custom_queries = self.check.init_config.get('global_custom_queries', [])

        # Override statement queries if only running custom queries
        if only_custom_queries:
            self.queries = []

        # Deduplicate
        for i, custom_query in enumerate(iter_unique(custom_queries), 1):
            query = Query(custom_query)
            query.query_data.setdefault('name', 'custom query #{}'.format(i))
            self.queries.append(query)

        if len(self.queries) == 0:
            self.logger.warning('QueryManager initialized with no query')

    def execute(self, extra_tags=None):
        # This needs to stay here b/c when we construct a QueryManager in a check's __init__
        # there is no check ID at that point
        self.logger = self.check.log

        return super(QueryManager, self).execute(extra_tags)