@databricks/sql
Version:
Driver for connection to Databricks SQL via Thrift API.
90 lines • 4.01 kB
JavaScript
"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.FetchType = void 0;
const node_int64_1 = __importDefault(require("node-int64"));
const TCLIService_types_1 = require("../../thrift/TCLIService_types");
const Status_1 = __importDefault(require("../dto/Status"));
const utils_1 = require("./utils");
var FetchType;
(function (FetchType) {
FetchType[FetchType["Data"] = 0] = "Data";
FetchType[FetchType["Logs"] = 1] = "Logs";
})(FetchType = exports.FetchType || (exports.FetchType = {}));
function checkIfOperationHasMoreRows(response) {
var _a, _b, _c;
if (response.hasMoreRows) {
return true;
}
const columns = ((_a = response.results) === null || _a === void 0 ? void 0 : _a.columns) || [];
const columnValue = (0, utils_1.getColumnValue)(columns[0]);
return ((_c = (_b = columnValue === null || columnValue === void 0 ? void 0 : columnValue.values) === null || _b === void 0 ? void 0 : _b.length) !== null && _c !== void 0 ? _c : 0) > 0;
}
class RowSetProvider {
get hasMoreRows() {
var _a;
// `hasMoreRowsFlag` is populated only after fetching the first row set.
// Prior to that, we use a `operationHandle.hasResultSet` flag which
// is set if there are any data at all. Also, we have to choose appropriate
// flag in a getter because both `hasMoreRowsFlag` and `operationHandle.hasResultSet`
// may change between this getter calls
return (_a = this.hasMoreRowsFlag) !== null && _a !== void 0 ? _a : this.operationHandle.hasResultSet;
}
constructor(context, operationHandle, prefetchedResults, returnOnlyPrefetchedResults) {
this.fetchOrientation = TCLIService_types_1.TFetchOrientation.FETCH_FIRST;
this.prefetchedResults = [];
this.hasMoreRowsFlag = undefined;
this.context = context;
this.operationHandle = operationHandle;
prefetchedResults.forEach((item) => {
if (item) {
this.prefetchedResults.push(item);
}
});
this.returnOnlyPrefetchedResults = returnOnlyPrefetchedResults;
}
processFetchResponse(response) {
Status_1.default.assert(response.status);
this.fetchOrientation = TCLIService_types_1.TFetchOrientation.FETCH_NEXT;
this.hasMoreRowsFlag = checkIfOperationHasMoreRows(response);
return response.results;
}
async fetchNext({ limit }) {
const prefetchedResponse = this.prefetchedResults.shift();
if (prefetchedResponse) {
return this.processFetchResponse(prefetchedResponse);
}
// We end up here if no more prefetched results available (checked above)
if (this.returnOnlyPrefetchedResults) {
return undefined;
}
// Don't fetch next chunk if there are no more data available
if (!this.hasMoreRows) {
return undefined;
}
const driver = await this.context.getDriver();
const response = await driver.fetchResults({
operationHandle: this.operationHandle,
orientation: this.fetchOrientation,
maxRows: new node_int64_1.default(limit),
fetchType: FetchType.Data,
});
return this.processFetchResponse(response);
}
async hasMore() {
// If there are prefetched results available - return `true` regardless of
// the actual state of `hasMoreRows` flag (because we actually have some data)
if (this.prefetchedResults.length > 0) {
return true;
}
// We end up here if no more prefetched results available (checked above)
if (this.returnOnlyPrefetchedResults) {
return false;
}
return this.hasMoreRows;
}
}
exports.default = RowSetProvider;
//# sourceMappingURL=RowSetProvider.js.map