dbt-labs / dbt-rpc

A server that can compile and run queries in the context of a dbt project. Additionally, it provides methods that can be used to list and terminate running processes.
https://docs.getdbt.com/reference/commands/rpc
Apache License 2.0
19 stars 7 forks source link

Cl/fix poll list result main #138

Closed ChenyuLInx closed 1 year ago

ChenyuLInx commented 1 year ago

Tested by adding the following code to builtins.py at class Poll(RemoteBuiltinMethod[PollParameters, PollResult]): to swap out task id

        # switch to the task id of list command
        test = [key for key in self.task_manager.active_tasks]
        test.remove(task_id)
        task_id = test[0]

And the test we run looks like

import pytest
from .util import (
    get_querier,
    ProjectDefinition,
)

@pytest.mark.supported('postgres')
def test_rpc_run_threads(
    project_root, profiles_root, dbt_profile, unique_schema
):
    project = ProjectDefinition(
        models={'my_model.sql': 'select 1 as id'}
    )
    querier_ctx = get_querier(
        project_def=project,
        project_dir=project_root,
        profiles_dir=profiles_root,
        schema=unique_schema,
    )
    with querier_ctx as querier:
        # run a run so that we have a request token
        run_result = querier.run()
        # run a list command so we have a task going
        result = querier.list()
        # there's code added in poll to swap out the request token so we are actually
        # getting the list command result back
        result1 = querier.poll(run_result['result']['request_token'])
        assert 'error' not in result1