@stackbit/cms-core 0.1.6 → 0.1.7

package/src/content-store.ts

@@ -4,7 +4,21 @@ import sanitizeFilename from 'sanitize-filename';
 
 import * as CSITypes from '@stackbit/types';
 import { getLocalizedFieldForLocale, UserCommandSpawner } from '@stackbit/types';
-import { Config, extendConfig, Field, loadConfigFromDir, Model, ImageModel, isImageModel, Preset, RawConfigWithPaths } from '@stackbit/sdk';
+import {
+    Config,
+    Field,
+    Model,
+    ImageModel,
+    Preset,
+    PresetMap,
+    loadPresets,
+    getYamlModelDirs,
+    getPresetDirs,
+    loadYamlModelsFromFiles,
+    mergeConfigModelsWithModelsFromFiles,
+    extendModelsWithPresetsIds,
+    mergeConfigModelsWithExternalModels
+} from '@stackbit/sdk';
 import { deferWhileRunning, mapPromise, reducePromise } from '@stackbit/utils';
 
 import * as ContentStoreTypes from './content-store-types';
@@ -12,9 +26,20 @@ import { Timer } from './utils/timer';
 import { SearchFilter } from './types/search-filter';
 import { searchDocuments } from './utils/search-utils';
 import { mapCSIAssetsToStoreAssets, mapCSIDocumentsToStoreDocuments } from './utils/csi-to-store-docs-converter';
-import { getContentSourceId, getContentSourceIdForContentSource, getModelFieldForFieldAtPath, getUserContextForSrcType } from './content-store-utils';
+import {
+    getContentSourceId,
+    getContentSourceIdForContentSource,
+    getModelFieldForFieldAtPath,
+    getUserContextForSrcType,
+    groupModelsByContentSource,
+    groupDocumentsByContentSource
+} from './content-store-utils';
 import { mapAssetsToLocalizedApiImages, mapDocumentsToLocalizedApiObjects, mapStoreAssetsToAPIAssets } from './utils/store-to-api-docs-converter';
 import { convertOperationField, createDocumentRecursively, getCreateDocumentThunk } from './utils/create-update-csi-docs';
