556 lines
17 KiB
TypeScript
556 lines
17 KiB
TypeScript
import { RuntimeError } from '@botpress/sdk'
|
|
import { Readable, Stream } from 'stream'
|
|
import { v4 as uuidv4 } from 'uuid'
|
|
import { getAuthenticatedGoogleClient } from './auth'
|
|
import { handleNotFoundError, handleRateLimitError, isGaxiosError } from './error-handling'
|
|
import { serializeToken } from './file-notification-token'
|
|
import { FilesCache } from './files-cache'
|
|
import { APP_GOOGLE_FOLDER_MIMETYPE, APP_GOOGLE_SHORTCUT_MIMETYPE, INDEXABLE_MIMETYPES } from './mime-types'
|
|
import {
|
|
BaseDiscriminatedFile,
|
|
GoogleDriveClient,
|
|
BaseNormalFile,
|
|
File,
|
|
BaseFolderFile,
|
|
Folder,
|
|
ListFilesOutput,
|
|
ListFoldersOutput,
|
|
CreateFileArgs,
|
|
UpdateFileArgs,
|
|
ListItemsInput,
|
|
ListItemsOutput,
|
|
BaseGenericFileUnion,
|
|
FileChannel,
|
|
GenericFile,
|
|
} from './types'
|
|
import { listItemsAndProcess, ListFunction, streamToBuffer, ListItemsInputWithArgs, listAllItems } from './utils'
|
|
import {
|
|
getFileTypeFromMimeType,
|
|
parseChannel,
|
|
parseBaseGeneric,
|
|
parseBaseGenerics,
|
|
parseBaseNormal,
|
|
} from './validation'
|
|
import * as bp from '.botpress'
|
|
|
|
type DownloadFileDataClientOutput = {
|
|
mimeType: string
|
|
dataSize: number
|
|
} & (
|
|
| {
|
|
dataType: 'buffer'
|
|
data: Buffer
|
|
}
|
|
| {
|
|
dataType: 'stream'
|
|
data: Readable
|
|
}
|
|
)
|
|
type TryWatchAllOutput = {
|
|
fileChannels: FileChannel[]
|
|
hasError: boolean
|
|
}
|
|
|
|
const MAX_RESOURCE_WATCH_EXPIRATION_DELAY_MS = 86400 * 1000 // 24 hours
|
|
const MAX_EXPORT_FILE_SIZE_BYTES = 10000000 // 10MB, as per the Google Drive API doc
|
|
const MYDRIVE_ID_ALIAS = 'root'
|
|
const PAGE_SIZE = 100
|
|
const GOOGLE_API_EXPORTFORMATS_FIELDS = 'exportFormats'
|
|
const GOOGLE_API_FILE_FIELDS =
|
|
'id, name, mimeType, parents, size, sha256Checksum, md5Checksum, version, trashed, modifiedTime, driveId, sharedWithMeTime'
|
|
const GOOGLE_API_FILELIST_FIELDS = `files(${GOOGLE_API_FILE_FIELDS}), nextPageToken`
|
|
|
|
const INCLUDE_FILES_FROM_ALL_DRIVES = {
|
|
includeItemsFromAllDrives: true,
|
|
supportsAllDrives: true,
|
|
} as const
|
|
|
|
export class Client {
|
|
private constructor(
|
|
private _ctx: bp.Context,
|
|
private _googleClient: GoogleDriveClient,
|
|
private _filesCache: FilesCache,
|
|
private _logger: bp.Logger
|
|
) {}
|
|
|
|
public static async create({
|
|
client,
|
|
ctx,
|
|
logger,
|
|
}: {
|
|
client: bp.Client
|
|
ctx: bp.Context
|
|
logger: bp.Logger
|
|
}): Promise<Client> {
|
|
const googleClient = await getAuthenticatedGoogleClient({
|
|
client,
|
|
ctx,
|
|
})
|
|
const filesCache = new FilesCache(client, ctx)
|
|
return new Client(ctx, googleClient, filesCache, logger)
|
|
}
|
|
|
|
public setCache(filesCache: FilesCache) {
|
|
this._filesCache = filesCache
|
|
}
|
|
|
|
public async getRootFolderId(): Promise<string> {
|
|
try {
|
|
const response = await this._googleClient.files.get({ fileId: MYDRIVE_ID_ALIAS })
|
|
return response.data.id!
|
|
} catch (thrown: unknown) {
|
|
if (isGaxiosError(thrown) && thrown.toString().includes('File not found: ')) {
|
|
return thrown.toString().split('File not found: ')[1]!.slice(0, -1)
|
|
}
|
|
throw thrown
|
|
}
|
|
}
|
|
|
|
public async listFiles({ nextToken }: ListItemsInput): Promise<ListFilesOutput> {
|
|
const { items: baseFiles, meta } = await this._listBaseNormalFiles({ nextToken })
|
|
const completeFilesPromises = baseFiles.map((f) => this._getCompleteFileFromBaseFile(f))
|
|
const items = await Promise.all(completeFilesPromises)
|
|
return {
|
|
items,
|
|
meta,
|
|
}
|
|
}
|
|
|
|
private async _listBaseNormalFiles({ nextToken }: ListItemsInput): Promise<ListItemsOutput<BaseNormalFile>> {
|
|
const {
|
|
items: newFiles,
|
|
meta: { nextToken: newNextToken },
|
|
} = await this._listBaseGenericFiles({
|
|
nextToken,
|
|
args: {
|
|
searchQuery: `mimeType != '${APP_GOOGLE_FOLDER_MIMETYPE}' and mimeType != '${APP_GOOGLE_SHORTCUT_MIMETYPE}'`,
|
|
},
|
|
})
|
|
|
|
const items = newFiles.filter((f) => f.type === 'normal')
|
|
return {
|
|
items,
|
|
meta: {
|
|
nextToken: newNextToken,
|
|
},
|
|
}
|
|
}
|
|
|
|
public async listFolders({ nextToken }: ListItemsInput): Promise<ListFoldersOutput> {
|
|
const { items: baseFolders, meta } = await this._listBaseFolderFiles({ nextToken })
|
|
|
|
const completeFoldersPromises = baseFolders.map((f) => this._getCompleteFolderFromBaseFolder(f))
|
|
const items = await Promise.all(completeFoldersPromises)
|
|
return {
|
|
items,
|
|
meta,
|
|
}
|
|
}
|
|
|
|
private async _listBaseFolderFiles({ nextToken }: ListItemsInput): Promise<ListItemsOutput<BaseFolderFile>> {
|
|
const {
|
|
items: newFiles,
|
|
meta: { nextToken: newNextToken },
|
|
} = await this._listBaseGenericFiles({
|
|
nextToken,
|
|
args: {
|
|
searchQuery: `mimeType = '${APP_GOOGLE_FOLDER_MIMETYPE}'`,
|
|
},
|
|
})
|
|
if (nextToken === undefined) {
|
|
// My Drive is not returned by list operation but needs to be part of list, so we add it to first page
|
|
const myDriveFile = await this._fetchFile(MYDRIVE_ID_ALIAS)
|
|
newFiles.push(myDriveFile)
|
|
}
|
|
const items = newFiles.filter((f) => f.type === 'folder')
|
|
return {
|
|
items,
|
|
meta: {
|
|
nextToken: newNextToken,
|
|
},
|
|
}
|
|
}
|
|
|
|
public async getChildren(folderId: string): Promise<GenericFile[]> {
|
|
const files = await listAllItems(this._listBaseGenericFiles.bind(this), {
|
|
searchQuery: this._getParentsFilter(folderId),
|
|
})
|
|
return await Promise.all(files.map((f) => this._getCompleteFile(f)))
|
|
}
|
|
|
|
public async getChildrenSubset({
|
|
folderId,
|
|
extraQuery,
|
|
nextToken,
|
|
}: {
|
|
folderId: string
|
|
extraQuery?: string
|
|
nextToken?: string
|
|
}) {
|
|
const searchQuery = this._getParentsFilter(folderId) + (extraQuery ? ` and ${extraQuery}` : '')
|
|
const listResponse = await this._googleClient.files.list({
|
|
corpora: 'user',
|
|
fields: GOOGLE_API_FILELIST_FIELDS,
|
|
q: `${searchQuery} and trashed != true`,
|
|
pageToken: nextToken,
|
|
pageSize: PAGE_SIZE,
|
|
spaces: 'drive',
|
|
...INCLUDE_FILES_FROM_ALL_DRIVES,
|
|
})
|
|
return { files: listResponse.data.files, nextToken: listResponse.data.nextPageToken ?? undefined }
|
|
}
|
|
|
|
private _getParentsFilter(parentId: string): string {
|
|
return parentId === MYDRIVE_ID_ALIAS ? 'not trashed' : `'${parentId}' in parents`
|
|
}
|
|
|
|
public async createFile({ name, parentId, mimeType }: CreateFileArgs): Promise<File> {
|
|
const response = await this._googleClient.files.create({
|
|
fields: GOOGLE_API_FILE_FIELDS,
|
|
requestBody: {
|
|
name,
|
|
parents: parentId ? [parentId] : undefined,
|
|
mimeType,
|
|
},
|
|
})
|
|
const file = parseBaseNormal(response.data)
|
|
this._filesCache.set({ type: 'normal', ...file })
|
|
return await this._getCompleteFileFromBaseFile(file)
|
|
}
|
|
|
|
public async readGenericFile(id: string): Promise<GenericFile> {
|
|
const file = await this._fetchFile(id)
|
|
return await this._getCompleteFile(file)
|
|
}
|
|
|
|
public async readFile(id: string): Promise<File> {
|
|
const file = await this._fetchFile(id)
|
|
if (file.type !== 'normal') {
|
|
throw new RuntimeError(`Attempted to read a file of type ${file.type}`)
|
|
}
|
|
return await this._getCompleteFileFromBaseFile(file)
|
|
}
|
|
|
|
public async updateFile({ id: fileId, name, parentId }: UpdateFileArgs): Promise<File> {
|
|
const addParents = parentId ? `${parentId}` : undefined
|
|
const response = await this._googleClient.files.update({
|
|
fields: GOOGLE_API_FILE_FIELDS,
|
|
fileId,
|
|
addParents, // Also removes old parents
|
|
requestBody: {
|
|
name,
|
|
},
|
|
})
|
|
const file = parseBaseNormal(response.data)
|
|
this._filesCache.set({ type: 'normal', ...file })
|
|
return await this._getCompleteFileFromBaseFile(file)
|
|
}
|
|
|
|
public async deleteFile(id: string) {
|
|
await this._googleClient.files.delete({
|
|
fileId: id,
|
|
})
|
|
}
|
|
|
|
public async uploadFileData({ id, mimeType, data }: { id: string; mimeType?: string; data: Stream }) {
|
|
await this._googleClient.files.update({
|
|
fileId: id,
|
|
media: {
|
|
body: data,
|
|
mimeType,
|
|
},
|
|
})
|
|
}
|
|
|
|
public async downloadFileData({ id }: { id: string }): Promise<DownloadFileDataClientOutput> {
|
|
const file = await this._fetchFile(id)
|
|
if (file.type !== 'normal') {
|
|
throw new RuntimeError(`Attempted to download a file of type ${file.type}`)
|
|
}
|
|
|
|
const exportType = await this._findExportType(file.mimeType)
|
|
let output: DownloadFileDataClientOutput
|
|
if (exportType) {
|
|
// File size is unknown when exporting, download all data to buffer to know size
|
|
const fileDownloadStream = await this._exportFileData(file, exportType)
|
|
const buffer = await streamToBuffer(fileDownloadStream, MAX_EXPORT_FILE_SIZE_BYTES)
|
|
output = {
|
|
mimeType: exportType,
|
|
dataSize: buffer.length,
|
|
dataType: 'buffer',
|
|
data: buffer,
|
|
}
|
|
} else {
|
|
output = {
|
|
mimeType: file.mimeType,
|
|
dataSize: file.size,
|
|
dataType: 'stream',
|
|
data: await this._fetchFileData(file),
|
|
}
|
|
}
|
|
return output
|
|
}
|
|
|
|
private _getRateLimitErrorHandler(): (error: unknown) => Promise<undefined> {
|
|
return async (error: unknown) => {
|
|
return handleRateLimitError(error, this._logger)
|
|
}
|
|
}
|
|
|
|
private _getNotFoundErrorHandler(): (error: unknown) => Promise<undefined> {
|
|
return async (error: unknown) => {
|
|
return handleNotFoundError(error, this._logger)
|
|
}
|
|
}
|
|
|
|
private async _tryWatchAllListableGenericFiles<T extends BaseGenericFileUnion>(
|
|
listFn: ListFunction<T>
|
|
): Promise<TryWatchAllOutput> {
|
|
const fileChannels: FileChannel[] = []
|
|
let hasError = false
|
|
await listItemsAndProcess(listFn, async (item) => {
|
|
const channel = await this._watch(item).catch(this._getRateLimitErrorHandler())
|
|
if (channel) {
|
|
fileChannels.push(channel)
|
|
} else {
|
|
hasError = true
|
|
}
|
|
})
|
|
return {
|
|
fileChannels,
|
|
hasError,
|
|
}
|
|
}
|
|
|
|
public async watch(id: string): Promise<FileChannel> {
|
|
const file = await this._fetchFile(id)
|
|
return await this._watch(file)
|
|
}
|
|
|
|
/**
|
|
* @returns Channel if successful, undefined if the subscription rate limit is exceeded
|
|
*/
|
|
public async tryWatch(id: string): Promise<FileChannel | undefined> {
|
|
return await this.watch(id).catch(this._getRateLimitErrorHandler())
|
|
}
|
|
|
|
private async _watch(file: BaseGenericFileUnion): Promise<FileChannel> {
|
|
const absoluteExpirationTimeMs: number = Date.now() + MAX_RESOURCE_WATCH_EXPIRATION_DELAY_MS
|
|
const { id: fileId, mimeType } = file
|
|
const token = serializeToken(
|
|
{
|
|
fileId,
|
|
fileType: getFileTypeFromMimeType(mimeType),
|
|
},
|
|
bp.secrets.WEBHOOK_SECRET
|
|
)
|
|
const response = await this._googleClient.files.watch({
|
|
fileId,
|
|
requestBody: {
|
|
id: uuidv4(),
|
|
type: 'web_hook',
|
|
address: `${process.env.BP_WEBHOOK_URL}/${this._ctx.webhookId}`,
|
|
token,
|
|
expiration: absoluteExpirationTimeMs.toString(),
|
|
},
|
|
})
|
|
const baseChannel = parseChannel(response.data)
|
|
this._logger.forBot().debug(`Watching file '${file.name}' (${file.id}): channel ID = ${baseChannel.id}`)
|
|
return {
|
|
...baseChannel,
|
|
fileId,
|
|
}
|
|
}
|
|
|
|
public async tryWatchAllFiles(): Promise<TryWatchAllOutput> {
|
|
return await this._tryWatchAllListableGenericFiles(this._listBaseNormalFiles.bind(this))
|
|
}
|
|
|
|
public async tryWatchAllFolders(): Promise<TryWatchAllOutput> {
|
|
return await this._tryWatchAllListableGenericFiles(this._listBaseFolderFiles.bind(this))
|
|
}
|
|
|
|
public async tryWatchAll(): Promise<TryWatchAllOutput> {
|
|
const [filesResult, foldersResult] = await Promise.all([this.tryWatchAllFiles(), this.tryWatchAllFolders()])
|
|
return {
|
|
fileChannels: [...filesResult.fileChannels, ...foldersResult.fileChannels],
|
|
hasError: filesResult.hasError || foldersResult.hasError,
|
|
}
|
|
}
|
|
|
|
public async unwatch(channels: FileChannel | FileChannel[]) {
|
|
if (!Array.isArray(channels)) {
|
|
channels = [channels]
|
|
}
|
|
const unwatchPromises = channels.map((channel) => {
|
|
const fileName = this._filesCache.find(channel.fileId)?.name ?? '[unknown]'
|
|
this._logger.forBot().debug(`Unwatching file ${fileName} (${channel.fileId}) with channel ID = ${channel.id}`)
|
|
const { id, resourceId } = channel
|
|
return this._googleClient.channels.stop({
|
|
requestBody: {
|
|
id,
|
|
resourceId,
|
|
},
|
|
})
|
|
})
|
|
await Promise.all(unwatchPromises)
|
|
}
|
|
|
|
public async tryUnwatch(channels: FileChannel | FileChannel[]) {
|
|
await this.unwatch(channels).catch(this._getNotFoundErrorHandler())
|
|
}
|
|
|
|
/**
|
|
* Removes internal fields and adds computed attributes
|
|
*/
|
|
private async _getCompleteFileFromBaseFile(file: BaseNormalFile): Promise<File> {
|
|
return {
|
|
...file,
|
|
path: await this._getFilePath({ type: 'normal', ...file }),
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Removes internal fields and adds computed attributes
|
|
*/
|
|
private async _getCompleteFolderFromBaseFolder(file: BaseFolderFile): Promise<Folder> {
|
|
const { id, mimeType, name, parentId } = file
|
|
return {
|
|
id,
|
|
mimeType,
|
|
name,
|
|
parentId,
|
|
path: await this._getFilePath({ type: 'folder', ...file }),
|
|
}
|
|
}
|
|
|
|
private async _getCompleteFile(file: BaseDiscriminatedFile): Promise<GenericFile> {
|
|
return {
|
|
...file,
|
|
path: await this._getFilePath(file),
|
|
}
|
|
}
|
|
|
|
private async _listBaseGenericFiles({
|
|
nextToken,
|
|
args,
|
|
}: ListItemsInputWithArgs<{ searchQuery?: string }>): Promise<ListItemsOutput<BaseDiscriminatedFile>> {
|
|
const searchQuery = args?.searchQuery
|
|
const listResponse = await this._googleClient.files.list({
|
|
corpora: 'user',
|
|
fields: GOOGLE_API_FILELIST_FIELDS,
|
|
q: (searchQuery ?? '') + (searchQuery?.length ? ' and ' : '') + 'trashed != true',
|
|
pageToken: nextToken,
|
|
pageSize: PAGE_SIZE,
|
|
spaces: 'drive',
|
|
...INCLUDE_FILES_FROM_ALL_DRIVES,
|
|
})
|
|
|
|
const newNextToken = listResponse.data.nextPageToken ?? undefined
|
|
const unvalidatedDriveFiles = listResponse.data.files
|
|
if (!unvalidatedDriveFiles) {
|
|
throw new RuntimeError('No files were returned by the API')
|
|
}
|
|
const newFiles = parseBaseGenerics(unvalidatedDriveFiles)
|
|
for (const newFile of newFiles) {
|
|
this._filesCache.set(newFile)
|
|
}
|
|
|
|
return {
|
|
items: newFiles,
|
|
meta: {
|
|
nextToken: newNextToken,
|
|
},
|
|
}
|
|
}
|
|
|
|
private async _getOrFetchFile(id: string): Promise<BaseDiscriminatedFile> {
|
|
let file = this._filesCache.find(id)
|
|
if (!file) {
|
|
file = await this._fetchFile(id)
|
|
}
|
|
return file
|
|
}
|
|
|
|
private async _fetchFile(id: string): Promise<BaseDiscriminatedFile> {
|
|
const response = await this._googleClient.files.get({
|
|
fileId: id,
|
|
fields: GOOGLE_API_FILE_FIELDS,
|
|
...INCLUDE_FILES_FROM_ALL_DRIVES,
|
|
})
|
|
const file = parseBaseGeneric(response.data)
|
|
this._filesCache.set(file)
|
|
return file
|
|
}
|
|
|
|
private async _fetchFileData({ id: fileId }: BaseNormalFile): Promise<Readable> {
|
|
const fileDownloadResponse = await this._googleClient.files.get(
|
|
{
|
|
fileId,
|
|
alt: 'media',
|
|
...INCLUDE_FILES_FROM_ALL_DRIVES,
|
|
},
|
|
{
|
|
responseType: 'stream',
|
|
}
|
|
)
|
|
return fileDownloadResponse.data
|
|
}
|
|
|
|
private async _exportFileData({ id: fileId }: BaseNormalFile, mimeType: string): Promise<Readable> {
|
|
const fileExportResponse = await this._googleClient.files.export(
|
|
{
|
|
fileId,
|
|
mimeType,
|
|
...INCLUDE_FILES_FROM_ALL_DRIVES,
|
|
},
|
|
{
|
|
responseType: 'stream',
|
|
}
|
|
)
|
|
return fileExportResponse.data
|
|
}
|
|
|
|
private async _fetchExportFormatMap(): Promise<Record<string, string[]>> {
|
|
const response = await this._googleClient.about.get({
|
|
fields: GOOGLE_API_EXPORTFORMATS_FIELDS,
|
|
})
|
|
const { exportFormats } = response.data
|
|
if (!exportFormats) {
|
|
throw new RuntimeError('Export formats are missing in Schema$About from the API response')
|
|
}
|
|
return exportFormats
|
|
}
|
|
|
|
/**
|
|
* @returns The export type to use, or undefined if the file cannot be exported
|
|
*/
|
|
private async _findExportType(originalContentType: string): Promise<string | undefined> {
|
|
const exportFormatMap = await this._fetchExportFormatMap()
|
|
const exportContentTypes = exportFormatMap[originalContentType]
|
|
if (!exportContentTypes) {
|
|
return undefined
|
|
}
|
|
|
|
const indexableContentType = INDEXABLE_MIMETYPES.find((type) => exportContentTypes.includes(type))
|
|
const defaultContentType = exportContentTypes[0]
|
|
return indexableContentType ?? defaultContentType
|
|
}
|
|
|
|
private _getFilePath = async (file: BaseDiscriminatedFile, pathAcc?: string[]): Promise<string[]> => {
|
|
const path = [file.name, ...(pathAcc ?? [])]
|
|
|
|
if (!file.parentId) {
|
|
return path
|
|
}
|
|
|
|
try {
|
|
const parent = await this._getOrFetchFile(file.parentId)
|
|
|
|
return await this._getFilePath(parent, path)
|
|
} catch {
|
|
return path
|
|
}
|
|
}
|
|
}
|