chore(data-warehouse): Update self managed sources schema (#24613)

PostHog · Aug 28, 2024 · 1927d8d · 1927d8d
1 parent 04d32a6
commit 1927d8d
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 2 deletions.
diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts
@@ -2043,6 +2043,9 @@ const api = {
         ): Promise<void> {
             await new ApiRequest().dataWarehouseTable(tableId).withAction('update_schema').create({ data: { updates } })
         },
+        async refreshSchema(tableId: DataWarehouseTable['id']): Promise<void> {
+            await new ApiRequest().dataWarehouseTable(tableId).withAction('refresh_schema').create()
+        },
     },
 
     dataWarehouseSavedQueries: {

diff --git a/frontend/src/scenes/data-warehouse/settings/DataWarehouseSelfManagedSourcesTable.tsx b/frontend/src/scenes/data-warehouse/settings/DataWarehouseSelfManagedSourcesTable.tsx
@@ -7,7 +7,7 @@ import { dataWarehouseSettingsLogic } from './dataWarehouseSettingsLogic'
 
 export function DataWarehouseSelfManagedSourcesTable(): JSX.Element {
     const { selfManagedTables } = useValues(dataWarehouseSettingsLogic)
-    const { deleteSelfManagedTable } = useActions(dataWarehouseSettingsLogic)
+    const { deleteSelfManagedTable, refreshSelfManagedTableSchema } = useActions(dataWarehouseSettingsLogic)
 
     return (
         <LemonTable
@@ -24,6 +24,13 @@ export function DataWarehouseSelfManagedSourcesTable(): JSX.Element {
                     render: (_, item: DatabaseSchemaDataWarehouseTable) => {
                         return (
                             <div className="flex flex-row justify-end">
+                                <LemonButton
+                                    data-attr={`refresh-data-warehouse-${item.name}`}
+                                    key={`refresh-data-warehouse-${item.name}`}
+                                    onClick={() => refreshSelfManagedTableSchema(item.id)}
+                                >
+                                    Update schema from source
+                                </LemonButton>
                                 <LemonButton
                                     status="danger"
                                     data-attr={`delete-data-warehouse-${item.name}`}

diff --git a/frontend/src/scenes/data-warehouse/settings/dataWarehouseSettingsLogic.ts b/frontend/src/scenes/data-warehouse/settings/dataWarehouseSettingsLogic.ts
@@ -37,6 +37,7 @@ export const dataWarehouseSettingsLogic = kea<dataWarehouseSettingsLogicType>([
         schemaLoadingFinished: (schema: ExternalDataSourceSchema) => ({ schema }),
         abortAnyRunningQuery: true,
         deleteSelfManagedTable: (tableId: string) => ({ tableId }),
+        refreshSelfManagedTableSchema: (tableId: string) => ({ tableId }),
     }),
     loaders(({ cache, actions, values }) => ({
         dataWarehouseSources: [
@@ -148,6 +149,12 @@ export const dataWarehouseSettingsLogic = kea<dataWarehouseSettingsLogicType>([
             await api.dataWarehouseTables.delete(tableId)
             actions.loadDatabase()
         },
+        refreshSelfManagedTableSchema: async ({ tableId }) => {
+            lemonToast.info('Updating schema...')
+            await api.dataWarehouseTables.refreshSchema(tableId)
+            lemonToast.success('Schema updated')
+            actions.loadDatabase()
+        },
         deleteSource: async ({ source }) => {
             await api.externalDataSources.delete(source.id)
             actions.loadSources(null)

diff --git a/mypy-baseline.txt b/mypy-baseline.txt
@@ -804,6 +804,8 @@ posthog/warehouse/api/external_data_schema.py:0: note: def [_T] get(self, Type,
 posthog/warehouse/api/external_data_source.py:0: error: Incompatible return value type (got "tuple[ExternalDataSource, dict[str, list[tuple[str, str]]]]", expected "tuple[ExternalDataSource, list[Any]]")  [return-value]
 posthog/warehouse/api/external_data_source.py:0: error: Incompatible return value type (got "tuple[ExternalDataSource, dict[str, list[tuple[str, str]]]]", expected "tuple[ExternalDataSource, list[Any]]")  [return-value]
 posthog/warehouse/api/table.py:0: error: Unused "type: ignore" comment  [unused-ignore]
+posthog/warehouse/api/table.py:0: error: Unused "type: ignore" comment  [unused-ignore]
+posthog/warehouse/api/table.py:0: error: Unused "type: ignore" comment  [unused-ignore]
 posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py:0: error: Need type annotation for "_execute_calls" (hint: "_execute_calls: list[<type>] = ...")  [var-annotated]
 posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py:0: error: Need type annotation for "_execute_async_calls" (hint: "_execute_async_calls: list[<type>] = ...")  [var-annotated]
 posthog/temporal/tests/batch_exports/test_snowflake_batch_export_workflow.py:0: error: Need type annotation for "_cursors" (hint: "_cursors: list[<type>] = ...")  [var-annotated]

diff --git a/posthog/warehouse/api/table.py b/posthog/warehouse/api/table.py
@@ -113,7 +113,7 @@ def create(self, validated_data):
             )
         table = DataWarehouseTable(**validated_data)
         try:
-            table.columns = table.get_columns()
+            table.columns = table.get_columns()  # type: ignore
         except Exception as err:
             raise serializers.ValidationError(str(err))
         table.save()
@@ -240,3 +240,12 @@ def update_schema(self, request: request.Request, *args: Any, **kwargs: Any) ->
         table.save()
 
         return response.Response(status=status.HTTP_200_OK)
+
+    @action(methods=["POST"], detail=True)
+    def refresh_schema(self, request: request.Request, *args: Any, **kwargs: Any) -> response.Response:
+        table: DataWarehouseTable = self.get_object()
+
+        table.columns = table.get_columns()  # type: ignore
+        table.save()
+
+        return response.Response(status=status.HTTP_200_OK)