Example #1
0
    def __init__(self, report):
        """Initialize.

        Parameters
        ----------
        report : Report
            The report with the script to run.

        """
        self._report = report
        self._event_queue = ScriptEventQueue()
        self._state = ScriptState.STOPPED
        self._last_run_data = RerunData(argv=report.argv,
                                        widget_state=WidgetStates())
        self._widgets = Widgets()

        self.run_on_save = config.get_option('server.runOnSave')

        self.on_state_changed = Signal(
            doc="""Emitted when the script's execution state state changes.

            Parameters
            ----------
            state : ScriptState
            """)

        self.on_file_change_not_handled = Signal(
            doc="Emitted when the file is modified and we haven't handled it.")

        self.on_script_compile_error = Signal(
            doc="""Emitted if our script fails to compile.  (*Not* emitted
            for normal exceptions thrown while a script is running.)

            Parameters
            ----------
            exception : Exception
                The exception that was thrown
            """)

        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed)

        # Will be set to true when we process a SHUTDOWN event
        self._shutdown_requested = False

        self._script_thread = None
        self._ctx = None

        # Set to true while we're executing. Used by
        # maybe_handle_execution_control_request.
        self._execing = False
Example #2
0
    def __init__(
        self, is_preheat, ioloop, script_path, command_line, uploaded_file_manager
    ):
        """Initialize the ReportSession.

        Parameters
        ----------
        is_preheat : bool
            True if this is the "preheated" session.

        ioloop : tornado.ioloop.IOLoop
            The Tornado IOLoop that we're running within.

        script_path : str
            Path of the Python file from which this report is generated.

        command_line : str
            Command line as input by the user.

        uploaded_file_manager : UploadedFileManager
            The server's UploadedFileManager.

        """
        # Each ReportSession has a unique string ID.
        if is_preheat:
            self.id = PREHEATED_ID
        else:
            self.id = str(uuid.uuid4())

        self._ioloop = ioloop
        self._report = Report(script_path, command_line)
        self._uploaded_file_mgr = uploaded_file_manager

        self._state = ReportSessionState.REPORT_NOT_RUNNING

        self._widget_states = WidgetStates()
        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed
        )
        self._sent_initialize_message = False
        self._storage = None
        self._maybe_reuse_previous_run = False
        self._run_on_save = config.get_option("server.runOnSave")

        # The ScriptRequestQueue is the means by which we communicate
        # with the active ScriptRunner.
        self._script_request_queue = ScriptRequestQueue()

        self._scriptrunner = None

        LOGGER.debug("ReportSession initialized (id=%s)", self.id)
Example #3
0
    def __init__(self, ioloop, script_path, command_line):
        """Initialize the ReportSession.

        Parameters
        ----------
        ioloop : tornado.ioloop.IOLoop
            The Tornado IOLoop that we're running within.

        script_path : str
            Path of the Python file from which this report is generated.

        command_line : str
            Command line as input by the user.

        """
        # Each ReportSession gets a unique ID
        self.id = ReportSession._next_id
        ReportSession._next_id += 1

        self._ioloop = ioloop
        self._report = Report(script_path, command_line)

        self._state = ReportSessionState.REPORT_NOT_RUNNING

        self._uploaded_file_mgr = UploadedFileManager()

        self._main_dg = DeltaGenerator(enqueue=self.enqueue, container=BlockPath.MAIN)
        self._sidebar_dg = DeltaGenerator(
            enqueue=self.enqueue, container=BlockPath.SIDEBAR
        )

        self._widget_states = WidgetStates()
        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed
        )
        self._sent_initialize_message = False
        self._storage = None
        self._maybe_reuse_previous_run = False
        self._run_on_save = config.get_option("server.runOnSave")

        # The ScriptRequestQueue is the means by which we communicate
        # with the active ScriptRunner.
        self._script_request_queue = ScriptRequestQueue()

        self._scriptrunner = None

        LOGGER.debug("ReportSession initialized (id=%s)", self.id)
