Skip to content

Commit 697b53f

Browse files
paveltiunovpaveltiunov
committed
feat: Report query usage for Athena and BigQuery
1 parent 20d3f98 commit 697b53f

File tree

6 files changed

+81
-40
lines changed

6 files changed

+81
-40
lines changed

‎packages/cubejs-athena-driver/driver/AthenaDriver.js‎

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@ const { promisify } = require('util');
33
const BaseDriver = require('@cubejs-backend/query-orchestrator/driver/BaseDriver');
44
const SqlString = require('sqlstring');
55

6-
const applyParams = (query, params) => {
7-
return SqlString.format(query, params);
8-
};
6+
const applyParams = (query, params) => SqlString.format(query, params);
97

108
class AthenaDriver extends BaseDriver {
119
constructor(config) {
@@ -34,7 +32,7 @@ class AthenaDriver extends BaseDriver {
3432
});
3533
}
3634

37-
async query(query, values) {
35+
async query(query, values, options) {
3836
const queryString = applyParams(
3937
query,
4038
(values || []).map(s => (typeof s === 'string' ? {
@@ -63,25 +61,26 @@ class AthenaDriver extends BaseDriver {
6361
) {
6462
const allRows = [];
6563
let columnInfo;
64+
this.reportQueryUsage({
65+
dataScannedInBytes: queryExecution.QueryExecution.Statistics.DataScannedInBytes
66+
}, options);
6667
for (
6768
let results = await this.athena.getQueryResultsAsync({ QueryExecutionId });
68-
!!results;
69+
results;
6970
results = results.NextToken && (await this.athena.getQueryResultsAsync({
7071
QueryExecutionId, NextToken: results.NextToken
7172
}))
7273
) {
7374
const [header, ...tableRows] = results.ResultSet.Rows;
7475
allRows.push(...(allRows.length ? results.ResultSet.Rows : tableRows));
7576
if (!columnInfo) {
76-
columnInfo = results.ResultSet.ResultSetMetadata.ColumnInfo
77+
columnInfo = results.ResultSet.ResultSetMetadata.ColumnInfo;
7778
}
7879
}
7980

80-
return allRows.map(r =>
81-
columnInfo
82-
.map((c, i) => ({ [c.Name]: r.Data[i].VarCharValue }))
83-
.reduce((a, b) => ({...a, ...b}), {})
84-
);
81+
return allRows.map(r => columnInfo
82+
.map((c, i) => ({ [c.Name]: r.Data[i].VarCharValue }))
83+
.reduce((a, b) => ({ ...a, ...b }), {}));
8584
}
8685
await this.sleep(500);
8786
}

‎packages/cubejs-bigquery-driver/driver/BigQueryDriver.js‎

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,16 @@ class BigQueryDriver extends BaseDriver {
4141
});
4242
}
4343

44-
query(query, values) {
45-
return this.bigquery.query({
44+
async query(query, values, options) {
45+
const data = await this.runQueryJob({
4646
query,
47-
params: values
48-
}).then(
49-
(data) => data[0] && data[0].map(
50-
row => R.map(value => (value && value.value && typeof value.value === 'string' ? value.value : value), row)
51-
)
47+
params: values,
48+
parameterMode: 'positional',
49+
useLegacySql: false
50+
}, options);
51+
52+
return data[0] && data[0].map(
53+
row => R.map(value => (value && value.value && typeof value.value === 'string' ? value.value : value), row)
5254
);
5355
}
5456

@@ -144,16 +146,21 @@ class BigQueryDriver extends BaseDriver {
144146
await this.bigquery.dataset(schemaName).get({ autoCreate: true });
145147
}
146148

147-
async loadPreAggregationIntoTable(preAggregationTableName, loadSql, params) {
149+
async loadPreAggregationIntoTable(preAggregationTableName, loadSql, params, options) {
148150
const [dataSet, tableName] = preAggregationTableName.split('.');
149-
const [job] = await this.bigquery.createQueryJob({
151+
const bigQueryQuery = {
150152
query: loadSql,
151153
params,
152154
parameterMode: 'positional',
153155
destination: this.bigquery.dataset(dataSet).table(tableName),
154156
createDisposition: "CREATE_IF_NEEDED",
155157
useLegacySql: false
156-
});
158+
};
159+
return this.runQueryJob(bigQueryQuery, options, false);
160+
}
161+
162+
async runQueryJob(bigQueryQuery, options, withResults = true) {
163+
const [job] = await this.bigquery.createQueryJob(bigQueryQuery);
157164
const awaitForJobStatus = async () => {
158165
const [result] = await job.getMetadata();
159166
if (result.status && result.status.state === 'DONE') {
@@ -164,17 +171,20 @@ class BigQueryDriver extends BaseDriver {
164171
JSON.stringify(result.status.errorResult)
165172
);
166173
}
174+
this.reportQueryUsage(result.statistics, options);
167175
} else {
168-
return false;
176+
return null;
169177
}
170-
return true;
178+
return withResults ? job.getQueryResults() : true;
171179
};
172-
for (let i = 0; i < 10 * 60 / 5; i++) {
173-
if (await awaitForJobStatus()) {
174-
break;
180+
for (let i = 0; i < 15 * 60 / 5; i++) {
181+
const result = await awaitForJobStatus();
182+
if (result) {
183+
return result;
175184
}
176-
await pause(5000);
185+
await pause(Math.min(5000, 200 * i));
177186
}
187+
throw new Error('BigQuery job timeout');
178188
}
179189

180190
quoteIdentifier(identifier) {

‎packages/cubejs-query-orchestrator/driver/BaseDriver.js‎

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,12 +93,12 @@ class BaseDriver {
9393
);
9494
}
9595

96-
loadPreAggregationIntoTable(preAggregationTableName, loadSql, params, tx) {
97-
return this.query(loadSql, params, tx);
96+
loadPreAggregationIntoTable(preAggregationTableName, loadSql, params, options) {
97+
return this.query(loadSql, params, options);
9898
}
9999

100-
dropTable(tableName, tx) {
101-
return this.query(`DROP TABLE ${tableName}`, [], tx);
100+
dropTable(tableName, options) {
101+
return this.query(`DROP TABLE ${tableName}`, [], options);
102102
}
103103

104104
param(/* paramIndex */) {
@@ -180,6 +180,19 @@ class BaseDriver {
180180
cancelCombinator(fn) {
181181
return cancelCombinator(fn);
182182
}
183+
184+
setLogger(logger) {
185+
this.logger = logger;
186+
}
187+
188+
reportQueryUsage(usage, queryOptions) {
189+
if (this.logger) {
190+
this.logger('SQL Query Usage', {
191+
...usage,
192+
...queryOptions
193+
});
194+
}
195+
}
183196
}
184197

185198
module.exports = BaseDriver;

‎packages/cubejs-query-orchestrator/orchestrator/PreAggregations.js‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -411,14 +411,21 @@ class PreAggregationLoader {
411411
}
412412

413413
logExecutingSql(invalidationKeys, query, params, targetTableName, newVersionEntry) {
414-
this.logger('Executing Load Pre Aggregation SQL', {
414+
this.logger(
415+
'Executing Load Pre Aggregation SQL',
416+
this.queryOptions(invalidationKeys, query, params, targetTableName, newVersionEntry)
417+
);
418+
}
419+
420+
queryOptions(invalidationKeys, query, params, targetTableName, newVersionEntry) {
421+
return {
415422
queryKey: this.preAggregationQueryKey(invalidationKeys),
416423
query,
417424
values: params,
418425
targetTableName,
419426
requestId: this.requestId,
420427
newVersionEntry,
421-
});
428+
};
422429
}
423430

424431
async refreshImplStoreInSourceStrategy(client, newVersionEntry, saveCancelFn, invalidationKeys) {
@@ -434,7 +441,8 @@ class PreAggregationLoader {
434441
await saveCancelFn(client.loadPreAggregationIntoTable(
435442
targetTableName,
436443
query,
437-
params
444+
params,
445+
this.queryOptions(invalidationKeys, query, params, targetTableName, newVersionEntry)
438446
));
439447
await this.createIndexes(client, newVersionEntry, saveCancelFn);
440448
await this.loadCache.reset(this.preAggregation);
@@ -456,7 +464,8 @@ class PreAggregationLoader {
456464
await saveCancelFn(client.loadPreAggregationIntoTable(
457465
targetTableName,
458466
query,
459-
params
467+
params,
468+
this.queryOptions(invalidationKeys, query, params, targetTableName, newVersionEntry)
460469
));
461470
const tableData = await this.downloadTempExternalPreAggregation(client, newVersionEntry, saveCancelFn);
462471
await this.uploadExternalPreAggregation(tableData, newVersionEntry, saveCancelFn);
@@ -476,7 +485,11 @@ class PreAggregationLoader {
476485
preAggregation: this.preAggregation,
477486
requestId: this.requestId
478487
});
479-
const tableData = await saveCancelFn(client.downloadQueryResults(sql, params));
488+
const tableData = await saveCancelFn(client.downloadQueryResults(
489+
sql,
490+
params,
491+
this.queryOptions(invalidationKeys, sql, params, this.targetTableName(newVersionEntry), newVersionEntry)
492+
));
480493
await this.uploadExternalPreAggregation(tableData, newVersionEntry, saveCancelFn);
481494
await this.loadCache.reset(this.preAggregation);
482495
}

‎packages/cubejs-query-orchestrator/orchestrator/QueryCache.js‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ class QueryCache {
148148
this.logger('Executing SQL', {
149149
...q
150150
});
151-
return client.query(q.query, q.values);
151+
return client.query(q.query, q.values, q);
152152
}, {
153153
logger: this.logger,
154154
cacheAndQueueDriver: this.options.cacheAndQueueDriver,
@@ -169,7 +169,7 @@ class QueryCache {
169169
this.logger('Executing SQL', {
170170
...q
171171
});
172-
return client.query(q.query, q.values);
172+
return client.query(q.query, q.values, q);
173173
},
174174
{
175175
logger: this.logger,

‎packages/cubejs-server-core/core/index.js‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ const devLogger = (level) => (type, { error, warning, ...message }) => {
8585
`${withColor(type, colors.yellow)}: ${format({ ...message, allSqlLines: true, showRestParams: true })} \n${withColor(warning, colors.yellow)}`
8686
);
8787
const logError = () => console.log(`${withColor(type, colors.red)}: ${format({ ...message, allSqlLines: true, showRestParams: true })} \n${error}`);
88-
const logDetails = () => console.log(`${withColor(type)}: ${format(message)}`);
88+
const logDetails = (showRestParams) => console.log(`${withColor(type)}: ${format({ ...message, showRestParams })}`);
8989

9090
if (error) {
9191
logError();
@@ -96,7 +96,7 @@ const devLogger = (level) => (type, { error, warning, ...message }) => {
9696
switch ((level || 'info').toLowerCase()) {
9797
case "trace": {
9898
if (!error && !warning) {
99-
logDetails();
99+
logDetails(true);
100100
break;
101101
}
102102
}
@@ -462,6 +462,9 @@ class CubejsServerCore {
462462
getDriver: async () => {
463463
if (!driverPromise) {
464464
const driver = await this.driverFactory(context);
465+
if (driver.setLogger) {
466+
driver.setLogger(this.logger);
467+
}
465468
driverPromise = driver.testConnection().then(() => driver).catch(e => {
466469
driverPromise = null;
467470
throw e;
@@ -472,6 +475,9 @@ class CubejsServerCore {
472475
getExternalDriverFactory: this.externalDriverFactory && (async () => {
473476
if (!externalPreAggregationsDriverPromise) {
474477
const driver = await this.externalDriverFactory(context);
478+
if (driver.setLogger) {
479+
driver.setLogger(this.logger);
480+
}
475481
externalPreAggregationsDriverPromise = driver.testConnection().then(() => driver).catch(e => {
476482
externalPreAggregationsDriverPromise = null;
477483
throw e;

0 commit comments

Comments
 (0)