From 9890e6bc95e00c1ee1007a4cff1578383949a746 Mon Sep 17 00:00:00 2001
From: James Daniels <james@jamesdaniels.net>
Date: Thu, 5 Apr 2018 15:54:41 -0400
Subject: [PATCH] fix(storage): Add our zone wrappers for download url and
 meta-data

---
 src/storage/ref.ts     | 17 +++++++++++++----
 src/storage/storage.ts |  8 +++++---
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/storage/ref.ts b/src/storage/ref.ts
index 6ca26600f..16a02df9c 100644
--- a/src/storage/ref.ts
+++ b/src/storage/ref.ts
@@ -2,6 +2,7 @@ import { SettableMetadata, UploadMetadata, Reference, StringFormat } from '@fire
 import { createUploadTask, AngularFireUploadTask } from './task';
 import { Observable } from 'rxjs/Observable';
 import { from } from 'rxjs/observable/from';
+import { FirebaseZoneScheduler } from 'angularfire2';
 
 export interface AngularFireStorageReference {
   getDownloadURL(): Observable<any>;
@@ -18,12 +19,20 @@ export interface AngularFireStorageReference {
  * creates observable methods from promise based methods.
  * @param ref
  */
-export function createStorageRef(ref: Reference): AngularFireStorageReference {
+export function createStorageRef(ref: Reference, scheduler: FirebaseZoneScheduler): AngularFireStorageReference {
   return {
-    getDownloadURL: () => from(ref.getDownloadURL()),
-    getMetadata: () => from(ref.getMetadata()),
+    getDownloadURL: () => scheduler.keepUnstableUntilFirst(
+      scheduler.runOutsideAngular(
+        from(ref.getDownloadURL())
+      )
+    ),
+    getMetadata: () => scheduler.keepUnstableUntilFirst(
+      scheduler.runOutsideAngular(
+        from(ref.getMetadata())
+      )
+    ),
     delete: () => from(ref.delete()),
-    child: (path: string) => createStorageRef(ref.child(path)),
+    child: (path: string) => createStorageRef(ref.child(path), scheduler),
     updateMetatdata: (meta: SettableMetadata) => from(ref.updateMetadata(meta)),
     put: (data: any, metadata?: UploadMetadata) => {
       const task = ref.put(data, metadata);
diff --git a/src/storage/storage.ts b/src/storage/storage.ts
index 6a659096c..d482870e0 100644
--- a/src/storage/storage.ts
+++ b/src/storage/storage.ts
@@ -3,7 +3,7 @@ import { FirebaseStorage, UploadMetadata } from '@firebase/storage-types';
 import { createStorageRef, AngularFireStorageReference } from './ref';
 import { createUploadTask, AngularFireUploadTask } from './task';
 import { Observable } from 'rxjs/Observable';
-import { FirebaseAppConfig, FirebaseAppName, _firebaseAppFactory } from 'angularfire2';
+import { FirebaseAppConfig, FirebaseAppName, FirebaseZoneScheduler, _firebaseAppFactory } from 'angularfire2';
 import { FirebaseOptions } from '@firebase/app-types';
 
 export const StorageBucket = new InjectionToken<string>('angularfire2.storageBucket');
@@ -18,6 +18,7 @@ export const StorageBucket = new InjectionToken<string>('angularfire2.storageBuc
 @Injectable()
 export class AngularFireStorage {
   public readonly storage: FirebaseStorage;
+  public readonly scheduler: FirebaseZoneScheduler;
 
   constructor(
     @Inject(FirebaseAppConfig) config:FirebaseOptions,
@@ -25,6 +26,7 @@ export class AngularFireStorage {
     @Optional() @Inject(StorageBucket) storageBucket:string,
     zone: NgZone
   ) {
+    this.scheduler = new FirebaseZoneScheduler(zone);
     this.storage = zone.runOutsideAngular(() => {
       const app = _firebaseAppFactory(config, name);
       return app.storage(storageBucket || undefined);
@@ -32,12 +34,12 @@ export class AngularFireStorage {
   }
 
   ref(path: string) {
-    return createStorageRef(this.storage.ref(path));
+    return createStorageRef(this.storage.ref(path), this.scheduler);
   }
 
   upload(path: string, data: any, metadata?: UploadMetadata) {
     const storageRef = this.storage.ref(path);
-    const ref = createStorageRef(storageRef);
+    const ref = createStorageRef(storageRef, this.scheduler);
     return ref.put(data, metadata);
   }