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

# import these so we can find them
from . import sql_commands  # noqa
from . import project_commands  # noqa
from . import deps  # noqa
import multiprocessing.queues  # noqa - https://bugs.python.org/issue41567
import json
import os
import signal
from contextlib import contextmanager
from typing import Iterator, Optional, List, Type

from werkzeug.middleware.dispatcher import DispatcherMiddleware
from werkzeug.wrappers import Request, Response
from werkzeug.serving import run_simple
from werkzeug.exceptions import NotFound

from dbt.exceptions import RuntimeException
from dbt.logger import (
    GLOBAL_LOGGER as logger,
    log_manager,
)
from dbt.rpc.logger import ServerContext, HTTPRequest, RPCResponse
from dbt.rpc.method import TaskTypes, RemoteMethod
from dbt.rpc.response_manager import ResponseManager
from dbt.rpc.task_manager import TaskManager
from dbt.task.base import ConfiguredTask
from dbt.utils import ForgivingJSONEncoder


# SIG_DFL ends up killing the process if multiple build up, but SIG_IGN just
# peacefully carries on
SIG_IGN = signal.SIG_IGN


@contextmanager
def signhup_replace() -> Iterator[bool]:
    """A context manager. Replace the current sighup handler with SIG_IGN on
    entering, and (if the current handler was not SIG_IGN) replace it on
    leaving. This is meant to be used inside a sighup handler itself to
    provide. a sort of locking model.

    This relies on the fact that 1) signals are only handled by the main thread
    (the default in Python) and 2) signal.signal() is "atomic" (only C
    instructions). I'm pretty sure that's reliable on posix.

    This shouldn't replace if the handler wasn't already SIG_IGN, and should
    yield whether it has the lock as its value. Callers shouldn't do
    singal-handling things inside this context manager if it does not have the
    lock (they should just exit the context).
    """
    # Don't use locks here! This is called from inside a signal handler

    # set our handler to ignore signals, capturing the existing one
    current_handler = signal.signal(signal.SIGHUP, SIG_IGN)

    # current_handler should be the handler unless we're already loading a
    # new manifest. So if the current handler is the ignore, there was a
    # double-hup! We should exit and not touch the signal handler, to make
    # sure we let the other signal handler fix it
    is_current_handler = current_handler is not SIG_IGN

    # if we got here, we're the ones in charge of configuring! Yield.
    try:
        yield is_current_handler
    finally:
        if is_current_handler:
            # the signal handler that successfully changed the handler is
            # responsible for resetting, and can't be re-called until it's
            # fixed, so no locking needed

            signal.signal(signal.SIGHUP, current_handler)


class RPCServerTask(ConfiguredTask):
    DEFAULT_LOG_FORMAT = 'json'

    def __init__(
        self, args, config, tasks: Optional[List[Type[RemoteMethod]]] = None
    ) -> None:
        if os.name == 'nt':
            raise RuntimeException(
                'The dbt RPC server is not supported on windows'
            )
        super().__init__(args, config)
        self.task_manager = TaskManager(
            self.args, self.config, TaskTypes(tasks)
        )
        signal.signal(signal.SIGHUP, self._sighup_handler)

    @classmethod
    def pre_init_hook(cls, args):
        """A hook called before the task is initialized."""
        if args.log_format == 'text':
            log_manager.format_text()
        else:
            log_manager.format_json()

    def _sighup_handler(self, signum, frame):
        with signhup_replace() as run_task_manger:
            if not run_task_manger:
                # a sighup handler is already active.
                return
            self.task_manager.reload_config()
            self.task_manager.reload_manifest()

    def run_forever(self):
        host = self.args.host
        port = self.args.port
        addr = (host, port)

        display_host = host
        if host == '0.0.0.0':
            display_host = 'localhost'

        logger.info(
            'Serving RPC server at {}:{}, pid={}'.format(
                *addr, os.getpid()
            )
        )

        logger.info(
            'Supported methods: {}'.format(sorted(self.task_manager.methods()))
        )

        logger.info(
            'Send requests to http://{}:{}/jsonrpc'.format(display_host, port)
        )

        app = DispatcherMiddleware(self.handle_request, {
            '/jsonrpc': self.handle_jsonrpc_request,
        })

        # we have to run in threaded mode if we want to share subprocess
        # handles, which is the easiest way to implement `kill` (it makes
        # `ps` easier as well). The alternative involves tracking
        # metadata+state in a multiprocessing.Manager, adds polling the
        # manager to the request  task handler and in general gets messy
        # fast.
        run_simple(
            host,
            port,
            app,
            threaded=not self.task_manager.single_threaded(),
        )

    def run(self):
        with ServerContext().applicationbound():
            self.run_forever()

    @Request.application
    def handle_jsonrpc_request(self, request):
        with HTTPRequest(request):
            jsonrpc_response = ResponseManager.handle(
                request, self.task_manager
            )
            json_data = json.dumps(
                jsonrpc_response.data,
                cls=ForgivingJSONEncoder,
            )
            response = Response(json_data, mimetype='application/json')
            # this looks and feels dumb, but our json encoder converts decimals
            # and datetimes, and if we use the json_data itself the output
            # looks silly because of escapes, so re-serialize it into valid
            # JSON types for logging.
            with RPCResponse(jsonrpc_response):
                logger.info('sending response ({}) to {}'.format(
                    response, request.remote_addr)
                )
            return response

    @Request.application
    def handle_request(self, request):
        raise NotFound()