chore(data-warehouse): Enable schema specific syncs during the onboar…

…ding wizard (#21797) * Enable schema specific syncs during the onboarding wizard * Update query snapshots * Update query snapshots * Fixed tests --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>
PostHog · Apr 25, 2024 · c7d506b · c7d506b
1 parent f96ec7d
commit c7d506b
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 12 deletions.
diff --git 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,22 +45,22 @@ export const SyncProgressStep = (): JSX.Element => {
             <div>
                 <LemonTable
                     emptyState="No schemas selected"
-                    dataSource={databaseSchema}
+                    dataSource={schemas}
                     loading={dataWarehouseSourcesLoading}
                     disableTableWhileLoading={false}
                     columns={[
                         {
                             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
@@ -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
@@ -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)

diff --git a/posthog/warehouse/api/test/test_external_data_source.py b/posthog/warehouse/api/test/test_external_data_source.py
@@ -228,7 +228,7 @@ def test_internal_postgres(self, patch_get_postgres_schemas):
                 },
             )
             self.assertEqual(response.status_code, 200)
-            self.assertEqual(response.json(), [{"should_sync": False, "table": "table_1"}])
+            self.assertEqual(response.json(), [{"should_sync": True, "table": "table_1"}])
 
             new_team = Team.objects.create(name="new_team", organization=self.team.organization)
 
@@ -262,7 +262,7 @@ def test_internal_postgres(self, patch_get_postgres_schemas):
                 },
             )
             self.assertEqual(response.status_code, 200)
-            self.assertEqual(response.json(), [{"should_sync": False, "table": "table_1"}])
+            self.assertEqual(response.json(), [{"should_sync": True, "table": "table_1"}])
 
             new_team = Team.objects.create(name="new_team", organization=self.team.organization)