license import '../flags' ; import { env } from '../environment' ; import { getModelArtifactsInfoForJSON } from './io_utils' ; import { IORouter , IORouterRegistry } from './router_registry' ; import { IOHandler , ModelArtifacts , ModelArtifactsInfo , ModelStoreManager , SaveResult } from './types' ; const DATABASE_NAME = 'tensorflowjs' ; const DATABASE_VERSION = 1 ; const MODEL_STORE_NAME = 'models_store' ; const INFO_STORE_NAME = 'model_info_store' ; /**
 * Delete the entire database for tensorflow.js, including the models store.
 */ export async function deleteDatabase ( ) { const idbFactory = getIndexedDBFactory ( ) ; return new Promise < void > ( ( resolve , reject ) => { const deleteRequest = idbFactory . deleteDatabase ( DATABASE_NAME ) ; deleteRequest . onsuccess = ( ) => resolve ( ) ; deleteRequest . onerror = error => reject ( error ) ; } ) ; } function getIndexedDBFactory ( ) { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) ) { throw new Error ( 'Failed to obtain IndexedDB factory because the current environment' + 'is not a web browser.' ) ; } const theWindow = typeof window === 'undefined' ? self : window ; const factory = theWindow . indexedDB || theWindow . mozIndexedDB || theWindow . webkitIndexedDB || theWindow . msIndexedDB || theWindow . shimIndexedDB ; if ( factory == null ) { throw new Error ( 'The current browser does not appear to support IndexedDB.' ) ; } return factory ; } function setUpDatabase ( openRequest ) { const db = openRequest . result as IDBDatabase ; db . createObjectStore ( MODEL_STORE_NAME , { keyPath : 'modelPath' } ) ; db . createObjectStore ( INFO_STORE_NAME , { keyPath : 'modelPath' } ) ; } /**
 * IOHandler subclass: Browser IndexedDB.
 *
 * See the doc string of `browserIndexedDB` for more details.
 */ export class BrowserIndexedDB implements IOHandler { protected readonly indexedDB : IDBFactory ; protected readonly modelPath : string ; static readonly URL_SCHEME = 'indexeddb://' ; constructor ( modelPath ) { this . indexedDB = getIndexedDBFactory ( ) ; if ( modelPath == null || ! modelPath ) { throw new Error ( 'For IndexedDB, modelPath must not be null, undefined or empty.' ) ; } this . modelPath = modelPath ; } async save ( modelArtifacts ) { if ( modelArtifacts . modelTopology instanceof ArrayBuffer ) { throw new Error ( 'BrowserLocalStorage.save() does not support saving model topology ' + 'in binary formats yet.' ) ; } return this . databaseAction ( this . modelPath , modelArtifacts ) as Promise < SaveResult > ; } async load ( ) { return this . databaseAction ( this . modelPath ) as Promise < ModelArtifacts > ; } param modelPath param modelArtifacts returns private databaseAction ( modelPath , modelArtifacts ? ) { return new Promise < ModelArtifacts | SaveResult > ( ( resolve , reject ) => { const openRequest = this . indexedDB . open ( DATABASE_NAME , DATABASE_VERSION ) ; openRequest . onupgradeneeded = ( ) => setUpDatabase ( openRequest ) ; openRequest . onsuccess = ( ) => { const db = openRequest . result ; if ( modelArtifacts == null ) { const modelTx = db . transaction ( MODEL_STORE_NAME , 'readonly' ) ; const modelStore = modelTx . objectStore ( MODEL_STORE_NAME ) ; const getRequest = modelStore . get ( this . modelPath ) ; getRequest . onsuccess = ( ) => { if ( getRequest . result == null ) { db . close ( ) ; return reject ( new Error ( `Cannot find model with path '${ this . modelPath }' ` + `in IndexedDB.` ) ) ; } else { resolve ( getRequest . result . modelArtifacts ) ; } } ; getRequest . onerror = error => { db . close ( ) ; return reject ( getRequest . error ) ; } ; modelTx . oncomplete = ( ) => db . close ( ) ; } else { const modelArtifactsInfo = getModelArtifactsInfoForJSON ( modelArtifacts ) ; const infoTx = db . transaction ( INFO_STORE_NAME , 'readwrite' ) ; let infoStore = infoTx . objectStore ( INFO_STORE_NAME ) ; const putInfoRequest = infoStore . put ( { modelPath : this . modelPath , modelArtifactsInfo } ) ; let modelTx ; putInfoRequest . onsuccess = ( ) => { modelTx = db . transaction ( MODEL_STORE_NAME , 'readwrite' ) ; const modelStore = modelTx . objectStore ( MODEL_STORE_NAME ) ; const putModelRequest = modelStore . put ( { modelPath : this . modelPath , modelArtifacts , modelArtifactsInfo } ) ; putModelRequest . onsuccess = ( ) => resolve ( { modelArtifactsInfo } ) ; putModelRequest . onerror = error => { infoStore = infoTx . objectStore ( INFO_STORE_NAME ) ; const deleteInfoRequest = infoStore . delete ( this . modelPath ) ; deleteInfoRequest . onsuccess = ( ) => { db . close ( ) ; return reject ( putModelRequest . error ) ; } ; deleteInfoRequest . onerror = error => { db . close ( ) ; return reject ( putModelRequest . error ) ; } ; } ; } ; putInfoRequest . onerror = error => { db . close ( ) ; return reject ( putInfoRequest . error ) ; } ; infoTx . oncomplete = ( ) => { if ( modelTx == null ) { db . close ( ) ; } else { modelTx . oncomplete = ( ) => db . close ( ) ; } } ; } } ; openRequest . onerror = error => reject ( openRequest . error ) ; } ) ; } } export const indexedDBRouter = ( url ) => { if ( ! env ( ) . getBool ( 'IS_BROWSER' ) ) { return null ; } else { if ( ! Array . isArray ( url ) && url . startsWith ( BrowserIndexedDB . URL_SCHEME ) ) { return browserIndexedDB ( url . slice ( BrowserIndexedDB . URL_SCHEME . length ) ) ; } else { return null ; } } } ; IORouterRegistry . registerSaveRouter ( indexedDBRouter ) ; IORouterRegistry . registerLoadRouter ( indexedDBRouter ) ; param modelPath returns export function browserIndexedDB ( modelPath ) { return new BrowserIndexedDB ( modelPath ) ; } function maybeStripScheme ( key ) { return key . startsWith ( BrowserIndexedDB . URL_SCHEME ) ? key . slice ( BrowserIndexedDB . URL_SCHEME . length ) : key ; } export class BrowserIndexedDBManager implements ModelStoreManager { private indexedDB : IDBFactory ; constructor ( ) { this . indexedDB = getIndexedDBFactory ( ) ; } async listModels ( ) { return new Promise < { [ path ] : ModelArtifactsInfo ; } > ( ( resolve , reject ) => { const openRequest = this . indexedDB . open ( DATABASE_NAME , DATABASE_VERSION ) ; openRequest . onupgradeneeded = ( ) => setUpDatabase ( openRequest ) ; openRequest . onsuccess = ( ) => { const db = openRequest . result ; const tx = db . transaction ( INFO_STORE_NAME , 'readonly' ) ; const store = tx . objectStore ( INFO_STORE_NAME ) ; const getAllInfoRequest = ( store as any ) . getAll ( ) as IDBRequest ; getAllInfoRequest . onsuccess = ( ) => { const out = { } ; for ( const item of getAllInfoRequest . result ) { out [ item . modelPath ] = item . modelArtifactsInfo ; } resolve ( out ) ; } ; getAllInfoRequest . onerror = error => { db . close ( ) ; return reject ( getAllInfoRequest . error ) ; } ; tx . oncomplete = ( ) => db . close ( ) ; } ; openRequest . onerror = error => reject ( openRequest . error ) ; } ) ; } async removeModel ( path ) { path = maybeStripScheme ( path ) ; return new Promise < ModelArtifactsInfo > ( ( resolve , reject ) => { const openRequest = this . indexedDB . open ( DATABASE_NAME , DATABASE_VERSION ) ; openRequest . onupgradeneeded = ( ) => setUpDatabase ( openRequest ) ; openRequest . onsuccess = ( ) => { const db = openRequest . result ; const infoTx = db . transaction ( INFO_STORE_NAME , 'readwrite' ) ; const infoStore = infoTx . objectStore ( INFO_STORE_NAME ) ; const getInfoRequest = infoStore . get ( path ) ; let modelTx ; getInfoRequest . onsuccess = ( ) => { if ( getInfoRequest . result == null ) { db . close ( ) ; return reject ( new Error ( `Cannot find model with path '${ path }' ` + `in IndexedDB.` ) ) ; } else { const deleteInfoRequest = infoStore . delete ( path ) ; const deleteModelData = ( ) => { modelTx = db . transaction ( MODEL_STORE_NAME , 'readwrite' ) ; const modelStore = modelTx . objectStore ( MODEL_STORE_NAME ) ; const deleteModelRequest = modelStore . delete ( path ) ; deleteModelRequest . onsuccess = ( ) => resolve ( getInfoRequest . result . modelArtifactsInfo ) ; deleteModelRequest . onerror = error => reject ( getInfoRequest . error ) ; } ; deleteInfoRequest . onsuccess = deleteModelData ; deleteInfoRequest . onerror = error => { deleteModelData ( ) ; db . close ( ) ; return reject ( getInfoRequest . error ) ; } ; } } ; getInfoRequest . onerror = error => { db . close ( ) ; return reject ( getInfoRequest . error ) ; } ; infoTx . oncomplete = ( ) => { if ( modelTx == null ) { db . close ( ) ; } else { modelTx . oncomplete = ( ) => db . close ( ) ; } } ; } ; openRequest . onerror = error => reject ( openRequest . error ) ; } ) ; } }