+import { normalizeModels, validateModels } from './utils/model-utils';
+import { IMAGE_MODEL } from './common/common-schema';
+
+export type HandleConfigAssets = <T extends Model>({ models, presets }: { models?: T[]; presets?: PresetMap }) => Promise<{ models: T[]; presets: PresetMap }>;
 
 export interface ContentSourceOptions {
     logger: ContentStoreTypes.Logger;
@@ -24,21 +49,23 @@ export interface ContentSourceOptions {
     userCommandSpawner?: UserCommandSpawner;
     onSchemaChangeCallback: () => void;
     onContentChangeCallback: (contentChanges: ContentStoreTypes.ContentChangeResult) => void;
-    handleConfigAssets: (config: Config) => Promise<Config>;
+    handleConfigAssets: HandleConfigAssets;
     devAppRestartNeeded?: () => void;
 }
 
 interface ContentSourceData {
     id: string;
     instance: CSITypes.ContentSourceInterface;
-    type: string;
-    projectId: string;
+    srcType: string;
+    srcProjectId: string;
     locales?: CSITypes.Locale[];
     defaultLocaleCode?: string;
     /* Array of extended and validated Models */
     models: Model[];
     /* Map of extended and validated Models by model name */
     modelMap: Record<string, Model>;
+    /* Array of original Models (as provided by content source)  */
+    csiModels: CSITypes.Model[];
     /* Map of original Models (as provided by content source) by model name */
     csiModelMap: Record<string, CSITypes.Model>;
     /* Array of original content source Documents */
@@ -59,6 +86,8 @@ interface ContentSourceData {
     assetMap: Record<string, ContentStoreTypes.Asset>;
 }
 
+type ContentSourceRawData = Omit<ContentSourceData, 'models' | 'modelMap' | 'documents' | 'documentMap'>;
+
 export class ContentStore {
     private readonly logger: ContentStoreTypes.Logger;
     private readonly userLogger: ContentStoreTypes.Logger;
@@ -67,13 +96,15 @@ export class ContentStore {
     private readonly webhookUrl?: string;
     private readonly onSchemaChangeCallback: () => void;
     private readonly onContentChangeCallback: (contentChanges: ContentStoreTypes.ContentChangeResult) => void;
-    private readonly handleConfigAssets: (config: Config) => Promise<Config>;
+    private readonly handleConfigAssets: HandleConfigAssets;
     private readonly devAppRestartNeeded?: () => void;
     private contentSources: CSITypes.ContentSourceInterface[] = [];
     private contentSourceDataById: Record<string, ContentSourceData> = {};
     private contentUpdatesWatchTimer: Timer;
-    private rawStackbitConfig: RawConfigWithPaths | null = null;
-    private presets?: Record<string, any>;
+    private stackbitConfig: Config | null = null;
+    private yamlModels: Model[] = [];
+    private configModels: Model[] = [];
+    private presets: PresetMap = {};
 
     constructor(options: ContentSourceOptions) {
         this.logger = options.logger.createLogger({ label: 'content-store' });
@@ -120,6 +151,10 @@ export class ContentStore {
 
     async init({ stackbitConfig }: { stackbitConfig: Config | null }): Promise<void> {
         this.logger.debug('init');
+        if (stackbitConfig) {
+            this.yamlModels = await this.loadYamlModels({ stackbitConfig });
+            this.presets = await this.loadPresets({ stackbitConfig });
+        }
         await this.setStackbitConfig({ stackbitConfig, init: true });
     }
 
@@ -129,37 +164,19 @@ export class ContentStore {
     }
 
     private async setStackbitConfig({ stackbitConfig, init }: { stackbitConfig: Config | null; init: boolean }) {
-        // TODO: Use stackbitConfig instead of loading rawStackbitConfig here
-        //  by splitting config loader into independent phases:
-        //  1. load and validate only the root config props
-        //  2. load content-source models
-        //  3. merge content-source models with config.models or via config.mapModels, sanitize and validate
-        //  4. load presets, adjust presets to have srcType and srcProjectId
-        if (!stackbitConfig) {
-            this.rawStackbitConfig = null;
-        } else {
-            const rawConfigResult = await loadConfigFromDir({
-                dirPath: stackbitConfig.dirPath,
-                logger: this.logger
-            });
-            for (const error of rawConfigResult.errors) {
-                this.userLogger.warn(error.message);
-            }
-            if (rawConfigResult.config) {
-                this.rawStackbitConfig = {
-                    ...rawConfigResult.config,
-                    contentSources: stackbitConfig.contentSources
-                };
-            } else {
-                this.rawStackbitConfig = null;
-            }
-        }
-        await this.loadContentSources({ init });
+        this.stackbitConfig = stackbitConfig;
+        this.configModels = this.mergeConfigModels(this.stackbitConfig, this.yamlModels);
+
+        await this.loadAllContentSourcesAndProcessData({ init });
     }
 
     /**
      * This method is called when contentUpdatesWatchTimer receives timeout.
-     * It then notifies all content sources to stop watching for content changes.
+     * This happens when the user is not using the Stackbit app for some time
+     * but container is not hibernated.
+     * It then notifies all content sources to stop watching for content
+     * changes, which in turn stops polling CMS for content changes and helps
+     * reducing the CMS API usage.
      */
     private handleTimerTimeout() {
         for (const contentSourceInstance of this.contentSources) {
@@ -179,7 +196,7 @@ export class ContentStore {
         }
 
         this.logger.debug('keepAlive => contentUpdatesWatchTimer is not running => load content source data');
-        await this.loadContentSources({ init: false });
+        await this.loadAllContentSourcesAndProcessData({ init: false });
     }
 
     /**
@@ -193,18 +210,55 @@ export class ContentStore {
      */
     async onContentSourceSchemaChange({ contentSourceId }: { contentSourceId: string }) {
         this.logger.debug('onContentSourceSchemaChange', { contentSourceId });
-        const contentSourceData = this.getContentSourceDataByIdOrThrow(contentSourceId);
-        this.contentSourceDataById[contentSourceId] = await this.loadContentSourceData({
-            contentSourceInstance: contentSourceData.instance,
-            init: false
-        });
-        this.onSchemaChangeCallback();
+        await this.reloadContentSourcesByIdAndProcessData({ contentSourceIds: [contentSourceId] });
     }
 
     async onFilesChange(updatedFiles: string[]): Promise<{ schemaChanged?: boolean; contentChanges: ContentStoreTypes.ContentChangeResult }> {
         this.logger.debug('onFilesChange');
 
-        let someContentSourceSchemaUpdated = false;
+        let schemaChanged = false;
+
+        if (this.stackbitConfig) {
+            // Check if any of the yaml models files were changed. If yaml model files were changed,
+            // reload them and merge them with models defined in stackbit config.
+            const modelDirs = getYamlModelDirs(this.stackbitConfig);
+            const yamlModelsChanged = updatedFiles.find((updatedFile) => _.some(modelDirs, (modelDir) => updatedFile.startsWith(modelDir)));
+            if (yamlModelsChanged) {
+                this.logger.debug('identified change in stackbit model files');
+                schemaChanged = true;
+                this.yamlModels = await this.loadYamlModels({ stackbitConfig: this.stackbitConfig });
+                this.configModels = this.mergeConfigModels(this.stackbitConfig, this.yamlModels);
+            }
+
+            // Check if any of the preset files were changed. If presets were changed, reload them.
+            const presetDirs = getPresetDirs(this.stackbitConfig);
+            const presetsChanged = updatedFiles.find((updatedFile) => _.some(presetDirs, (presetDir) => updatedFile.startsWith(presetDir)));
+            if (presetsChanged) {
+                this.logger.debug('identified change in stackbit preset files');
+                schemaChanged = true;
+                this.presets = await this.loadPresets({ stackbitConfig: this.stackbitConfig });
+            }
+        }
+
+        const contentSourceIdsWithChangedSchema: string[] = [];
+        const contentChangeEvents: { contentSourceId: string; contentChangeEvent: CSITypes.ContentChangeEvent }[] = [];
+
+        for (const contentSourceInstance of this.contentSources) {
+            const contentSourceId = getContentSourceIdForContentSource(contentSourceInstance);
+            this.logger.debug(`call onFilesChange for contentSource: ${contentSourceId}`);
+            const onFilesChangeResult = (await contentSourceInstance.onFilesChange?.({ updatedFiles: updatedFiles })) ?? {};
+            this.logger.debug(`schemaChanged: ${onFilesChangeResult.schemaChanged}, has contentChangeEvent: ${!!onFilesChangeResult.contentChangeEvent}`);
+
+            // if schema is changed, there is no need to return contentChanges
+            // because schema changes reloads everything and implies content changes
+            if (onFilesChangeResult.schemaChanged) {
+                schemaChanged = true;
+                contentSourceIdsWithChangedSchema.push(contentSourceId);
+            } else if (onFilesChangeResult.contentChangeEvent) {
+                contentChangeEvents.push({ contentSourceId, contentChangeEvent: onFilesChangeResult.contentChangeEvent });
+            }
+        }
+
         const contentChanges: ContentStoreTypes.ContentChangeResult = {
             updatedDocuments: [],
             updatedAssets: [],
@@ -212,49 +266,115 @@ export class ContentStore {
             deletedAssets: []
         };
 
-        for (const contentSourceInstance of this.contentSources) {
-            const contentSourceId = getContentSourceIdForContentSource(contentSourceInstance);
-            this.logger.debug(`call onFilesChange for contentSource: ${contentSourceId}`);
-            const { schemaChanged, contentChangeEvent } = (await contentSourceInstance.onFilesChange?.({ updatedFiles: updatedFiles })) ?? {};
-            this.logger.debug(`schemaChanged: ${schemaChanged}, has contentChangeEvent: ${!!contentChangeEvent}`);
-            // if schema is changed, there is no need to return contentChanges
-            // because schema changes reloads everything and implies content changes
-            if (schemaChanged) {
-                someContentSourceSchemaUpdated = true;
-                this.contentSourceDataById[contentSourceId] = await this.loadContentSourceData({ contentSourceInstance, init: false });
-            } else if (contentChangeEvent) {
+        // If the schema was changed, there is no need to accumulate or notify about content changes.
+        // The processData will update the store with the latest data. And once the Studio receives
+        // the schemaChanged notification it will reload all the models and the documents with their latest state.
+        if (schemaChanged) {
+            await this.reloadContentSourcesByIdAndProcessData({ contentSourceIds: contentSourceIdsWithChangedSchema });
+        } else {
+            contentChangeEvents.reduce((contentChanges, { contentSourceId, contentChangeEvent }) => {
                 const contentChangeResult = this.onContentChange(contentSourceId, contentChangeEvent);
                 contentChanges.updatedDocuments = contentChanges.updatedDocuments.concat(contentChangeResult.updatedDocuments);
                 contentChanges.updatedAssets = contentChanges.updatedAssets.concat(contentChangeResult.updatedAssets);
                 contentChanges.deletedDocuments = contentChanges.deletedDocuments.concat(contentChangeResult.deletedDocuments);
                 contentChanges.deletedAssets = contentChanges.deletedAssets.concat(contentChangeResult.deletedAssets);
-            }
+                return contentChanges;
+            }, contentChanges);
         }
 
         return {
-            schemaChanged: someContentSourceSchemaUpdated,
+            schemaChanged: schemaChanged,
             contentChanges: contentChanges
         };
     }
 
-    private async loadContentSources({ init }: { init: boolean }) {
-        this.logger.debug('loadContentSources', { init });
+    private async loadYamlModels({ stackbitConfig }: { stackbitConfig: Config }): Promise<Model[]> {
+        const yamlModelsResult = await loadYamlModelsFromFiles(stackbitConfig);
+        for (const error of yamlModelsResult.errors) {
+            this.userLogger.warn(error.message);
+        }
+        return yamlModelsResult.models;
+    }
+
+    private mergeConfigModels(stackbitConfig: Config | null, modelsFromFiles: Model[]) {
+        const configModelsResult = mergeConfigModelsWithModelsFromFiles(stackbitConfig?.models ?? [], modelsFromFiles);
+        for (const error of configModelsResult.errors) {
+            this.userLogger.warn(error.message);
+        }
+        return configModelsResult.models;
+    }
+
+    private async loadPresets({ stackbitConfig }: { stackbitConfig: Config }): Promise<Record<string, Preset>> {
+        const contentSources = stackbitConfig?.contentSources ?? [];
+        const singleContentSource = contentSources.length === 1 ? contentSources[0] : null;
+        const presetResult = await loadPresets({
+            config: stackbitConfig,
+            ...(singleContentSource
+                ? {
+                      fallbackSrcType: singleContentSource.getContentSourceType(),
+                      fallbackSrcProjectId: singleContentSource.getProjectId()
+                  }
+                : null)
+        });
+        for (const error of presetResult.errors) {
+            this.userLogger.warn(error.message);
+        }
+        const { presets } = await this.handleConfigAssets({ presets: presetResult.presets });
+        return presets;
+    }
+
+    private async loadAllContentSourcesAndProcessData({ init }: { init: boolean }) {
+        this.logger.debug('loadAllContentSourcesAndProcessData', { init });
 
         this.contentUpdatesWatchTimer.stopTimer();
 
-        // TODO: move CSITypes to separate package so Config will be able to use them
-        const contentSources: CSITypes.ContentSourceInterface[] = (this.rawStackbitConfig?.contentSources ?? []) as CSITypes.ContentSourceInterface[];
+        const contentSources = this.stackbitConfig?.contentSources ?? [];
 
         const promises = contentSources.map((contentSourceInstance) => {
             return this.loadContentSourceData({ contentSourceInstance, init });
         });
 
-        const contentSourceDataArr = await Promise.all(promises);
-        const contentSourceDataById: Record<string, ContentSourceData> = _.keyBy(contentSourceDataArr, 'id');
+        const contentSourceRawDataArr = await Promise.all(promises);
 
         // update all content sources at once to prevent race conditions
+        this.contentSourceDataById = await this.processData({
+            stackbitConfig: this.stackbitConfig,
+            configModels: this.configModels,
+            presets: this.presets,
+            contentSourceRawDataArr: contentSourceRawDataArr
+        });
         this.contentSources = contentSources;
-        this.contentSourceDataById = contentSourceDataById;
+
+        this.contentUpdatesWatchTimer.startTimer();
+    }
+
+    private async reloadContentSourcesByIdAndProcessData({ contentSourceIds }: { contentSourceIds: string[] }) {
+        this.logger.debug('reloadContentSourcesByIdAndProcessData', { contentSourceIds });
+
+        this.contentUpdatesWatchTimer.stopTimer();
+
+        const promises = this.contentSources.map(
+            (contentSourceInstance): Promise<ContentSourceRawData> => {
+                const contentSourceId = getContentSourceIdForContentSource(contentSourceInstance);
+                if (contentSourceIds.includes(contentSourceId)) {
+                    return this.loadContentSourceData({
+                        contentSourceInstance: contentSourceInstance,
+                        init: false
+                    });
+                } else {
+                    return Promise.resolve(_.omit(this.contentSourceDataById[contentSourceId], ['models', 'modelMap', 'documents', 'documentMap']));
+                }
+            }
+        );
+
+        const contentSourceRawDataArr = await Promise.all(promises);
+
+        this.contentSourceDataById = await this.processData({
+            stackbitConfig: this.stackbitConfig,
+            configModels: this.configModels,
+            presets: this.presets,
+            contentSourceRawDataArr: contentSourceRawDataArr
+        });
 
         this.contentUpdatesWatchTimer.startTimer();
     }
@@ -265,7 +385,7 @@ export class ContentStore {
     }: {
         contentSourceInstance: CSITypes.ContentSourceInterface;
         init: boolean;
-    }): Promise<ContentSourceData> {
+    }): Promise<ContentSourceRawData> {
         const contentSourceId = getContentSourceIdForContentSource(contentSourceInstance);
         this.logger.debug('loadContentSourceData', { contentSourceId, init });
 
@@ -283,98 +403,31 @@ export class ContentStore {
             await contentSourceInstance.reset();
         }
 
-        // TODO: introduce optimization: don't fetch content source models,
-        //  documents, assets if only stackbitConfig was changed
-
         const csiModels = await contentSourceInstance.getModels();
-        const locales = await contentSourceInstance?.getLocales();
-        const defaultLocaleCode = locales?.find((locale) => locale.default)?.code;
-
-        // for older versions of stackbit, it uses models to extend content source models
-        let modelsNoImage: Exclude<Model, ImageModel>[] = [];
-        let imageModel: ImageModel | undefined;
-        if (this.rawStackbitConfig) {
-            const result = await extendConfig({
-                config: this.rawStackbitConfig,
-                externalModels: csiModels
-            });
-            for (const error of result?.errors ?? []) {
-                this.userLogger.warn(error.message);
-            }
-            const config = await this.handleConfigAssets(result.config);
-            const modelsWithImageModel = config?.models ?? [];
-            const imageModelIndex = modelsWithImageModel.findIndex((model) => isImageModel(model));
-            if (imageModelIndex > -1) {
-                imageModel = modelsWithImageModel[imageModelIndex] as ImageModel;
-                modelsWithImageModel.splice(imageModelIndex, 1);
-            }
-            modelsNoImage = modelsWithImageModel as Exclude<Model, ImageModel>[];
-
-            // TODO: load presets externally from config, and create additional map
-            //  that maps presetIds by model name instead of storing that map inside every model
-
-            // Augment presets with srcType and srcProjectId if they don't exist
-            this.presets = _.reduce(
-                Object.keys(config?.presets ?? {}),
-                (accum: Record<string, Preset>, presetId) => {
-                    const preset = config?.presets?.[presetId];
-                    _.set(accum, [presetId], {
-                        ...preset,
-                        srcType: preset?.srcType ?? contentSourceInstance.getContentSourceType(),
-                        srcProjectId: preset?.srcProjectId ?? contentSourceInstance.getProjectId()
-                    });
-                    return accum;
-                },
-                {}
-            );
-        }
+        const csiModelMap = _.keyBy(csiModels, 'name');
 
-        if (this.rawStackbitConfig?.mapModels) {
-            const srcType = contentSourceInstance.getContentSourceType();
-            const srcProjectId = contentSourceInstance.getProjectId();
-            const modelsWithSource = modelsNoImage.map((model) => ({
-                srcType,
-                srcProjectId,
-                ...model
-            }));
-            const mappedModels = this.rawStackbitConfig.mapModels({
-                models: modelsWithSource
-            });
-            modelsNoImage = mappedModels.map((model) => {
-                const { srcType, srcProjectId, ...rest } = model;
-                return rest;
-            });
-        }
+        const locales = await contentSourceInstance.getLocales();
+        const defaultLocaleCode = locales?.find((locale) => locale.default)?.code;
 
-        const models: Model[] = imageModel ? [...modelsNoImage, imageModel] : modelsNoImage;
-        const modelMap = _.keyBy(models, 'name');
-        const csiModelMap = _.keyBy(csiModels, 'name');
         const csiDocuments = await contentSourceInstance.getDocuments({ modelMap: csiModelMap });
         const csiAssets = await contentSourceInstance.getAssets();
         const csiDocumentMap = _.keyBy(csiDocuments, 'id');
         const csiAssetMap = _.keyBy(csiAssets, 'id');
 
-        const contentStoreDocuments = mapCSIDocumentsToStoreDocuments({
-            csiDocuments,
-            contentSourceInstance,
-            modelMap,
-            defaultLocaleCode
-        });
         const contentStoreAssets = mapCSIAssetsToStoreAssets({
             csiAssets,
             contentSourceInstance,
             defaultLocaleCode
         });
-        const documentMap = _.keyBy(contentStoreDocuments, 'srcObjectId');
         const assetMap = _.keyBy(contentStoreAssets, 'srcObjectId');
 
         this.logger.debug('loaded content source data', {
             contentSourceId,
-            locales,
             defaultLocaleCode,
-            modelCount: models.length,
-            documentCount: contentStoreDocuments.length,
-            assetCount: contentStoreAssets.length
+            localesCount: locales.length,
+            modelCount: csiModels.length,
+            documentCount: csiDocuments.length,
+            assetCount: csiAssets.length
         });
 
         contentSourceInstance.startWatchingContentUpdates({
@@ -394,28 +447,21 @@ export class ContentStore {
             },
             onSchemaChange: async () => {
                 this.logger.debug('content source called onSchemaChange', { contentSourceId });
-                this.contentSourceDataById[contentSourceId] = await this.loadContentSourceData({
-                    contentSourceInstance: contentSourceInstance,
-                    init: false
-                });
-                this.onSchemaChangeCallback();
+                await this.reloadContentSourcesByIdAndProcessData({ contentSourceIds: [contentSourceId] });
             }
         });
 
         return {
             id: contentSourceId,
-            type: contentSourceInstance.getContentSourceType(),
-            projectId: contentSourceInstance.getProjectId(),
+            srcType: contentSourceInstance.getContentSourceType(),
+            srcProjectId: contentSourceInstance.getProjectId(),
             instance: contentSourceInstance,
             locales: locales,
             defaultLocaleCode: defaultLocaleCode,
-            models: models,
-            modelMap: modelMap,
+            csiModels: csiModels,
             csiModelMap: csiModelMap,
             csiDocuments: csiDocuments,
             csiDocumentMap: csiDocumentMap,
-            documents: contentStoreDocuments,
-            documentMap: documentMap,
             csiAssets: csiAssets,
             csiAssetMap: csiAssetMap,
             assets: contentStoreAssets,
@@ -458,8 +504,8 @@ export class ContentStore {
             }
 
             result.deletedDocuments.push({
-                srcType: contentSourceData.type,
-                srcProjectId: contentSourceData.projectId,
+                srcType: contentSourceData.srcType,
+                srcProjectId: contentSourceData.srcProjectId,
                 srcObjectId: docId
             });
         });
@@ -479,15 +525,43 @@ export class ContentStore {
             }
 
             result.deletedAssets.push({
-                srcType: contentSourceData.type,
-                srcProjectId: contentSourceData.projectId,
+                srcType: contentSourceData.srcType,
+                srcProjectId: contentSourceData.srcProjectId,
                 srcObjectId: assetId
             });
         });
 
+        // map csi documents through stackbitConfig.mapDocuments
+        let mappedDocs = contentChangeEvent.documents;
+        if (this.stackbitConfig?.mapDocuments) {
+            const csiDocumentsWithSource = contentChangeEvent.documents.map(
+                (csiDocument): CSITypes.DocumentWithSource => ({
+                    srcType: contentSourceData.srcType,
+                    srcProjectId: contentSourceData.srcProjectId,
+                    ...csiDocument
+                })
+            );
+
+            const modelsWithSource = contentSourceData.models.map(
+                (model): CSITypes.ModelWithSource => {
+                    return {
+                        srcType: contentSourceData.srcType,
+                        srcProjectId: contentSourceData.srcProjectId,
+                        ...model
+                    };
+                }
+            );
+
+            mappedDocs =
+                this.stackbitConfig?.mapDocuments?.({
+                    documents: _.cloneDeep(csiDocumentsWithSource),
+                    models: _.cloneDeep(modelsWithSource)
+                }) ?? csiDocumentsWithSource;
+        }
+
         // map csi documents and assets to content store documents and assets
         const documents = mapCSIDocumentsToStoreDocuments({
-            csiDocuments: contentChangeEvent.documents,
+            csiDocuments: mappedDocs,
             contentSourceInstance: contentSourceData.instance,
             modelMap: contentSourceData.modelMap,
             defaultLocaleCode: contentSourceData.defaultLocaleCode
@@ -500,7 +574,7 @@ export class ContentStore {
 
         // update contentSourceData with new or updated documents and assets
         Object.assign(contentSourceData.csiDocumentMap, _.keyBy(contentChangeEvent.documents, 'id'));
-        Object.assign(contentSourceData.csiAssets, _.keyBy(contentChangeEvent.assets, 'id'));
+        Object.assign(contentSourceData.csiAssetMap, _.keyBy(contentChangeEvent.assets, 'id'));
         Object.assign(contentSourceData.documentMap, _.keyBy(documents, 'srcObjectId'));
         Object.assign(contentSourceData.assetMap, _.keyBy(assets, 'srcObjectId'));
 
@@ -518,8 +592,8 @@ export class ContentStore {
                 contentSourceData.csiDocuments.splice(dataIndex, 1, csiDocument);
             }
             result.updatedDocuments.push({
-                srcType: contentSourceData.type,
-                srcProjectId: contentSourceData.projectId,
+                srcType: contentSourceData.srcType,
+                srcProjectId: contentSourceData.srcProjectId,
                 srcObjectId: document.srcObjectId
             });
         }
@@ -538,8 +612,8 @@ export class ContentStore {
                 contentSourceData.csiAssets.splice(index, 1, csiAsset);
             }
             result.updatedAssets.push({
-                srcType: contentSourceData.type,
-                srcProjectId: contentSourceData.projectId,
+                srcType: contentSourceData.srcType,
+                srcProjectId: contentSourceData.srcProjectId,
                 srcObjectId: asset.srcObjectId
             });
         }
@@ -547,13 +621,95 @@ export class ContentStore {
         return result;
     }
 
-    getModels(): Record<string, Record<string, Record<string, Model>>> {
+    private async processData({
+        stackbitConfig,
+        configModels,
+        presets,
+        contentSourceRawDataArr
+    }: {
+        stackbitConfig: Config | null;
+        configModels: Model[];
+        presets: Record<string, Preset>;
+        contentSourceRawDataArr: ContentSourceRawData[];
+    }): Promise<Record<string, ContentSourceData>> {
+        const modelsWithSource = contentSourceRawDataArr.reduce((accum: CSITypes.ModelWithSource[], csData) => {
+            const mergedModels = mergeConfigModelsWithExternalModels({
+                configModels: configModels,
+                externalModels: csData.csiModels
+            });
+            const modelsWithSource = mergedModels.map(
+                (model): CSITypes.ModelWithSource => {
+                    return {
+                        srcType: csData.srcType,
+                        srcProjectId: csData.id,
+                        ...model
+                    };
+                }
+            );
+            return accum.concat(modelsWithSource);
+        }, []);
+
+        // TODO: Is there a better way than deep cloning objects before passing them to user methods?
+        //   Not cloning mutable objects will break the internal state if user mutates the objects.
+        const mappedModels = stackbitConfig?.mapModels?.({ models: _.cloneDeep(modelsWithSource) }) ?? modelsWithSource;
+        const normalizedModels = normalizeModels({ models: mappedModels, logger: this.userLogger });
+        const validatedModels = validateModels({ models: normalizedModels, logger: this.userLogger });
+        const modelsWithPresetsIds = extendModelsWithPresetsIds({ models: validatedModels, presets });
+        const { models } = await this.handleConfigAssets({ models: modelsWithPresetsIds });
+
+        const csiDocumentsWithSource = contentSourceRawDataArr.reduce((accum: CSITypes.DocumentWithSource[], csData) => {
+            const csiDocumentsWithSource = csData.csiDocuments.map(
+                (csiDocument): CSITypes.DocumentWithSource => ({
+                    srcType: csData.srcType,
+                    srcProjectId: csData.id,
+                    ...csiDocument
+                })
+            );
+            return accum.concat(csiDocumentsWithSource);
+        }, []);
+
+        // TODO: Is there a better way than deep cloning objects before passing them to user methods?
+        //   Not cloning mutable objects will break the internal state if user mutates the objects.
+        const mappedDocs =
+            stackbitConfig?.mapDocuments?.({
+                documents: _.cloneDeep(csiDocumentsWithSource),
+                models: _.cloneDeep(models)
+            }) ?? csiDocumentsWithSource;
+
+        const modelMapByContentSource = groupModelsByContentSource({ models: models });
+        const documentMapByContentSource = groupDocumentsByContentSource({ documents: mappedDocs });
+
+        const contentSourceDataArr = contentSourceRawDataArr.map(
+            (csData): ContentSourceData => {
+                const modelMap = _.get(modelMapByContentSource, [csData.srcType, csData.id], {});
+                const mappedCSIDocuments = _.get(documentMapByContentSource, [csData.srcType, csData.id]);
+                const documents = mapCSIDocumentsToStoreDocuments({
+                    csiDocuments: mappedCSIDocuments,
+                    contentSourceInstance: csData.instance,
+                    defaultLocaleCode: csData.defaultLocaleCode,
+                    modelMap: modelMap
+                });
+                return {
+                    ...csData,
+                    models: Object.values(modelMap),
+                    modelMap,
+                    documents,
+                    documentMap: _.keyBy(documents, 'srcObjectId')
+                };
+            }
+        );
+
+        return _.keyBy(contentSourceDataArr, 'id');
+    }
+
+    getModels(): Record<string, Record<string, Record<string, Model | ImageModel>>> {
         return _.reduce(
             this.contentSourceDataById,
-            (result: Record<string, Record<string, Record<string, Model>>>, contentSourceData) => {
+            (result: Record<string, Record<string, Record<string, Model | ImageModel>>>, contentSourceData) => {
                 const contentSourceType = contentSourceData.instance.getContentSourceType();
                 const srcProjectId = contentSourceData.instance.getProjectId();
                 _.set(result, [contentSourceType, srcProjectId], contentSourceData.modelMap);
+                _.set(result, [contentSourceType, srcProjectId, '__image_model'], IMAGE_MODEL);
                 return result;
             },
             {}
@@ -601,8 +757,8 @@ export class ContentStore {
         return reducePromise(
             contentSourceDataArr,
             async (accum: ContentStoreTypes.HasAccessResult, contentSourceData) => {
-                const srcType = contentSourceData.type;
-                const srcProjectId = contentSourceData.projectId;
+                const srcType = contentSourceData.srcType;
+                const srcProjectId = contentSourceData.srcProjectId;
                 const userContext = getUserContextForSrcType(srcType, user);
                 let result = await contentSourceData.instance.hasAccess({ userContext });
                 // backwards compatibility with older CSI version
@@ -1286,15 +1442,15 @@ export class ContentStore {
             const contentSourceData = this.getContentSourceDataByIdOrThrow(contentSourceId);
             locale = locale ?? contentSourceData.defaultLocaleCode;
             const { documents, assets } = getCSIDocumentsAndAssetsFromContentSourceDataByIds(contentSourceData, contentSourceObjects);
-            const userContext = getUserContextForSrcType(contentSourceData.type, user);
+            const userContext = getUserContextForSrcType(contentSourceData.srcType, user);
             const internalValidationErrors = internalValidateContent(documents, assets, contentSourceData);
             const validationResult = await contentSourceData.instance.validateDocuments({ documents, assets, locale, userContext });
             errors = errors.concat(
                 internalValidationErrors,
                 validationResult.errors.map((validationError) => ({
                     message: validationError.message,
-                    srcType: contentSourceData.type,
-                    srcProjectId: contentSourceData.projectId,
+                    srcType: contentSourceData.srcType,
+                    srcProjectId: contentSourceData.srcProjectId,
                     srcObjectType: validationError.objectType,
                     srcObjectId: validationError.objectId,
                     fieldPath: validationError.fieldPath,
@@ -1344,7 +1500,7 @@ export class ContentStore {
             const contentSourceData = this.getContentSourceDataByIdOrThrow(contentSourceId);
 
             documents.push(...contentSourceData.documents);
-            _.set(schema, [contentSourceData.type, contentSourceData.projectId], contentSourceData.modelMap);
+            _.set(schema, [contentSourceData.srcType, contentSourceData.srcProjectId], contentSourceData.modelMap);
         });
 
         return searchDocuments({
@@ -1360,7 +1516,7 @@ export class ContentStore {
         const objectsBySourceId = _.groupBy(objects, (object) => getContentSourceId(object.srcType, object.srcProjectId));
         for (const [contentSourceId, contentSourceObjects] of Object.entries(objectsBySourceId)) {
             const contentSourceData = this.getContentSourceDataByIdOrThrow(contentSourceId);
-            const userContext = getUserContextForSrcType(contentSourceData.type, user);
+            const userContext = getUserContextForSrcType(contentSourceData.srcType, user);
             const { documents, assets } = getCSIDocumentsAndAssetsFromContentSourceDataByIds(contentSourceData, contentSourceObjects);
             await contentSourceData.instance.publishDocuments({ documents, assets, userContext });
         }
@@ -1461,8 +1617,8 @@ function validateDocumentFields(
         if (!objRef) {
             errors.push({
                 fieldPath,
-                srcType: contentSourceData.type,
-                srcProjectId: contentSourceData.projectId,
+                srcType: contentSourceData.srcType,
+                srcProjectId: contentSourceData.srcProjectId,
                 srcObjectType: documentField.refType,
                 srcObjectId: document.id,
                 message: `Can't find referenced ${documentField.refType}: ${documentField.refId}`