Example #4
0
class ReportSession(object):
    """
    Contains session data for a single "user" of an active report
    (that is, a connected browser tab).

    Each ReportSession has its own Report, root DeltaGenerator, ScriptRunner,
    and widget state.

    A ReportSession is attached to each thread involved in running its Report.
    """

    _next_id = 0

    def __init__(self, ioloop, script_path, script_argv):
        """Initialize the ReportSession.

        Parameters
        ----------
        ioloop : tornado.ioloop.IOLoop
            The Tornado IOLoop that we're running within

        script_path : str
            Path of the Python file from which this report is generated.

        script_argv : list of str
            Command-line arguments to run the script with.

        """
        # Each ReportSession gets a unique ID
        self.id = ReportSession._next_id
        ReportSession._next_id += 1

        self._ioloop = ioloop
        self._report = Report(script_path, script_argv)

        self._state = ReportSessionState.REPORT_NOT_RUNNING

        self._main_dg = DeltaGenerator(self.enqueue, container=BlockPath.MAIN)
        self._sidebar_dg = DeltaGenerator(self.enqueue,
                                          container=BlockPath.SIDEBAR)

        self._widget_states = WidgetStates()
        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed)
        self._sent_initialize_message = False
        self._storage = None
        self._maybe_reuse_previous_run = False
        self._run_on_save = config.get_option('server.runOnSave')

        # The ScriptRequestQueue is the means by which we communicate
        # with the active ScriptRunner.
        self._script_request_queue = ScriptRequestQueue()

        self._scriptrunner = None

        LOGGER.debug('ReportSession initialized (id=%s)', self.id)

    def flush_browser_queue(self):
        """Clears the report queue and returns the messages it contained.

        The Server calls this periodically to deliver new messages
        to the browser connected to this report.

        Returns
        -------
        list[ForwardMsg]
            The messages that were removed from the queue and should
            be delivered to the browser.

        """
        return self._report.flush_browser_queue()

    def shutdown(self):
        """Shuts down the ReportSession.

        It's an error to use a ReportSession after it's been shut down.

        """
        if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
            LOGGER.debug('Shutting down (id=%s)', self.id)

            # Shut down the ScriptRunner, if one is active.
            # self._state must not be set to SHUTDOWN_REQUESTED until
            # after this is called.
            if self._scriptrunner is not None:
                self._enqueue_script_request(ScriptRequest.SHUTDOWN)

            self._state = ReportSessionState.SHUTDOWN_REQUESTED
            self._local_sources_watcher.close()

    def enqueue(self, msg):
        """Enqueues a new ForwardMsg to our browser queue.

        This can be called on both the main thread and a ScriptRunner
        run thread.

        Parameters
        ----------
        msg : ForwardMsg
            The message to enqueue

        Returns
        -------
        bool
            True if the message was enqueued, or False if
            client.displayEnabled is not set

        """
        if not config.get_option('client.displayEnabled'):
            return False

        # If we have an active ScriptRunner, signal that it can handle
        # an execution control request. (Copy the scriptrunner reference
        # to avoid it being unset from underneath us, as this function can
        # be called outside the main thread.)
        scriptrunner = self._scriptrunner
        if scriptrunner is not None:
            scriptrunner.maybe_handle_execution_control_request()

        self._report.enqueue(msg)
        return True

    def enqueue_exception(self, e):
        """Enqueues an Exception message.

        Parameters
        ----------
        e : BaseException

        """
        # This does a few things:
        # 1) Clears the current report in the browser.
        # 2) Marks the current report as "stopped" in the browser.
        # 3) HACK: Resets any script params that may have been broken (e.g. the
        # command-line when rerunning with wrong argv[0])
        self._on_scriptrunner_event(
            ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS)
        self._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED)
        self._on_scriptrunner_event(
            ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS)

        msg = ForwardMsg()
        msg.delta.id = 0
        exception_proto.marshall(msg.delta.new_element.exception, e)

        self.enqueue(msg)

    def request_rerun(self, argv=None, widget_state=None):
        """Signal that we're interested in running the script.

        If the script is not already running, it will be started immediately.
        Otherwise, a rerun will be requested.

        Parameters
        ----------
        argv : dict | None
            The new command line arguments to run the script with, or None
            to use the argv from the previous run of the script.
        widget_state : dict | None
            The widget state dictionary to run the script with, or None
            to use the widget state from the previous run of the script.

        """
        self._enqueue_script_request(ScriptRequest.RERUN,
                                     RerunData(argv, widget_state))

    def _on_source_file_changed(self):
        """One of our source files changed. Schedule a rerun if appropriate."""
        if self._run_on_save:
            self.request_rerun()
        else:
            self._enqueue_file_change_message()

    def _clear_queue(self):
        self._report.clear()

    def _on_scriptrunner_event(self,
                               event,
                               exception=None,
                               widget_states=None):
        """Called when our ScriptRunner emits an event.

        This is *not* called on the main thread.

        Parameters
        ----------
        event : ScriptRunnerEvent

        exception : BaseException | None
            An exception thrown during compilation. Set only for the
            SCRIPT_STOPPED_WITH_COMPILE_ERROR event.

        widget_states : streamlit.proto.Widget_pb2.WidgetStates | None
            The ScriptRunner's final WidgetStates. Set only for the
            SHUTDOWN event.

        """
        LOGGER.debug('OnScriptRunnerEvent: %s', event)

        prev_state = self._state

        if event == ScriptRunnerEvent.SCRIPT_STARTED:
            if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
                self._state = ReportSessionState.REPORT_IS_RUNNING

            if config.get_option('server.liveSave'):
                # Enqueue into the IOLoop so it runs without blocking AND runs
                # on the main thread.
                self._ioloop.spawn_callback(self._save_running_report)

            self._clear_queue()
            self._maybe_enqueue_initialize_message()
            self._enqueue_new_report_message()

        elif (event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS
              or event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR):

            if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
                self._state = ReportSessionState.REPORT_NOT_RUNNING

            self._enqueue_report_finished_message()

            if config.get_option('server.liveSave'):
                # Enqueue into the IOLoop so it runs without blocking AND runs
                # on the main thread.
                self._ioloop.spawn_callback(self._save_final_report_and_quit)

            script_succeeded = \
                event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS

            if script_succeeded:
                # When a script completes successfully, we update our
                # LocalSourcesWatcher to account for any source code changes
                # that change which modules should be watched. (This is run on
                # the main thread, because LocalSourcesWatcher is not
                # thread safe.)
                self._ioloop.spawn_callback(
                    self._local_sources_watcher.update_watched_modules)
            else:
                # When a script fails to compile, we send along the exception.
                from streamlit.elements import exception_proto
                msg = ForwardMsg()
                exception_proto.marshall(
                    msg.session_event.script_compilation_exception, exception)
                self.enqueue(msg)

        elif event == ScriptRunnerEvent.SHUTDOWN:
            # When ScriptRunner shuts down, update our local reference to it,
            # and check to see if we need to spawn a new one. (This is run on
            # the main thread.)
            def on_shutdown():
                self._widget_states = widget_states
                self._scriptrunner = None
                # Because a new ScriptEvent could have been enqueued while the
                # scriptrunner was shutting down, we check to see if we should
                # create a new one. (Otherwise, a newly-enqueued ScriptEvent
                # won't be processed until another event is enqueued.)
                self._maybe_create_scriptrunner()

            self._ioloop.spawn_callback(on_shutdown)

        # Send a message if our run state changed
        report_was_running = prev_state == ReportSessionState.REPORT_IS_RUNNING
        report_is_running = self._state == ReportSessionState.REPORT_IS_RUNNING
        if report_is_running != report_was_running:
            self._enqueue_session_state_changed_message()

    def _enqueue_session_state_changed_message(self):
        msg = ForwardMsg()
        msg.session_state_changed.run_on_save = self._run_on_save
        msg.session_state_changed.report_is_running = \
            self._state == ReportSessionState.REPORT_IS_RUNNING
        self.enqueue(msg)

    def _enqueue_file_change_message(self):
        LOGGER.debug('Enqueuing report_changed message (id=%s)', self.id)
        msg = ForwardMsg()
        msg.session_event.report_changed_on_disk = True
        self.enqueue(msg)

    def _maybe_enqueue_initialize_message(self):
        if self._sent_initialize_message:
            return

        self._sent_initialize_message = True

        msg = ForwardMsg()
        imsg = msg.initialize

        imsg.config.sharing_enabled = (config.get_option('global.sharingMode')
                                       != 'off')
        LOGGER.debug('New browser connection: sharing_enabled=%s',
                     imsg.config.sharing_enabled)

        imsg.config.gather_usage_stats = (
            config.get_option('browser.gatherUsageStats'))
        LOGGER.debug('New browser connection: gather_usage_stats=%s',
                     imsg.config.gather_usage_stats)

        imsg.environment_info.streamlit_version = __version__
        imsg.environment_info.python_version = ('.'.join(
            map(str, sys.version_info)))

        imsg.session_state.run_on_save = self._run_on_save
        imsg.session_state.report_is_running = (
            self._state == ReportSessionState.REPORT_IS_RUNNING)

        imsg.user_info.installation_id = __installation_id__
        imsg.user_info.email = Credentials.get_current().activation.email

        self.enqueue(msg)

    def _enqueue_new_report_message(self):
        self._report.generate_new_id()
        msg = ForwardMsg()
        msg.new_report.id = self._report.report_id
        msg.new_report.command_line.extend(self._report.argv)
        msg.new_report.name = self._report.name
        msg.new_report.script_path = self._report.script_path
        self.enqueue(msg)

    def _enqueue_report_finished_message(self):
        msg = ForwardMsg()
        msg.report_finished = True
        self.enqueue(msg)

    def handle_rerun_script_request(self,
                                    command_line=None,
                                    widget_state=None,
                                    is_preheat=False):
        """Tells the ScriptRunner to re-run its report.

        Parameters
        ----------
        command_line : str | None
            The new command line arguments to run the script with, or None
            to use its previous command line value.
        widget_state : WidgetStates | None
            The WidgetStates protobuf to run the script with, or None
            to use its previous widget states.
        is_preheat: boolean
            True if this ReportSession should run the script immediately, and
            then ignore the next rerun request if it matches the already-ran
            argv and widget state.

        """
        old_argv = self._report.argv

        if command_line is not None:
            self._report.parse_argv_from_command_line(command_line)

        if is_preheat:
            self._maybe_reuse_previous_run = True  # For next time.

        elif self._maybe_reuse_previous_run:
            # If this is a "preheated" ReportSession, reuse the previous run if
            # the argv and widget state matches. But only do this one time
            # ever.
            self._maybe_reuse_previous_run = False

            has_widget_state = (widget_state is not None
                                and len(widget_state.widgets) > 0)
            has_new_argv = old_argv != self._report.argv

            if not has_widget_state and not has_new_argv:
                LOGGER.debug(
                    'Skipping rerun since the preheated run is the same')
                return

        self.request_rerun(self._report.argv, widget_state)

    def handle_stop_script_request(self):
        """Tells the ScriptRunner to stop running its report."""
        self._enqueue_script_request(ScriptRequest.STOP)

    def handle_clear_cache_request(self):
        """Clears this report's cache.

        Because this cache is global, it will be cleared for all users.

        """
        # Setting verbose=True causes clear_cache to print to stdout.
        # Since this command was initiated from the browser, the user
        # doesn't need to see the results of the command in their
        # terminal.
        caching.clear_cache()

    def handle_set_run_on_save_request(self, new_value):
        """Changes our run_on_save flag to the given value.

        The browser will be notified of the change.

        Parameters
        ----------
        new_value : bool
            New run_on_save value

        """
        self._run_on_save = new_value
        self._enqueue_session_state_changed_message()

    def _enqueue_script_request(self, request, data=None):
        """Enqueue a ScriptEvent into our ScriptEventQueue.

        If a script thread is not already running, one will be created
        to handle the event.

        Parameters
        ----------
        request : ScriptRequest
            The type of request.

        data : Any
            Data associated with the request, if any.

        """
        if self._state == ReportSessionState.SHUTDOWN_REQUESTED:
            LOGGER.warning('Discarding %s request after shutdown' % request)
            return

        self._script_request_queue.enqueue(request, data)
        self._maybe_create_scriptrunner()

    def _maybe_create_scriptrunner(self):
        """Create a new ScriptRunner if we have unprocessed script requests.

        This is called every time a ScriptRequest is enqueued, and also
        after a ScriptRunner shuts down, in case new requests were enqueued
        during its termination.

        This function should only be called on the main thread.

        """
        if (self._state == ReportSessionState.SHUTDOWN_REQUESTED
                or self._scriptrunner is not None
                or not self._script_request_queue.has_request):
            return

        # Create the ScriptRunner, attach event handlers, and start it
        self._scriptrunner = ScriptRunner(
            report=self._report,
            main_dg=self._main_dg,
            sidebar_dg=self._sidebar_dg,
            widget_states=self._widget_states,
            request_queue=self._script_request_queue)
        self._scriptrunner.on_event.connect(self._on_scriptrunner_event)
        self._scriptrunner.start()

    @tornado.gen.coroutine
    def handle_save_request(self, ws):
        """Save serialized version of report deltas to the cloud."""
        @tornado.gen.coroutine
        def progress(percent):
            progress_msg = ForwardMsg()
            progress_msg.upload_report_progress = percent
            yield ws.write_message(progress_msg.SerializeToString(),
                                   binary=True)

        # Indicate that the save is starting.
        try:
            yield progress(0)

            url = yield self._save_final_report(progress)

            # Indicate that the save is done.
            progress_msg = ForwardMsg()
            progress_msg.report_uploaded = url
            yield ws.write_message(progress_msg.SerializeToString(),
                                   binary=True)

        except Exception as e:
            # Horrible hack to show something if something breaks.
            err_msg = '%s: %s' % (type(e).__name__, str(e)
                                  or 'No further details.')
            progress_msg = ForwardMsg()
            progress_msg.report_uploaded = err_msg
            yield ws.write_message(progress_msg.SerializeToString(),
                                   binary=True)
            raise e

    @tornado.gen.coroutine
    def _save_running_report(self):
        files = self._report.serialize_running_report_to_files()
        url = yield self._get_storage().save_report_files(
            self._report.report_id, files)

        if config.get_option('server.liveSave'):
            util.print_url('Saved running report', url)

        raise tornado.gen.Return(url)

    @tornado.gen.coroutine
    def _save_final_report(self, progress=None):
        files = self._report.serialize_final_report_to_files()
        url = yield self._get_storage().save_report_files(
            self._report.report_id, files, progress)

        if config.get_option('server.liveSave'):
            util.print_url('Saved final report', url)

        raise tornado.gen.Return(url)

    @tornado.gen.coroutine
    def _save_final_report_and_quit(self):
        yield self._save_final_report()
        self._ioloop.stop()

    def _get_storage(self):
        if self._storage is None:
            self._storage = Storage()
        return self._storage
