Creating a widget for running SPL (Splunk) queries and mo.state

[kyrre]

I'm building a custom widget to execute Splunk Processing Language (SPL) queries.

So, essentially it's code editor that allows the user to:

• Write SPL queries with syntax highlighting
• Specify a date and time range for the query
• Execute queries against Splunk Enterprise/Cloud
• Automatically create DuckDB tables from the results for further analysis

The main challenge I'm facing is state management. When users are working interactively with SPL queries, and they re-run a cell containing the widget, it gets recreated and loses its state - query text and table names get reset, breaking the workflow.

I'm currently using marimo's state functionality to persist widget state across cell re-executions with mo.state(), which lets the widget maintain query content and settings even when cells are re-run. The front-end is just using codemirror.

class SplunkAPI:
    """Simple Splunk widget manager"""

    def __init__(self, token, host="localhost", port=8089, autoLogin=False, con=None):
        # intentionally leak strings to avoid race condition where __del__
        # cleanup destroys state of newly created widgets with same ID
        self._widget_states = {}

        # store the ibis connection - create default if none provided
        if con is None:
            import ibis
            import duckdb

            self.con = ibis.duckdb.from_connection(duckdb)
        else:
            self.con = con

        # connect to splunk
        self.service = client.connect(
            host=host, port=port, splunkToken=token, autologin=autoLogin
        )
        print(f"[SPL] created splunkapi connected to {host}:{port}")

    def spl(self, **kwargs):
        """Create a new SPL editor widget"""
        return SPLWidget(context=self, **kwargs)


class SPLWidget(anywidget.AnyWidget):
    _esm = pathlib.Path(__file__).parent / "static" / "widget.js"
    _css = pathlib.Path(__file__).parent / "static" / "widget.css"

    query = traitlets.Unicode("").tag(sync=True)
    table_name = traitlets.Unicode("").tag(sync=True)

    def __init__(self, context: SplunkAPI, **kwargs):
        # get the variable name as widget ID from marimo cell definitions
        # mo.defs() returns a list of variable names assigned in the current cell
        # this gives us a unique identifier for state persistence across cell re-runs
        cell_defs = mo.defs()
        if not cell_defs:
            # this fails when the widget isn't assigned to a variable
            # e.g., just calling splunk.spl() without assignment
            raise ValueError(
                "Widget must be assigned to a variable (e.g., 'my_widget = splunk.spl()')"
            )

        self.context = context
        self.widget_id = cell_defs[0]  # use the first (and usually only) variable name

        print(f"[SPL] creating widget with ID: {self.widget_id}")

        # create or get existing state for both query and table_name
        # if widget_id already exists, we reuse the same marimo state
        # this allows state persistence when cells are re-run
        if self.widget_id not in self.context._widget_states:
            print(f"[SPL] creating new state for {self.widget_id}")
            get_query, set_query = mo.state("")
            get_table_name, set_table_name = mo.state("")
            self.context._widget_states[self.widget_id] = {
                "query": (get_query, set_query),
                "table_name": (get_table_name, set_table_name),
            }
        else:
            print(f"[SPL] reusing existing state for {self.widget_id}")

        current_query = self._get_state_value("query")
        current_table_name = self._get_state_value("table_name")

        print(
            f"[SPL] initializing with query: '{current_query}', table_name: '{current_table_name}'"
        )

        super().__init__(query=current_query, table_name=current_table_name, **kwargs)
        self.observe(self._on_query_change, names=["query"])
        self.observe(self._on_table_name_change, names=["table_name"])

        # handle messages from frontend
        self.on_msg(self._handle_frontend_msg)

    def _on_query_change(self, change):
        """Handle query changes with logging"""
        print(f"[SPL] query changed for {self.widget_id}: '{change['new']}'")
        self._set_state_value("query", change["new"])

    def _on_table_name_change(self, change):
        """Handle table_name changes with logging"""
        print(f"[SPL] table name changed for {self.widget_id}: '{change['new']}'")
        self._set_state_value("table_name", change["new"])

    def search(self):
        """Execute search and create view"""
        if not self.query.strip():
            print(f"[SPL] no query to execute for {self.widget_id}")
            return None

        print(f"[SPL] executing search for {self.widget_id}: {self.query}")

        search_results = self.context.service.jobs.oneshot(
            self.query, output_mode="json"
        )

        reader = results.JSONResultsReader(search_results)
        data = pa.Table.from_pylist(
            [result for result in reader if isinstance(result, dict)]
        )

        print(f"[SPL] creating table '{self.table_name}' with {len(data)} rows")

        return self.context.con.create_table(self.table_name, data, overwrite=True)

    def _get_state(self, key):
        """Get state getter/setter tuple for a key"""
        return self.context._widget_states[self.widget_id][key]

    def _set_state_value(self, key, value):
        """Set a state value using the setter function"""
        _, setter = self._get_state(key)
        setter(value)

    def _get_state_value(self, key):
        """Get current state value using the getter function"""
        getter, _ = self._get_state(key)
        return getter()

    def _handle_frontend_msg(self, widget, content, buffers):
        """Handle messages from the frontend"""
        if content.get("type") == "execute_search":
            print(f"[SPL] execute button clicked for {self.widget_id}")
            table_name = self.table_name.strip()
            if table_name:
                try:
                    result = self.search()
                    print(
                        f"[SPL] search executed successfully, created table: {table_name}"
                    )
                except Exception as e:
                    print(f"[SPL] error executing search: {e}")
            else:
                print(f"[SPL] no table name provided for {self.widget_id}")

Would love any feedback on this approach! 🙏🏻

[dmadisetti]

You could use localstorage on the JS side to persist per user/ app.

State doesn't work well unless exposed as a toplevel variable. This works just by virtue of having hidden object references (even without state)- refactoring without state should make things a little cleaner

Another option is using mo.watch.file, which could save the user query if the queries are meant to be relatively persistent.

But cool! Feedback on DevX is appreciated

[kyrre]

Ohhhh. I think I originally had the "widget state" at top of the module,

# Move widget state to module level
_widget_states = {}

which deluded me into thinking I was doing something clever. 😅 Good thing I asked here!

The use case is companion notebooks for investigating a security incident or alert.

Such as a notebook will have some pre-defined information about the case, a description of what to do and some queries (Splunk, Databricks or Snowflake queries), e.g., for retrieving data in order to create a process tree visualisation.

The security engineer would then work on the case by doing a number of custom queries. By creating this widget I was hoping to avoid the engineers losing their work when refining their queries and to offer a similar UI to Splunk:

When I asked on discord two possible solutions were proposed:

1. anywidget
2. custom driver with sqlalchemy

I now realise it's kind of hard to achieve this workflow using anywidget. Do you have opinion on the second suggestion? It would solve the persistence issue at least, but you wouldn't have syntax highlighting nor the UI datepicker.

It might just make more sense to just continue with with what we have (string queries):

search(
con, "risk", f"""
 index=risk 
 _index_earliest=-36h@h _index_latest=-0@h  risk_object="{asset}"
 | table _time firstTime search_name annotations._all process process_id risk_message risk_object risk_score
""")

(then we also get access to the AI features!)