3 files changed, 22 insertions, 16 deletions

diff --git a/src/client/documents/Documents.ts b/src/client/documents/Documents.ts
index a7514f1d6..847f604d0 100644
--- a/src/client/documents/Documents.ts
+++ b/src/client/documents/Documents.ts
@@ -60,6 +60,7 @@ export interface DocumentOptions {
     copyDraggedItems?: boolean;
     documentText?: string;
     borderRounding?: number;
+    columnsKey?: Array<Key>;
 }
 
 export namespace Documents {
@@ -116,6 +117,7 @@ export namespace Documents {
         if (options.layoutKeys !== undefined) { doc.Set(KeyStore.LayoutKeys, new ListField(options.layoutKeys)); }
         if (options.copyDraggedItems !== undefined) { doc.SetBoolean(KeyStore.CopyDraggedItems, options.copyDraggedItems); }
         if (options.borderRounding !== undefined) { doc.SetNumber(KeyStore.BorderRounding, options.borderRounding); }
+        if (options.columnsKey !== undefined) { doc.SetData(KeyStore.ColumnsKey, options.columnsKey, ListField); }
         return doc;
     }
 
@@ -271,13 +273,13 @@ export namespace Documents {
         if (!makePrototype) {
             return SetInstanceOptions(collProto, { ...options, viewType: CollectionViewType.Freeform }, [documents, ListField], id);
         }
-        return assignToDelegate(SetInstanceOptions(collProto, { ...options, viewType: CollectionViewType.Freeform }, [documents, ListField], id).MakeDelegate(), options);
+        return assignToDelegate(SetInstanceOptions(collProto, { columnsKey: [KeyStore.Title], ...options, viewType: CollectionViewType.Freeform }, [documents, ListField], id).MakeDelegate(), options);
     }
     export function SchemaDocument(documents: Array<Document>, options: DocumentOptions, id?: string) {
-        return assignToDelegate(SetInstanceOptions(collProto, { ...options, viewType: CollectionViewType.Schema }, [documents, ListField], id), options);
+        return assignToDelegate(SetInstanceOptions(collProto, { columnsKey: [KeyStore.Title], ...options, viewType: CollectionViewType.Schema }, [documents, ListField], id), options);
     }
     export function TreeDocument(documents: Array<Document>, options: DocumentOptions, id?: string) {
-        return assignToDelegate(SetInstanceOptions(collProto, { ...options, viewType: CollectionViewType.Tree }, [documents, ListField], id), options);
+        return assignToDelegate(SetInstanceOptions(collProto, { columnsKey: [KeyStore.Title], ...options, viewType: CollectionViewType.Tree }, [documents, ListField], id), options);
     }
     export function DockDocument(config: string, options: DocumentOptions, id?: string) {
         return assignToDelegate(SetInstanceOptions(collProto, { ...options, viewType: CollectionViewType.Docking }, [config, TextField], id), options);
diff --git a/src/client/views/collections/CollectionSchemaView.tsx b/src/client/views/collections/CollectionSchemaView.tsx
index b6d5f1bfa..b3cc3fbd5 100644
--- a/src/client/views/collections/CollectionSchemaView.tsx
+++ b/src/client/views/collections/CollectionSchemaView.tsx
@@ -2,13 +2,13 @@ import React = require("react");
 import { library } from '@fortawesome/fontawesome-svg-core';
 import { faCog, faPlus } from '@fortawesome/free-solid-svg-icons';
 import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
-import { action, computed, observable, untracked } from "mobx";
+import { action, computed, observable, untracked, runInAction } from "mobx";
 import { observer } from "mobx-react";
 import ReactTable, { CellInfo, ComponentPropsGetterR, ReactTableDefaults } from "react-table";
 import { MAX_ROW_HEIGHT } from '../../views/globalCssVariables.scss'
 import "react-table/react-table.css";
 import { Document } from "../../../fields/Document";
-import { Field, Opt } from "../../../fields/Field";
+import { Field, Opt, FieldWaiting } from "../../../fields/Field";
 import { Key } from "../../../fields/Key";
 import { KeyStore } from "../../../fields/KeyStore";
 import { ListField } from "../../../fields/ListField";
@@ -165,18 +165,22 @@ export class CollectionSchemaView extends CollectionSubView {
         super.CreateDropTarget(ele);
     }
 
-    @action
     toggleKey = (key: Key) => {
-        this.props.Document.GetOrCreateAsync<ListField<Key>>(KeyStore.ColumnsKey, ListField,
-            (field) => {
-                const index = field.Data.indexOf(key);
-                if (index === -1) {
-                    this.columns.push(key);
-                } else {
-                    this.columns.splice(index, 1);
+        this.props.Document.GetTAsync<ListField<Key>>(KeyStore.ColumnsKey, ListField).then(field =>
+            runInAction(() => {
+                if (field !== FieldWaiting) {
+                    if (field) {
+                        const index = field.Data.indexOf(key);
+                        if (index === -1) {
+                            this.columns.push(key);
+                        } else {
+                            this.columns.splice(index, 1);
+                        }
+                    } else {
+                        this.props.Document.SetData(KeyStore.ColumnsKey, [key], ListField);
+                    }
                 }
-
-            });
+            }));
     }
 
     //toggles preview side-panel of schema
diff --git a/src/fields/Document.ts b/src/fields/Document.ts
index 2797efc09..55ac731dd 100644
--- a/src/fields/Document.ts
+++ b/src/fields/Document.ts
@@ -204,7 +204,7 @@ export class Document extends Field {
         callback: (field: T) => void
     ): void {
         //This currently doesn't deal with prototypes
-        if (this._proxies.has(key.Id)) {
+        if (this._proxies.has(key.Id) || this.fields.has(key.Id)) {
             Server.GetDocumentField(this, key, field => {
                 if (field && field instanceof ctor) {
                     callback(field);