fix: builders (query builder, notebooks) fetch all buckets now (#3137)

Author: hoorayimhelping

src/buckets/actions/thunks.test.ts

@@ -56,7 +56,7 @@ describe('buckets thunks', () => {
     it('should return results upon success', async () => {
       mockGetBuckets(true)
 
-      const bucks = await fetchAllBuckets('ord01')
+      const {normalizedBuckets: bucks} = await fetchAllBuckets('ord01')
       expect(bucks.result).toContain('custom-id')
       expect(bucks.result).toContain('demo-bucket')
     })

src/buckets/actions/thunks.ts

@@ -90,16 +90,17 @@ export const getBuckets = () => async (
     }
     const org = getOrg(state)
 
-    let buckets
+    let bucketsResponse
     if (isFlagEnabled('fetchAllBuckets')) {
       // a limit of -1 means fetch all buckets for this org
-      buckets = await fetchAllBuckets(org.id, -1)
+      bucketsResponse = await fetchAllBuckets(org.id, -1)
     } else {
-      buckets = await fetchAllBuckets(org.id)
+      bucketsResponse = await fetchAllBuckets(org.id)
     }
-    dispatch(setBuckets(RemoteDataState.Done, buckets))
+    dispatch(
+      setBuckets(RemoteDataState.Done, bucketsResponse.normalizedBuckets)
+    )
   } catch (error) {
-    console.error(error)
     dispatch(setBuckets(RemoteDataState.Error))
     dispatch(notify(getBucketsFailed()))
   }

src/buckets/api/index.ts

@@ -24,8 +24,12 @@ export const fetchAllBuckets = async (orgID: string, limit = BUCKET_LIMIT) => {
     demoDataBuckets = await fetchDemoDataBuckets()
   }
 
-  return normalize<Bucket, BucketEntities, string[]>(
-    [...resp.data.buckets, ...demoDataBuckets],
-    arrayOfBuckets
-  )
+  return {
+    buckets: resp.data.buckets,
+    demoDataBuckets: demoDataBuckets,
+    normalizedBuckets: normalize<Bucket, BucketEntities, string[]>(
+      [...resp.data.buckets, ...demoDataBuckets],
+      arrayOfBuckets
+    ),
+  }
 }

src/external/monaco.flux.server.ts

@@ -164,8 +164,8 @@ export class LSPServer {
         limit = -1
       }
 
-      const buckets = await fetchAllBuckets(org.id, limit)
-      return Object.values(buckets.entities.buckets).map(b => b.name)
+      const {normalizedBuckets} = await fetchAllBuckets(org.id, limit)
+      return Object.values(normalizedBuckets.entities.buckets).map(b => b.name)
     } catch (e) {
       console.error(e)
       return []

src/timeMachine/actions/queryBuilderThunks.ts

@@ -1,20 +1,19 @@
 import {Dispatch} from 'react'
+import {get} from 'lodash'
 
 import {queryBuilderFetcher} from 'src/timeMachine/apis/QueryBuilderFetcher'
 
 // Utils
 import {event} from 'src/cloud/utils/reporting'
 import {prohibitedDeselect} from 'src/shared/copy/notifications'
+import {isFlagEnabled} from 'src/shared/utils/featureFlag'
+
 // API
-import {normalize} from 'normalizr'
-import {get} from 'lodash'
-import {fetchDemoDataBuckets} from 'src/cloud/apis/demodata'
-import * as api from 'src/client'
+import {fetchAllBuckets} from 'src/buckets/api'
 
 // Types
 import {
   Bucket,
-  BucketEntities,
   GetState,
   RemoteDataState,
   ResourceType,
@@ -57,13 +56,10 @@ import {
 } from 'src/timeMachine/actions/queryBuilder'
 import {setBuckets} from 'src/buckets/actions/creators'
 import {notify} from 'src/shared/actions/notifications'
+
 // Constants
-import {LIMIT} from 'src/resources/constants'
 import {AGG_WINDOW_AUTO} from 'src/timeMachine/constants/queryBuilder'
 
-// Schemas
-import {arrayOfBuckets} from 'src/schemas'
-
 export const removeTagSelector = (index: number) => (
   dispatch: Dispatch<Action>
 ) => {
@@ -232,25 +228,23 @@ export const loadBuckets = () => async (
   const orgID = getOrg(getState()).id
   dispatch(setBuilderBucketsStatus(RemoteDataState.Loading))
 
+  let bucketsResponse
   try {
-    const resp = await api.getBuckets({query: {orgID, limit: LIMIT}})
-
-    if (resp.status !== 200) {
-      throw new Error(resp.data.message)
+    if (isFlagEnabled('fetchAllBuckets')) {
+      // a limit of -1 means fetch all buckets for this org
+      bucketsResponse = await fetchAllBuckets(orgID, -1)
+    } else {
+      bucketsResponse = await fetchAllBuckets(orgID)
     }
 
-    const demoDataBuckets = await fetchDemoDataBuckets()
-
-    const normalizedBuckets = normalize<Bucket, BucketEntities, string[]>(
-      [...resp.data.buckets, ...demoDataBuckets],
-      arrayOfBuckets
+    dispatch(
+      setBuckets(RemoteDataState.Done, bucketsResponse.normalizedBuckets)
     )
 
-    dispatch(setBuckets(RemoteDataState.Done, normalizedBuckets))
-
-    const allBuckets = [...resp.data.buckets, ...demoDataBuckets].map(
-      b => b.name
-    )
+    const allBuckets = [
+      ...bucketsResponse.buckets,
+      ...bucketsResponse.demoDataBuckets,
+    ].map(b => b.name)
 
     const systemBuckets = allBuckets.filter(b => b.startsWith('_'))
     const userBuckets = allBuckets.filter(b => !b.startsWith('_'))