diff --git a/packages/server/src/api/controllers/datasource.ts b/packages/server/src/api/controllers/datasource.ts
index 6961011b0d..a702a8cd84 100644
--- a/packages/server/src/api/controllers/datasource.ts
+++ b/packages/server/src/api/controllers/datasource.ts
@@ -63,24 +63,16 @@ export async function information(
   }
 }
 
-export async function buildSchemaFromDb(
+export async function buildSchemaFromSource(
   ctx: UserCtx<BuildSchemaFromSourceRequest, BuildSchemaFromSourceResponse>
 ) {
-  const db = context.getAppDB()
+  const datasourceId = ctx.params.datasourceId
   const tablesFilter = ctx.request.body.tablesFilter
-  const datasource = await sdk.datasources.get(ctx.params.datasourceId)
 
-  const { tables, errors } = await sdk.datasources.buildFilteredSchema(
-    datasource,
+  const { datasource, errors } = await sdk.datasources.buildSchemaFromSource(
+    datasourceId,
     tablesFilter
   )
-  datasource.entities = tables
-
-  setDefaultDisplayColumns(datasource)
-  const dbResp = await db.put(
-    sdk.tables.populateExternalTableSchemas(datasource)
-  )
-  datasource._rev = dbResp.rev
 
   ctx.body = {
     datasource: await sdk.datasources.removeSecretSingle(datasource),
@@ -88,24 +80,6 @@ export async function buildSchemaFromDb(
   }
 }
 
-/**
- * Make sure all datasource entities have a display name selected
- */
-function setDefaultDisplayColumns(datasource: Datasource) {
-  //
-  for (let entity of Object.values(datasource.entities || {})) {
-    if (entity.primaryDisplay) {
-      continue
-    }
-    const notAutoColumn = Object.values(entity.schema).find(
-      schema => !schema.autocolumn
-    )
-    if (notAutoColumn) {
-      entity.primaryDisplay = notAutoColumn.name
-    }
-  }
-}
-
 /**
  * Check for variables that have been updated or removed and invalidate them.
  */
diff --git a/packages/server/src/api/routes/datasource.ts b/packages/server/src/api/routes/datasource.ts
index 7b4945806a..755088c56c 100644
--- a/packages/server/src/api/routes/datasource.ts
+++ b/packages/server/src/api/routes/datasource.ts
@@ -53,7 +53,7 @@ router
   .post(
     "/api/datasources/:datasourceId/schema",
     authorized(permissions.BUILDER),
-    datasourceController.buildSchemaFromDb
+    datasourceController.buildSchemaFromSource
   )
   .post(
     "/api/datasources",
diff --git a/packages/server/src/sdk/app/datasources/datasources.ts b/packages/server/src/sdk/app/datasources/datasources.ts
index bb0d8bcb13..c71c3f1b31 100644
--- a/packages/server/src/sdk/app/datasources/datasources.ts
+++ b/packages/server/src/sdk/app/datasources/datasources.ts
@@ -337,7 +337,7 @@ const preSaveAction: Partial<Record<SourceName, any>> = {
 /**
  * Make sure all datasource entities have a display name selected
  */
-function setDefaultDisplayColumns(datasource: Datasource) {
+export function setDefaultDisplayColumns(datasource: Datasource) {
   //
   for (let entity of Object.values(datasource.entities || {})) {
     if (entity.primaryDisplay) {
diff --git a/packages/server/src/sdk/app/datasources/plus.ts b/packages/server/src/sdk/app/datasources/plus.ts
index 117d19a6a7..04cd508863 100644
--- a/packages/server/src/sdk/app/datasources/plus.ts
+++ b/packages/server/src/sdk/app/datasources/plus.ts
@@ -5,7 +5,9 @@ import {
   Schema,
 } from "@budibase/types"
 import * as datasources from "./datasources"
+import tableSdk from "../tables"
 import { getIntegration } from "../../../integrations"
+import { context } from "@budibase/backend-core"
 
 export async function buildFilteredSchema(
   datasource: Datasource,
@@ -60,3 +62,24 @@ export async function getAndMergeDatasource(datasource: Datasource) {
   }
   return await datasources.enrich(datasource)
 }
+
+export async function buildSchemaFromSource(
+  datasourceId: string,
+  tablesFilter?: string[]
+) {
+  const db = context.getAppDB()
+
+  const datasource = await datasources.get(datasourceId)
+
+  const { tables, errors } = await buildFilteredSchema(datasource, tablesFilter)
+  datasource.entities = tables
+
+  datasources.setDefaultDisplayColumns(datasource)
+  const dbResp = await db.put(tableSdk.populateExternalTableSchemas(datasource))
+  datasource._rev = dbResp.rev
+
+  return {
+    datasource,
+    errors,
+  }
+}