From 90a4fd79f1f2e4b157fc793349ca8781c4ff67e1 Mon Sep 17 00:00:00 2001
From: Tom Owers <owerstom@gmail.com>
Date: Wed, 24 Apr 2024 11:54:13 +0100
Subject: [PATCH] Enable schema specific syncs during the onboarding wizard

---
 .../external/forms/SyncProgressStep.tsx       | 19 +++++++++++--------
 .../data-warehouse/new/sourceWizardLogic.tsx  |  1 +
 posthog/warehouse/api/external_data_source.py |  4 ++--
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/frontend/src/scenes/data-warehouse/external/forms/SyncProgressStep.tsx b/frontend/src/scenes/data-warehouse/external/forms/SyncProgressStep.tsx
index 2432f349e3040..1e58314fdb862 100644
--- a/frontend/src/scenes/data-warehouse/external/forms/SyncProgressStep.tsx
+++ b/frontend/src/scenes/data-warehouse/external/forms/SyncProgressStep.tsx
@@ -3,28 +3,31 @@ import { useValues } from 'kea'
 import { sourceWizardLogic } from 'scenes/data-warehouse/new/sourceWizardLogic'
 import { dataWarehouseSettingsLogic } from 'scenes/data-warehouse/settings/dataWarehouseSettingsLogic'
 
+import { ExternalDataSourceSchema } from '~/types'
+
 export const SyncProgressStep = (): JSX.Element => {
-    const { databaseSchema, sourceId } = useValues(sourceWizardLogic)
+    const { sourceId } = useValues(sourceWizardLogic)
     const { dataWarehouseSources, dataWarehouseSourcesLoading } = useValues(dataWarehouseSettingsLogic)
 
     const source = dataWarehouseSources?.results.find((n) => n.id === sourceId)
+    const schemas = source?.schemas ?? []
 
-    const getSyncStatus = (shouldSync: boolean): { status: string; tagType: LemonTagType } => {
-        if (!shouldSync) {
+    const getSyncStatus = (schema: ExternalDataSourceSchema): { status: string; tagType: LemonTagType } => {
+        if (!schema.should_sync) {
             return {
                 status: 'Not synced',
                 tagType: 'default',
             }
         }
 
-        if (!source || source.status === 'Running') {
+        if (schema.status === 'Running') {
             return {
                 status: 'Syncing...',
                 tagType: 'primary',
             }
         }
 
-        if (source.status === 'Completed') {
+        if (schema.status === 'Completed') {
             return {
                 status: 'Completed',
                 tagType: 'success',
@@ -42,7 +45,7 @@ export const SyncProgressStep = (): JSX.Element => {
             <div>
                 <LemonTable
                     emptyState="No schemas selected"
-                    dataSource={databaseSchema}
+                    dataSource={schemas}
                     loading={dataWarehouseSourcesLoading}
                     disableTableWhileLoading={false}
                     columns={[
@@ -50,14 +53,14 @@ export const SyncProgressStep = (): JSX.Element => {
                             title: 'Table',
                             key: 'table',
                             render: function RenderTable(_, schema) {
-                                return schema.table
+                                return schema.name
                             },
                         },
                         {
                             title: 'Status',
                             key: 'status',
                             render: function RenderStatus(_, schema) {
-                                const { status, tagType } = getSyncStatus(schema.should_sync)
+                                const { status, tagType } = getSyncStatus(schema)
 
                                 return <LemonTag type={tagType}>{status}</LemonTag>
                             },
diff --git a/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx b/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx
index 2f400d9e14872..7ff85f473dd8f 100644
--- a/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx
+++ b/frontend/src/scenes/data-warehouse/new/sourceWizardLogic.tsx
@@ -448,6 +448,7 @@ export const sourceWizardLogic = kea<sourceWizardLogicType>([
                 lemonToast.success('New Data Resource Created')
                 actions.setSourceId(id)
                 actions.resetSourceConnectionDetails()
+                actions.loadSources(null)
                 actions.onNext()
             } catch (e: any) {
                 lemonToast.error(e.data?.message ?? e.message)
diff --git a/posthog/warehouse/api/external_data_source.py b/posthog/warehouse/api/external_data_source.py
index 36142a805938c..80576bc5da2d5 100644
--- a/posthog/warehouse/api/external_data_source.py
+++ b/posthog/warehouse/api/external_data_source.py
@@ -444,7 +444,7 @@ def database_schema(self, request: Request, *arg: Any, **kwargs: Any):
                     },
                 )
 
-            result_mapped_to_options = [{"table": row, "should_sync": False} for row in result]
+            result_mapped_to_options = [{"table": row, "should_sync": True} for row in result]
             return Response(status=status.HTTP_200_OK, data=result_mapped_to_options)
 
         # Return the possible endpoints for all other source types
@@ -455,7 +455,7 @@ def database_schema(self, request: Request, *arg: Any, **kwargs: Any):
                 data={"message": "Invalid parameter: source_type"},
             )
 
-        options = [{"table": row, "should_sync": False} for row in schemas]
+        options = [{"table": row, "should_sync": True} for row in schemas]
         return Response(status=status.HTTP_200_OK, data=options)
 
     @action(methods=["POST"], detail=False)