Example #5
0
class ScriptRunner(object):
    def __init__(self, report):
        """Initialize.

        Parameters
        ----------
        report : Report
            The report with the script to run.

        """
        self._report = report
        self._event_queue = ScriptEventQueue()
        self._state = ScriptState.STOPPED
        self._last_run_data = RerunData(argv=report.argv,
                                        widget_state=WidgetStates())
        self._widgets = Widgets()

        self.run_on_save = config.get_option('server.runOnSave')

        self.on_state_changed = Signal(
            doc="""Emitted when the script's execution state state changes.

            Parameters
            ----------
            state : ScriptState
            """)

        self.on_file_change_not_handled = Signal(
            doc="Emitted when the file is modified and we haven't handled it.")

        self.on_script_compile_error = Signal(
            doc="""Emitted if our script fails to compile.  (*Not* emitted
            for normal exceptions thrown while a script is running.)

            Parameters
            ----------
            exception : Exception
                The exception that was thrown
            """)

        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed)

        # Will be set to true when we process a SHUTDOWN event
        self._shutdown_requested = False

        self._script_thread = None
        self._ctx = None

        # Set to true while we're executing. Used by
        # maybe_handle_execution_control_request.
        self._execing = False

    @property
    def widgets(self):
        """
        Returns
        -------
        Widgets
            Our widget state dictionary

        """
        return self._widgets

    def start_run_loop(self, ctx):
        """Starts the ScriptRunner's thread.

        This must be called only once.

        Parameters
        ----------
        ctx : ReportContext
            The ReportContext that owns this ScriptRunner. This will be
            attached to the ScriptRunner's runloop thread.

        """

        assert self._script_thread is None, 'Already started!'
        self._ctx = ctx

        # start our thread
        self._script_thread = ReportThread(ctx,
                                           target=self._loop,
                                           name='ScriptRunner.loop')
        self._script_thread.start()

    def _set_state(self, new_state):
        if self._state == new_state:
            return

        LOGGER.debug('Scriptrunner state: %s -> %s' % (self._state, new_state))
        self._state = new_state
        self.on_state_changed.send(self._state)

    def is_running(self):
        return self._state == ScriptState.RUNNING

    def is_shutdown(self):
        return self._state == ScriptState.IS_SHUTDOWN

    def request_rerun(self, argv=None, widget_state=None):
        """Signal that we're interested in running the script.

        If the script is not already running, it will be started immediately.
        Otherwise, a rerun will be requested.

        Parameters
        ----------
        argv : dict | None
            The new command line arguments to run the script with, or None
            to use the argv from the previous run of the script.
        widget_state : dict | None
            The widget state dictionary to run the script with, or None
            to use the widget state from the previous run of the script.

        """
        if self.is_shutdown():
            LOGGER.warning('Discarding RERUN event after shutdown')
            return
        self._event_queue.enqueue(ScriptEvent.RERUN,
                                  RerunData(argv, widget_state))

    def request_stop(self):
        if self.is_shutdown():
            LOGGER.warning('Discarding STOP event after shutdown')
            return
        self._event_queue.enqueue(ScriptEvent.STOP)

    def request_shutdown(self):
        if self.is_shutdown():
            LOGGER.warning('Discarding SHUTDOWN event after shutdown')
            return
        self._event_queue.enqueue(ScriptEvent.SHUTDOWN)

    def maybe_handle_execution_control_request(self):
        if self._script_thread != threading.current_thread():
            # We can only handle execution_control_request if we're on the
            # script execution thread. However, it's possible for deltas to
            # be enqueued (and, therefore, for this function to be called)
            # in separate threads, so we check for that here.
            return

        if not self._execing:
            # If the _execing flag is not set, we're not actually inside
            # an exec() call. This happens when our script exec() completes,
            # we change our state to STOPPED, and a statechange-listener
            # enqueues a new ForwardEvent
            return

        # Pop the next event from our queue. Don't block if there's no event
        event, event_data = self._event_queue.dequeue_nowait()
        if event is None:
            return

        LOGGER.debug('Received ScriptEvent: %s', event)
        if event == ScriptEvent.STOP:
            raise StopException()
        elif event == ScriptEvent.SHUTDOWN:
            self._shutdown_requested = True
            raise StopException()
        elif event == ScriptEvent.RERUN:
            raise RerunException(event_data)
        else:
            raise RuntimeError('Unrecognized ScriptEvent: %s' % event)

    def _on_source_file_changed(self):
        """One of our source files changed. Schedule a rerun if appropriate."""
        if self.run_on_save:
            self._event_queue.enqueue(ScriptEvent.RERUN,
                                      RerunData(argv=None, widget_state=None))
        else:
            self.on_file_change_not_handled.send()

    def _loop(self):
        """Our run loop.

        Continually pops events from the event_queue. Ends when we receive
        a SHUTDOWN event.

        """
        while not self._shutdown_requested:
            assert self._state == ScriptState.STOPPED

            # Dequeue our next event. If the event queue is empty, the thread
            # will go to sleep, and awake when there's a new event.
            event, event_data = self._event_queue.dequeue()
            if event == ScriptEvent.STOP:
                LOGGER.debug('Ignoring STOP event while not running')
            elif event == ScriptEvent.SHUTDOWN:
                LOGGER.debug('Shutting down')
                self._shutdown_requested = True
            elif event == ScriptEvent.RERUN:
                self._run_script(event_data)
            else:
                raise RuntimeError('Unrecognized ScriptEvent: %s' % event)

        # Release resources
        self._local_sources_watcher.close()

        self._set_state(ScriptState.IS_SHUTDOWN)

    def _install_tracer(self):
        """Install function that runs before each line of the script."""
        def trace_calls(frame, event, arg):
            self.maybe_handle_execution_control_request()
            return trace_calls

        # Python interpreters are not required to implement sys.settrace.
        if hasattr(sys, 'settrace'):
            sys.settrace(trace_calls)

    @contextmanager
    def _set_execing_flag(self):
        """A context for setting the ScriptRunner._execing flag.

        Used by maybe_handle_execution_control_request to ensure that
        we only handle requests while we're inside an exec() call
        """
        if self._execing:
            raise RuntimeError('Nested set_execing_flag call')
        self._execing = True
        try:
            yield
        finally:
            self._execing = False

    def _run_script(self, rerun_data):
        """Run our script.

        Parameters
        ----------
        rerun_data: RerunData
            The RerunData to use.

        """
        assert self._state == ScriptState.STOPPED
        LOGGER.debug('Running script %s', rerun_data)

        # Reset delta generator so it starts from index 0.
        import streamlit as st
        st._reset()

        self._set_state(ScriptState.RUNNING)

        # Compile the script. Any errors thrown here will be surfaced
        # to the user via a modal dialog, and won't result in their
        # previous report disappearing.
        try:
            # Python 3 got rid of the native execfile() command, so we now read
            # the file, compile it, and exec() it. This implementation is
            # compatible with both 2 and 3.
            with open(self._report.script_path) as f:
                filebody = f.read()

            if config.get_option('runner.magicEnabled'):
                filebody = magic.add_magic(filebody, self._report.script_path)

            code = compile(
                filebody,
                # Pass in the file path so it can show up in exceptions.
                self._report.script_path,
                # We're compiling entire blocks of Python, so we need "exec"
                # mode (as opposed to "eval" or "single").
                'exec',
                # Don't inherit any flags or "future" statements.
                flags=0,
                dont_inherit=1,
                # Parameter not supported in Python2:
                # optimize=-1,
            )

        except BaseException as e:
            # We got a compile error. Send the exception onto the client
            # as a SessionEvent and bail immediately.
            LOGGER.debug('Fatal script error: %s' % e)
            self.on_script_compile_error.send(e)
            self._set_state(ScriptState.STOPPED)
            return

        # If we get here, we've successfully compiled our script. The next step
        # is to run it. Errors thrown during execution will be shown to the
        # user as ExceptionElements.

        # Get our argv and widget_state for this run, defaulting to
        # self._last_run_data for missing values.
        # Also update self._last_run_data for the next run.
        argv = rerun_data.argv or self._last_run_data.argv
        widget_state = rerun_data.widget_state or \
                       self._last_run_data.widget_state
        self._last_run_data = RerunData(argv,
                                        reset_widget_triggers(widget_state))

        # Update the Widget singleton with the new widget_state
        self._ctx.widgets.set_state(widget_state)

        if config.get_option('runner.installTracer'):
            self._install_tracer()

        # This will be set to a RerunData instance if our execution
        # is interrupted by a RerunException.
        rerun_with_data = None

        try:
            # Create fake module. This gives us a name global namespace to
            # execute the code in.
            module = _new_module('__main__')

            # Install the fake module as the __main__ module. This allows
            # the pickle module to work inside the user's code, since it now
            # can know the module where the pickled objects stem from.
            # IMPORTANT: This means we can't use "if __name__ == '__main__'" in
            # our code, as it will point to the wrong module!!!
            sys.modules['__main__'] = module

            # Make it look like command-line args were set to whatever the user
            # asked them to be via the GUI.
            # IMPORTANT: This means we can't count on sys.argv in our code ---
            # but we already knew it from the argv surgery in cli.py.
            # TODO: Remove this feature when we implement interactivity!
            #  This is not robust in a multi-user environment.
            sys.argv = argv

            # Add special variables to the module's globals dict.
            module.__dict__['__file__'] = self._report.script_path

            with modified_sys_path(self._report), self._set_execing_flag():
                exec(code, module.__dict__)

        except RerunException as e:
            rerun_with_data = e.rerun_data

        except StopException:
            pass

        except BaseException as e:
            # Show exceptions in the Streamlit report.
            LOGGER.debug(e)
            st.exception(e)  # This is OK because we're in the script thread.
            # TODO: Clean up the stack trace, so it doesn't include
            # ScriptRunner.

        finally:
            self._set_state(ScriptState.STOPPED)

        # Use _log_if_error() to make sure we never ever ever stop running the
        # script without meaning to.
        _log_if_error(self._local_sources_watcher.update_watched_modules)
        _log_if_error(_clean_problem_modules)

        if rerun_with_data is not None:
            self._run_script(rerun_with_data)
