diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts index cee55593de93f..369bc79edb9ea 100644 --- a/packages/cubejs-api-gateway/src/gateway.ts +++ b/packages/cubejs-api-gateway/src/gateway.ts @@ -1210,7 +1210,7 @@ class ApiGateway { ): Promise { let inQueue = false; let status: string = 'n/a'; - const queuedList = await orchestrator.getPreAggregationQueueStates(); + const queuedList = await orchestrator.getPreAggregationQueueStates(job.dataSource); queuedList.forEach((item) => { if ( item.queryHandler && diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts index 15fc50a0d4ed7..87d0825e0bdf7 100644 --- a/packages/cubejs-api-gateway/test/index.test.ts +++ b/packages/cubejs-api-gateway/test/index.test.ts @@ -18,6 +18,7 @@ import { AdapterApiMock } from './mocks'; import { ApiScopesTuple } from '../src/types/auth'; +import { PreAggJob } from '../src/types/request'; const logger = (type, message) => console.log({ type, ...message }); @@ -1184,6 +1185,57 @@ describe('API Gateway', () => { expect(res.status).toEqual(400); expect(res.body.error.includes('Cannot parse selector date range')).toBeTruthy(); }); + + // https://github.com/cube-js/cube/issues/11313 + test('job queue status is reported per job data source, not only for the default one', async () => { + // Constructed off the prototype (bypassing the constructor, which requires a full + // native SQL server) since getPreAggJobQueueStatus doesn't use instance state. + const apiGateway = Object.create(ApiGateway.prototype); + + const job: PreAggJob = { + request: 'request-id', + context: { securityContext: {} }, + preagg: 'orders_test.main', + table: 'orders_test_main', + target: 'orders_test_main_20200101', + structure: 'structure-version', + content: 'content-version', + updated: 1, + key: [], + status: 'scheduled', + timezone: 'UTC', + dataSource: 'test_ds', + }; + + // Mimics QueryOrchestrator#getPreAggregationQueueStates, which defaults + // its dataSource argument to 'default' and only returns queue entries + // that were scheduled on the requested data source's queue. + const orchestrator = { + getPreAggregationQueueStates: jest.fn((dataSource = 'default') => { + if (dataSource !== job.dataSource) { + return []; + } + return [{ + queryHandler: 'query', + query: { + requestId: job.request, + newVersionEntry: { + table_name: job.table, + structure_version: job.structure, + content_version: job.content, + last_updated_at: job.updated, + }, + }, + status: ['active'], + }]; + }), + }; + + const status = await (apiGateway as any).getPreAggJobQueueStatus(orchestrator, job); + + expect(orchestrator.getPreAggregationQueueStates).toHaveBeenCalledWith(job.dataSource); + expect(status).toEqual('processing'); + }); }); describe('healtchecks', () => { diff --git a/packages/cubejs-server-core/src/core/OrchestratorApi.ts b/packages/cubejs-server-core/src/core/OrchestratorApi.ts index 1d9fb163d7f05..12e1da161c4c9 100644 --- a/packages/cubejs-server-core/src/core/OrchestratorApi.ts +++ b/packages/cubejs-server-core/src/core/OrchestratorApi.ts @@ -289,8 +289,8 @@ export class OrchestratorApi { return this.orchestrator.checkPartitionsBuildRangeCache(queryBody); } - public async getPreAggregationQueueStates() { - return this.orchestrator.getPreAggregationQueueStates(); + public async getPreAggregationQueueStates(dataSource?: string) { + return this.orchestrator.getPreAggregationQueueStates(dataSource); } public async cancelPreAggregationQueriesFromQueue(queryKeys: string[], dataSource: string) { diff --git a/packages/cubejs-server-core/test/unit/OrchestratorApi.test.ts b/packages/cubejs-server-core/test/unit/OrchestratorApi.test.ts new file mode 100644 index 0000000000000..383bce2945256 --- /dev/null +++ b/packages/cubejs-server-core/test/unit/OrchestratorApi.test.ts @@ -0,0 +1,18 @@ +import { OrchestratorApi } from '../../src/core/OrchestratorApi'; + +describe('OrchestratorApi', () => { + // https://github.com/cube-js/cube/issues/11313 + test('getPreAggregationQueueStates forwards dataSource to the orchestrator', async () => { + const api = Object.create(OrchestratorApi.prototype); + const getPreAggregationQueueStates = jest.fn(async () => []); + api.orchestrator = { getPreAggregationQueueStates }; + + await api.getPreAggregationQueueStates('test_ds'); + expect(getPreAggregationQueueStates).toHaveBeenLastCalledWith('test_ds'); + + // QueryOrchestrator#getPreAggregationQueueStates defaults an undefined + // dataSource to 'default', so calls without one keep working. + await api.getPreAggregationQueueStates(); + expect(getPreAggregationQueueStates).toHaveBeenLastCalledWith(undefined); + }); +}); diff --git a/rust/cube-cli/src/main.rs b/rust/cube-cli/src/main.rs index 1f08936272037..7df4dd0b12b5d 100644 --- a/rust/cube-cli/src/main.rs +++ b/rust/cube-cli/src/main.rs @@ -21,7 +21,7 @@ struct Cli { #[command(flatten)] global: GlobalArgs, #[command(subcommand)] - command: Command, + command: Option, } #[derive(clap::Args, Clone)] @@ -255,21 +255,29 @@ async fn main() { unsafe { libc::signal(libc::SIGPIPE, libc::SIG_DFL); } - let cli = Cli::parse(); + let Cli { global, command } = Cli::parse(); commands::update::cleanup_stale_binary(); + // Bare `cube` with no subcommand: show the version, then the help text. + let Some(command) = command else { + println!("Cube CLI {}", env!("CUBE_CLI_VERSION")); + println!(); + let _ = cli_command().print_help(); + return; + }; + // Check for a newer release concurrently with the command; the notice // (if any) prints after the command output. Skipped for `update` itself // and `completion` (whose output is eval'd by shells). Completion also // skips telemetry — it runs on shell startup. - let is_completion = matches!(cli.command, Command::Completion(_)); - let check = match &cli.command { + let is_completion = matches!(command, Command::Completion(_)); + let check = match &command { Command::Update(_) | Command::Completion(_) => None, _ => Some(update::spawn_check()), }; - let command_name = cli.command.name(); + let command_name = command.name(); - let result = run(cli).await; + let result = run(global, command).await; if !is_completion { let mut props = serde_json::Map::new(); @@ -293,10 +301,10 @@ async fn main() { } } -async fn run(cli: Cli) -> Result<()> { - let mut ctx = Ctx::new(&cli.global)?; +async fn run(global: GlobalArgs, command: Command) -> Result<()> { + let mut ctx = Ctx::new(&global)?; use Command::*; - match cli.command { + match command { Login(args) => commands::login::command(args, &mut ctx).await, Logout(args) => commands::logout::command(args, &mut ctx).await, Whoami(args) => commands::whoami::command(args, &ctx).await,