Example #6
0
class ReportSession(object):
    """
    Contains session data for a single "user" of an active report
    (that is, a connected browser tab).

    Each ReportSession has its own Report, root DeltaGenerator, ScriptRunner,
    and widget state.

    A ReportSession is attached to each thread involved in running its Report.
    """

    _next_id = 0

    def __init__(self, ioloop, script_path, command_line):
        """Initialize the ReportSession.

        Parameters
        ----------
        ioloop : tornado.ioloop.IOLoop
            The Tornado IOLoop that we're running within.

        script_path : str
            Path of the Python file from which this report is generated.

        command_line : str
            Command line as input by the user.

        """
        # Each ReportSession gets a unique ID
        self.id = ReportSession._next_id
        ReportSession._next_id += 1

        self._ioloop = ioloop
        self._report = Report(script_path, command_line)

        self._state = ReportSessionState.REPORT_NOT_RUNNING

        self._uploaded_file_mgr = UploadedFileManager()

        self._widget_states = WidgetStates()
        self._local_sources_watcher = LocalSourcesWatcher(
            self._report, self._on_source_file_changed)
        self._sent_initialize_message = False
        self._storage = None
        self._maybe_reuse_previous_run = False
        self._run_on_save = config.get_option("server.runOnSave")

        # The ScriptRequestQueue is the means by which we communicate
        # with the active ScriptRunner.
        self._script_request_queue = ScriptRequestQueue()

        self._scriptrunner = None

        LOGGER.debug("ReportSession initialized (id=%s)", self.id)

    def flush_browser_queue(self):
        """Clears the report queue and returns the messages it contained.

        The Server calls this periodically to deliver new messages
        to the browser connected to this report.

        Returns
        -------
        list[ForwardMsg]
            The messages that were removed from the queue and should
            be delivered to the browser.

        """
        return self._report.flush_browser_queue()

    def shutdown(self):
        """Shuts down the ReportSession.

        It's an error to use a ReportSession after it's been shut down.

        """
        if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
            LOGGER.debug("Shutting down (id=%s)", self.id)
            self._uploaded_file_mgr.delete_all_files()

            # Shut down the ScriptRunner, if one is active.
            # self._state must not be set to SHUTDOWN_REQUESTED until
            # after this is called.
            if self._scriptrunner is not None:
                self._enqueue_script_request(ScriptRequest.SHUTDOWN)

            self._state = ReportSessionState.SHUTDOWN_REQUESTED
            self._local_sources_watcher.close()

    def enqueue(self, msg):
        """Enqueues a new ForwardMsg to our browser queue.

        This can be called on both the main thread and a ScriptRunner
        run thread.

        Parameters
        ----------
        msg : ForwardMsg
            The message to enqueue

        """
        if not config.get_option("client.displayEnabled"):
            return

        # Avoid having two maybe_handle_execution_control_request running on
        # top of each other when tracer is installed. This leads to a lock
        # contention.
        if not config.get_option("runner.installTracer"):
            # If we have an active ScriptRunner, signal that it can handle an
            # execution control request. (Copy the scriptrunner reference to
            # avoid it being unset from underneath us, as this function can be
            # called outside the main thread.)
            scriptrunner = self._scriptrunner

            if scriptrunner is not None:
                scriptrunner.maybe_handle_execution_control_request()

        self._report.enqueue(msg)

    def enqueue_exception(self, e):
        """Enqueues an Exception message.

        Parameters
        ----------
        e : BaseException

        """
        # This does a few things:
        # 1) Clears the current report in the browser.
        # 2) Marks the current report as "stopped" in the browser.
        # 3) HACK: Resets any script params that may have been broken (e.g. the
        # command-line when rerunning with wrong argv[0])
        self._on_scriptrunner_event(
            ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS)
        self._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED)
        self._on_scriptrunner_event(
            ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS)

        msg = ForwardMsg()
        msg.metadata.delta_id = 0
        exception_proto.marshall(msg.delta.new_element.exception, e)

        self.enqueue(msg)

    def request_rerun(self, widget_state=None):
        """Signal that we're interested in running the script.

        If the script is not already running, it will be started immediately.
        Otherwise, a rerun will be requested.

        Parameters
        ----------
        widget_state : dict | None
            The widget state dictionary to run the script with, or None
            to use the widget state from the previous run of the script.

        """
        self._enqueue_script_request(ScriptRequest.RERUN,
                                     RerunData(widget_state))

    def _on_source_file_changed(self):
        """One of our source files changed. Schedule a rerun if appropriate."""
        if self._run_on_save:
            self.request_rerun()
        else:
            self._enqueue_file_change_message()

    def _clear_queue(self):
        self._report.clear()

    def _on_scriptrunner_event(self,
                               event,
                               exception=None,
                               widget_states=None):
        """Called when our ScriptRunner emits an event.

        This is *not* called on the main thread.

        Parameters
        ----------
        event : ScriptRunnerEvent

        exception : BaseException | None
            An exception thrown during compilation. Set only for the
            SCRIPT_STOPPED_WITH_COMPILE_ERROR event.

        widget_states : streamlit.proto.Widget_pb2.WidgetStates | None
            The ScriptRunner's final WidgetStates. Set only for the
            SHUTDOWN event.

        """
        LOGGER.debug("OnScriptRunnerEvent: %s", event)

        prev_state = self._state

        if event == ScriptRunnerEvent.SCRIPT_STARTED:
            if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
                self._state = ReportSessionState.REPORT_IS_RUNNING

            if config.get_option("server.liveSave"):
                # Enqueue into the IOLoop so it runs without blocking AND runs
                # on the main thread.
                self._ioloop.spawn_callback(self._save_running_report)

            self._clear_queue()
            self._maybe_enqueue_initialize_message()
            self._enqueue_new_report_message()

        elif (event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS
              or event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_COMPILE_ERROR):

            if self._state != ReportSessionState.SHUTDOWN_REQUESTED:
                self._state = ReportSessionState.REPORT_NOT_RUNNING

            script_succeeded = event == ScriptRunnerEvent.SCRIPT_STOPPED_WITH_SUCCESS

            self._enqueue_report_finished_message(
                ForwardMsg.FINISHED_SUCCESSFULLY if script_succeeded else
                ForwardMsg.FINISHED_WITH_COMPILE_ERROR)

            if config.get_option("server.liveSave"):
                # Enqueue into the IOLoop so it runs without blocking AND runs
                # on the main thread.
                self._ioloop.spawn_callback(self._save_final_report_and_quit)

            if script_succeeded:
                # When a script completes successfully, we update our
                # LocalSourcesWatcher to account for any source code changes
                # that change which modules should be watched. (This is run on
                # the main thread, because LocalSourcesWatcher is not
                # thread safe.)
                self._ioloop.spawn_callback(
                    self._local_sources_watcher.update_watched_modules)
            else:
                # When a script fails to compile, we send along the exception.
                from streamlit.elements import exception_proto

                msg = ForwardMsg()
                exception_proto.marshall(
                    msg.session_event.script_compilation_exception, exception)
                self.enqueue(msg)

        elif event == ScriptRunnerEvent.SHUTDOWN:
            # When ScriptRunner shuts down, update our local reference to it,
            # and check to see if we need to spawn a new one. (This is run on
            # the main thread.)
            def on_shutdown():
                self._widget_states = widget_states
                self._scriptrunner = None
                # Because a new ScriptEvent could have been enqueued while the
                # scriptrunner was shutting down, we check to see if we should
                # create a new one. (Otherwise, a newly-enqueued ScriptEvent
                # won't be processed until another event is enqueued.)
                self._maybe_create_scriptrunner()

            self._ioloop.spawn_callback(on_shutdown)

        # Send a message if our run state changed
        report_was_running = prev_state == ReportSessionState.REPORT_IS_RUNNING
        report_is_running = self._state == ReportSessionState.REPORT_IS_RUNNING
        if report_is_running != report_was_running:
            self._enqueue_session_state_changed_message()

    def _enqueue_session_state_changed_message(self):
        msg = ForwardMsg()
        msg.session_state_changed.run_on_save = self._run_on_save
        msg.session_state_changed.report_is_running = (
            self._state == ReportSessionState.REPORT_IS_RUNNING)
        self.enqueue(msg)

    def _enqueue_file_change_message(self):
        LOGGER.debug("Enqueuing report_changed message (id=%s)", self.id)
        msg = ForwardMsg()
        msg.session_event.report_changed_on_disk = True
        self.enqueue(msg)

    def _maybe_enqueue_initialize_message(self):
        if self._sent_initialize_message:
            return

        self._sent_initialize_message = True

        msg = ForwardMsg()
        imsg = msg.initialize

        imsg.config.sharing_enabled = config.get_option(
            "global.sharingMode") != "off"

        imsg.config.gather_usage_stats = config.get_option(
            "browser.gatherUsageStats")

        imsg.config.max_cached_message_age = config.get_option(
            "global.maxCachedMessageAge")

        imsg.config.mapbox_token = config.get_option("mapbox.token")

        LOGGER.debug(
            "New browser connection: "
            "gather_usage_stats=%s, "
            "sharing_enabled=%s, "
            "max_cached_message_age=%s",
            imsg.config.gather_usage_stats,
            imsg.config.sharing_enabled,
            imsg.config.max_cached_message_age,
        )

        imsg.environment_info.streamlit_version = __version__
        imsg.environment_info.python_version = ".".join(
            map(str, sys.version_info))

        imsg.session_state.run_on_save = self._run_on_save
        imsg.session_state.report_is_running = (
            self._state == ReportSessionState.REPORT_IS_RUNNING)

        imsg.user_info.installation_id = __installation_id__
        if Credentials.get_current().activation:
            imsg.user_info.email = Credentials.get_current().activation.email
        else:
            imsg.user_info.email = ""

        imsg.command_line = self._report.command_line

        self.enqueue(msg)

    def _enqueue_new_report_message(self):
        self._report.generate_new_id()
        msg = ForwardMsg()
        msg.new_report.id = self._report.report_id
        msg.new_report.name = self._report.name
        msg.new_report.script_path = self._report.script_path
        self.enqueue(msg)

    def _enqueue_report_finished_message(self, status):
        """Enqueues a report_finished ForwardMsg.

        Parameters
        ----------
        status : ReportFinishedStatus

        """
        msg = ForwardMsg()
        msg.report_finished = status
        self.enqueue(msg)

    def handle_rerun_script_request(self,
                                    command_line=None,
                                    widget_state=None,
                                    is_preheat=False):
        """Tells the ScriptRunner to re-run its report.

        Parameters
        ----------
        command_line : str | None
            The new command line arguments to run the script with, or None
            to use its previous command line value.
        widget_state : WidgetStates | None
            The WidgetStates protobuf to run the script with, or None
            to use its previous widget states.
        is_preheat: boolean
            True if this ReportSession should run the script immediately, and
            then ignore the next rerun request if it matches the already-ran
            widget state.

        """
        if is_preheat:
            self._maybe_reuse_previous_run = True  # For next time.

        elif self._maybe_reuse_previous_run:
            # If this is a "preheated" ReportSession, reuse the previous run if
            # the widget state matches. But only do this one time ever.
            self._maybe_reuse_previous_run = False

            has_widget_state = (widget_state is not None
                                and len(widget_state.widgets) > 0)

            if not has_widget_state:
                LOGGER.debug(
                    "Skipping rerun since the preheated run is the same")
                return

        self.request_rerun(widget_state)

    def handle_upload_file(self, upload_file):
        self._uploaded_file_mgr.create_or_clear_file(
            widget_id=upload_file.widget_id,
            name=upload_file.name,
            size=upload_file.size,
            last_modified=upload_file.lastModified,
            chunks=upload_file.chunks,
        )

        self.handle_rerun_script_request(widget_state=self._widget_states)

    def handle_upload_file_chunk(self, upload_file_chunk):
        progress = self._uploaded_file_mgr.process_chunk(
            widget_id=upload_file_chunk.widget_id,
            index=upload_file_chunk.index,
            data=upload_file_chunk.data,
        )

        if progress == 1:
            self.handle_rerun_script_request(widget_state=self._widget_states)

    def handle_delete_uploaded_file(self, delete_uploaded_file):
        self._uploaded_file_mgr.delete_file(
            widget_id=delete_uploaded_file.widget_id)
        self.handle_rerun_script_request(widget_state=self._widget_states)

    def handle_stop_script_request(self):
        """Tells the ScriptRunner to stop running its report."""
        self._enqueue_script_request(ScriptRequest.STOP)

    def handle_clear_cache_request(self):
        """Clears this report's cache.

        Because this cache is global, it will be cleared for all users.

        """
        # Setting verbose=True causes clear_cache to print to stdout.
        # Since this command was initiated from the browser, the user
        # doesn't need to see the results of the command in their
        # terminal.
        caching.clear_cache()

    def handle_set_run_on_save_request(self, new_value):
        """Changes our run_on_save flag to the given value.

        The browser will be notified of the change.

        Parameters
        ----------
        new_value : bool
            New run_on_save value

        """
        self._run_on_save = new_value
        self._enqueue_session_state_changed_message()

    def _enqueue_script_request(self, request, data=None):
        """Enqueue a ScriptEvent into our ScriptEventQueue.

        If a script thread is not already running, one will be created
        to handle the event.

        Parameters
        ----------
        request : ScriptRequest
            The type of request.

        data : Any
            Data associated with the request, if any.

        """
        if self._state == ReportSessionState.SHUTDOWN_REQUESTED:
            LOGGER.warning("Discarding %s request after shutdown" % request)
            return

        self._script_request_queue.enqueue(request, data)
        self._maybe_create_scriptrunner()

    def _maybe_create_scriptrunner(self):
        """Create a new ScriptRunner if we have unprocessed script requests.

        This is called every time a ScriptRequest is enqueued, and also
        after a ScriptRunner shuts down, in case new requests were enqueued
        during its termination.

        This function should only be called on the main thread.

        """
        if (self._state == ReportSessionState.SHUTDOWN_REQUESTED
                or self._scriptrunner is not None
                or not self._script_request_queue.has_request):
            return

        # Create the ScriptRunner, attach event handlers, and start it
        self._scriptrunner = ScriptRunner(
            report=self._report,
            enqueue_forward_msg=self.enqueue,
            widget_states=self._widget_states,
            request_queue=self._script_request_queue,
            uploaded_file_mgr=self._uploaded_file_mgr,
        )
        self._scriptrunner.on_event.connect(self._on_scriptrunner_event)
        self._scriptrunner.start()

    @tornado.gen.coroutine
    def handle_save_request(self, ws):
        """Save serialized version of report deltas to the cloud.

        "Progress" ForwardMsgs will be sent to the client during the upload.
        These messages are sent "out of band" - that is, they don't get
        enqueued into the ReportQueue (because they're not part of the report).
        Instead, they're written directly to the report's WebSocket.

        Parameters
        ----------
        ws : _BrowserWebSocketHandler
            The report's websocket handler.

        """
        @tornado.gen.coroutine
        def progress(percent):
            progress_msg = ForwardMsg()
            progress_msg.upload_report_progress = percent
            yield ws.write_message(serialize_forward_msg(progress_msg),
                                   binary=True)

        # Indicate that the save is starting.
        try:
            yield progress(0)

            url = yield self._save_final_report(progress)

            # Indicate that the save is done.
            progress_msg = ForwardMsg()
            progress_msg.report_uploaded = url
            yield ws.write_message(serialize_forward_msg(progress_msg),
                                   binary=True)

        except Exception as e:
            # Horrible hack to show something if something breaks.
            err_msg = "%s: %s" % (type(e).__name__, str(e)
                                  or "No further details.")
            progress_msg = ForwardMsg()
            progress_msg.report_uploaded = err_msg
            yield ws.write_message(serialize_forward_msg(progress_msg),
                                   binary=True)

            LOGGER.warning("Failed to save report:", exc_info=e)

    @tornado.gen.coroutine
    def _save_running_report(self):
        files = self._report.serialize_running_report_to_files()
        url = yield self._get_storage().save_report_files(
            self._report.report_id, files)

        if config.get_option("server.liveSave"):
            url_util.print_url("Saved running app", url)

        raise tornado.gen.Return(url)

    @tornado.gen.coroutine
    def _save_final_report(self, progress_coroutine=None):
        files = self._report.serialize_final_report_to_files()
        url = yield self._get_storage().save_report_files(
            self._report.report_id, files, progress_coroutine)

        if config.get_option("server.liveSave"):
            url_util.print_url("Saved final app", url)

        raise tornado.gen.Return(url)

    @tornado.gen.coroutine
    def _save_final_report_and_quit(self):
        yield self._save_final_report()
        self._ioloop.stop()

    def _get_storage(self):
        if self._storage is None:
            sharing_mode = config.get_option("global.sharingMode")
            if sharing_mode == "s3":
                self._storage = S3Storage()
            elif sharing_mode == "file":
                self._storage = FileStorage()
            else:
                raise RuntimeError("Unsupported sharing mode '%s'" %
                                   sharing_mode)
        return self._storage