Show correct number of private snippets for SQL folders (#28411)

* Show correct number of private snippets for SQL folders

* Ensure that private snippet count is updated when sharing / unsharing queries

* Fix snippets caching issuie

* Fix move query modal showing more than just the project's folders

apps/studio/components/interfaces/SQLEditor/MoveQueryModal.tsx

@@ -10,7 +10,7 @@ import { getContentById } from 'data/content/content-id-query'
 import { useContentUpsertV2Mutation } from 'data/content/content-upsert-v2-mutation'
 import { useSQLSnippetFolderCreateMutation } from 'data/content/sql-folder-create-mutation'
 import { Snippet, SnippetDetail } from 'data/content/sql-folders-query'
-import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
+import { useSnippetFolders, useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
 import {
   Button,
   CommandEmpty_Shadcn_,
@@ -83,7 +83,7 @@ export const MoveQueryModal = ({ visible, snippets = [], onClose }: MoveQueryMod
     defaultValues: { name: '' },
   })
 
-  const folders = Object.values(snapV2.folders).map((x) => x.folder)
+  const folders = useSnippetFolders(ref as string)
   const selectedFolder =
     selectedId === 'root'
       ? 'Root of the editor'

apps/studio/components/layouts/SQLEditorLayout/SQLEditorNavV2/SQLEditorNav.tsx

@@ -9,7 +9,7 @@ import { MoveQueryModal } from 'components/interfaces/SQLEditor/MoveQueryModal'
 import RenameQueryModal from 'components/interfaces/SQLEditor/RenameQueryModal'
 import { createSqlSnippetSkeletonV2 } from 'components/interfaces/SQLEditor/SQLEditor.utils'
 import { useContentDeleteMutation } from 'data/content/content-delete-mutation'
-import { getContentById, useContentIdQuery } from 'data/content/content-id-query'
+import { getContentById } from 'data/content/content-id-query'
 import { useSQLSnippetFoldersDeleteMutation } from 'data/content/sql-folders-delete-mutation'
 import {
   getSQLSnippetFolders,
@@ -22,7 +22,12 @@ import { useLocalStorage } from 'hooks/misc/useLocalStorage'
 import { useSelectedProject } from 'hooks/misc/useSelectedProject'
 import { useProfile } from 'lib/profile'
 import uuidv4 from 'lib/uuid'
-import { useSqlEditorV2StateSnapshot } from 'state/sql-editor-v2'
+import {
+  useFavoriteSnippets,
+  useSnippetFolders,
+  useSnippets,
+  useSqlEditorV2StateSnapshot,
+} from 'state/sql-editor-v2'
 import {
   AlertDescription_Shadcn_,
   AlertTitle_Shadcn_,
@@ -39,6 +44,7 @@ import { ROOT_NODE, formatFolderResponseForTreeView } from './SQLEditorNav.utils
 import { SQLEditorTreeViewItem } from './SQLEditorTreeViewItem'
 import { untitledSnippetTitle } from 'components/interfaces/SQLEditor/SQLEditor.constants'
 import { useSqlSnippetsQuery } from 'data/content/sql-snippets-query'
+import { useContentCountQuery } from 'data/content/content-count-query'
 
 interface SQLEditorNavProps {
   searchText: string
@@ -77,20 +83,21 @@ export const SQLEditorNav = ({ searchText: _searchText }: SQLEditorNavProps) =>
   // =======================================================
   // [Joshen] Set up favorites, shared, and private snippets
   // =======================================================
-  const folders = Object.values(snapV2.folders).map((x) => x.folder)
-  const contents = snapV2.sortedSnippets.filter((x) =>
+  const snippets = useSnippets(projectRef as string)
+  const folders = useSnippetFolders(projectRef as string)
+  const contents = snippets.filter((x) =>
     searchText.length > 0 ? x.name.toLowerCase().includes(searchText.toLowerCase()) : true
   )
   const snippet = snapV2.snippets[id as string]?.snippet
 
   const privateSnippets = contents.filter((snippet) => snippet.visibility === 'user')
-  const numPrivateSnippets = Object.keys(snapV2.snippets).length
+  const numPrivateSnippets = snapV2.privateSnippetCount[projectRef as string]
   const privateSnippetsTreeState =
-    folders.length === 0 && numPrivateSnippets === 0
+    folders.length === 0 && snippets.length === 0
       ? [ROOT_NODE]
       : formatFolderResponseForTreeView({ folders, contents: privateSnippets })
 
-  const favoriteSnippets = snapV2.sortedFavoriteSnippets
+  const favoriteSnippets = useFavoriteSnippets(projectRef as string)
   const numFavoriteSnippets = favoriteSnippets.length
   const favoritesTreeState =
     numFavoriteSnippets === 0
@@ -130,6 +137,17 @@ export const SQLEditorNav = ({ searchText: _searchText }: SQLEditorNavProps) =>
     },
   })
 
+  useContentCountQuery(
+    { projectRef, type: 'sql' },
+    {
+      onSuccess(data) {
+        if (projectRef !== undefined) {
+          snapV2.setPrivateSnippetCount({ projectRef, value: data.count })
+        }
+      },
+    }
+  )
+
   const { mutate: deleteContent, isLoading: isDeleting } = useContentDeleteMutation({
     onError: (error, data) => {
       if (error.message.includes('Contents not found')) {
@@ -187,6 +205,13 @@ export const SQLEditorNav = ({ searchText: _searchText }: SQLEditorNavProps) =>
     snapV2.shareSnippet(selectedSnippetToShare.id, 'project')
     setSelectedSnippetToShare(undefined)
     setShowSharedSnippets(true)
+
+    if (projectRef !== undefined) {
+      snapV2.setPrivateSnippetCount({
+        projectRef,
+        value: snapV2.privateSnippetCount[projectRef] - 1,
+      })
+    }
   }
 
   const onConfirmUnshare = () => {
@@ -194,6 +219,13 @@ export const SQLEditorNav = ({ searchText: _searchText }: SQLEditorNavProps) =>
     snapV2.shareSnippet(selectedSnippetToUnshare.id, 'user')
     setSelectedSnippetToUnshare(undefined)
     setShowPrivateSnippets(true)
+
+    if (projectRef !== undefined) {
+      snapV2.setPrivateSnippetCount({
+        projectRef,
+        value: snapV2.privateSnippetCount[projectRef] + 1,
+      })
+    }
   }
 
   const onSelectCopyPersonal = async (snippet: Snippet) => {

apps/studio/data/content/content-count-query.ts

@@ -0,0 +1,41 @@
+import { useQuery, UseQueryOptions } from '@tanstack/react-query'
+
+import { get, handleError } from 'data/fetchers'
+import { ResponseError } from 'types'
+import { contentKeys } from './keys'
+
+interface getContentCountVariables {
+  projectRef?: string
+  type: 'sql' | 'report' | 'log_sql'
+}
+
+export async function getContentCount(
+  { projectRef, type }: getContentCountVariables,
+  signal?: AbortSignal
+) {
+  if (typeof projectRef === 'undefined') throw new Error('projectRef is required')
+
+  const { data, error } = await get('/platform/projects/{ref}/content/count', {
+    params: { path: { ref: projectRef }, query: { type } },
+    signal,
+  })
+
+  if (error) throw handleError(error)
+  return data
+}
+
+export type ContentIdData = Awaited<ReturnType<typeof getContentCount>>
+export type ContentIdError = ResponseError
+
+export const useContentCountQuery = <TData = ContentIdData>(
+  { projectRef, type }: getContentCountVariables,
+  { enabled = true, ...options }: UseQueryOptions<ContentIdData, ContentIdError, TData> = {}
+) =>
+  useQuery<ContentIdData, ContentIdError, TData>(
+    contentKeys.count(projectRef, type),
+    ({ signal }) => getContentCount({ projectRef, type }, signal),
+    {
+      enabled: enabled && typeof projectRef !== 'undefined',
+      ...options,
+    }
+  )

apps/studio/data/content/content-id-query.ts

@@ -3,7 +3,6 @@ import { useQuery, UseQueryOptions } from '@tanstack/react-query'
 import { get, handleError } from 'data/fetchers'
 import { ResponseError } from 'types'
 import { contentKeys } from './keys'
-import { components } from 'api-types'
 
 export async function getContentById(
   { projectRef, id }: { projectRef?: string; id?: string },

apps/studio/data/content/keys.ts

@@ -4,4 +4,6 @@ export const contentKeys = {
     ['projects', projectRef, 'content', 'folders', id].filter(Boolean),
   resource: (projectRef: string | undefined, id?: string) =>
     ['projects', projectRef, 'content', id] as const,
+  count: (projectRef: string | undefined, type?: string) =>
+    ['projects', projectRef, 'content', 'count', type].filter(Boolean),
 }

apps/studio/state/sql-editor-v2.ts

@@ -7,7 +7,9 @@ import { UpsertContentPayloadV2, upsertContent } from 'data/content/content-upse
 import { createSQLSnippetFolder } from 'data/content/sql-folder-create-mutation'
 import { updateSQLSnippetFolder } from 'data/content/sql-folder-update-mutation'
 import { Snippet, SnippetFolder, SnippetFolderResponse } from 'data/content/sql-folders-query'
-import { SqlSnippet, SqlSnippets } from 'data/content/sql-snippets-query'
+import { SqlSnippet } from 'data/content/sql-snippets-query'
+import { getQueryClient } from 'data/query-client'
+import { contentKeys } from 'data/content/keys'
 
 export type StateSnippetFolder = {
   projectRef: string
@@ -59,6 +61,9 @@ export const sqlEditorState = proxy({
   loaded: {} as {
     [projectRef: string]: boolean
   },
+  privateSnippetCount: {} as {
+    [projectRef: string]: number
+  },
   // Synchronous saving of folders and snippets (debounce behavior)
   needsSaving: proxySet<string>([]),
   // Stores the state of each snippet
@@ -68,25 +73,6 @@ export const sqlEditorState = proxy({
   limit: 100,
   order: 'inserted_at' as 'name' | 'inserted_at',
 
-  get sortedSnippets() {
-    return Object.values(sqlEditorState.snippets)
-      .map((x) => x.snippet)
-      .sort((a, b) => {
-        if (sqlEditorState.order === 'name') return a.name.localeCompare(b.name)
-        else return new Date(b.inserted_at).valueOf() - new Date(a.inserted_at).valueOf()
-      })
-  },
-
-  get sortedFavoriteSnippets() {
-    return Object.values(sqlEditorState.favoriteSnippets)
-      .map((x) => x.snippet)
-      .sort((a, b) => {
-        if (sqlEditorState.order === 'name') return a.name.localeCompare(b.name)
-        else
-          return new Date(b.inserted_at ?? '').valueOf() - new Date(a.inserted_at ?? '').valueOf()
-      })
-  },
-
   // ========================================================================
   // ## Methods to interact the store with
   // ========================================================================
@@ -130,6 +116,10 @@ export const sqlEditorState = proxy({
 
   setOrder: (value: 'name' | 'inserted_at') => (sqlEditorState.order = value),
 
+  setPrivateSnippetCount: ({ projectRef, value }: { projectRef: string; value: number }) => {
+    sqlEditorState.privateSnippetCount[projectRef] = value
+  },
+
   addSnippet: ({ projectRef, snippet }: { projectRef: string; snippet: Snippet }) => {
     if (snippet.id && sqlEditorState.snippets[snippet.id]?.snippet?.content === undefined) {
       sqlEditorState.snippets[snippet.id] = { projectRef, splitSizes: [50, 50], snippet }
@@ -315,6 +305,38 @@ export const getSqlEditorV2StateSnapshot = () => snapshot(sqlEditorState)
 export const useSqlEditorV2StateSnapshot = (options?: Parameters<typeof useSnapshot>[1]) =>
   useSnapshot(sqlEditorState, options)
 
+export const useSnippetFolders = (projectRef: string) => {
+  const snapshot = useSqlEditorV2StateSnapshot()
+
+  return Object.values(snapshot.folders)
+    .filter((x) => x.projectRef === projectRef)
+    .map((x) => x.folder)
+}
+
+export const useSnippets = (projectRef: string) => {
+  const snapshot = useSqlEditorV2StateSnapshot()
+
+  return Object.values(snapshot.snippets)
+    .filter((x) => x.projectRef === projectRef)
+    .map((x) => x.snippet)
+    .sort((a, b) => {
+      if (snapshot.order === 'name') return a.name.localeCompare(b.name)
+      else return new Date(b.inserted_at).valueOf() - new Date(a.inserted_at).valueOf()
+    })
+}
+
+export const useFavoriteSnippets = (projectRef: string) => {
+  const snapshot = useSqlEditorV2StateSnapshot()
+
+  return Object.values(snapshot.favoriteSnippets)
+    .filter((x) => x.projectRef === projectRef)
+    .map((x) => x.snippet)
+    .sort((a, b) => {
+      if (snapshot.order === 'name') return a.name.localeCompare(b.name)
+      else return new Date(b.inserted_at ?? '').valueOf() - new Date(a.inserted_at ?? '').valueOf()
+    })
+}
+
 // ========================================================================
 // ## Below are all the asynchronous saving logic for the SQL Editor
 // ========================================================================
@@ -322,10 +344,11 @@ export const useSqlEditorV2StateSnapshot = (options?: Parameters<typeof useSnaps
 async function upsertSnippet(id: string, projectRef: string, payload: UpsertContentPayloadV2) {
   try {
     sqlEditorState.savingStates[id] = 'UPDATING'
-    await upsertContent({
-      projectRef,
-      payload,
-    })
+    await upsertContent({ projectRef, payload })
+
+    const queryClient = getQueryClient()
+    await queryClient.invalidateQueries(contentKeys.count(projectRef, 'sql'))
+
     sqlEditorState.savingStates[id] = 'IDLE'
   } catch (error) {
     sqlEditorState.savingStates[id] = 'UPDATING_FAILED'

packages/api-types/types/api.d.ts

@@ -795,6 +795,10 @@ export interface paths {
      */
     patch: operations['ContentController_updateContent']
   }
+  '/platform/projects/{ref}/content/count': {
+    /** Gets the count of a user's content by type */
+    get: operations['ContentController_getContentCount']
+  }
   '/platform/projects/{ref}/content/folders': {
     /** Gets project's content root folder */
     get: operations['ContentFoldersController_getRootFolder']
@@ -1626,6 +1630,10 @@ export interface paths {
      */
     patch: operations['ContentController_updateContent']
   }
+  '/v0/projects/{ref}/content/count': {
+    /** Gets the count of a user's content by type */
+    get: operations['ContentController_getContentCount']
+  }
   '/v0/projects/{ref}/content/item/{id}': {
     /** Gets project's content by the given id */
     get: operations['ContentController_getContentById']
@@ -1868,7 +1876,7 @@ export interface paths {
   '/v1/projects/{ref}/config/auth/third-party-auth': {
     /** [Alpha] Lists all third-party auth integrations */
     get: operations['ThirdPartyAuthController_listTPAForProject']
-    /** [Alpha] Creates a new third-party auth integration */
+    /** Creates a new third-party auth integration */
     post: operations['ThirdPartyAuthController_createTPAForProject']
   }
   '/v1/projects/{ref}/config/auth/third-party-auth/{tpa_id}': {
@@ -2283,6 +2291,8 @@ export interface components {
       mailer_templates_reauthentication_content: string | null
       mailer_templates_recovery_content: string | null
       mfa_max_enrolled_factors: number | null
+      mfa_totp_enroll_enabled: boolean | null
+      mfa_totp_verify_enabled: boolean | null
       password_hibp_enabled: boolean | null
       password_min_length: number | null
       password_required_characters: string | null
@@ -3134,6 +3144,9 @@ export interface components {
       verify_jwt?: boolean
       version: number
     }
+    GetContentCountResponse: {
+      count: number
+    }
     GetMetricsBody: {
       /** @enum {string} */
       interval: '1d' | '3d' | '7d'
@@ -3452,6 +3465,8 @@ export interface components {
       MAILER_TEMPLATES_REAUTHENTICATION_CONTENT: string
       MAILER_TEMPLATES_RECOVERY_CONTENT: string
       MFA_MAX_ENROLLED_FACTORS: number
+      MFA_TOTP_ENROLL_ENABLED: boolean
+      MFA_TOTP_VERIFY_ENABLED: boolean
       PASSWORD_HIBP_ENABLED: boolean
       PASSWORD_MIN_LENGTH: number
       PASSWORD_REQUIRED_CHARACTERS: string
@@ -4451,6 +4466,7 @@ export interface components {
         | 'rls_references_user_metadata'
         | 'materialized_view_in_api'
         | 'foreign_table_in_api'
+        | 'unsupported_reg_types'
         | 'auth_otp_long_expiry'
         | 'auth_otp_short_length'
       remediation: Record<string, never>
@@ -5407,6 +5423,8 @@ export interface components {
       mailer_templates_reauthentication_content?: string
       mailer_templates_recovery_content?: string
       mfa_max_enrolled_factors?: number
+      mfa_totp_enroll_enabled?: boolean
+      mfa_totp_verify_enabled?: boolean
       password_hibp_enabled?: boolean
       password_min_length?: number
       /** @enum {string} */
@@ -5647,6 +5665,8 @@ export interface components {
       MAILER_TEMPLATES_REAUTHENTICATION_CONTENT?: string
       MAILER_TEMPLATES_RECOVERY_CONTENT?: string
       MFA_MAX_ENROLLED_FACTORS?: number
+      MFA_TOTP_ENROLL_ENABLED?: boolean
+      MFA_TOTP_VERIFY_ENABLED?: boolean
       PASSWORD_HIBP_ENABLED?: boolean
       PASSWORD_MIN_LENGTH?: number
       /** @enum {string} */
@@ -11649,6 +11669,29 @@ export interface operations {
       }
     }
   }
+  /** Gets the count of a user's content by type */
+  ContentController_getContentCount: {
+    parameters: {
+      query: {
+        type: string
+      }
+      path: {
+        /** @description Project ref */
+        ref: string
+      }
+    }
+    responses: {
+      200: {
+        content: {
+          'application/json': components['schemas']['GetContentCountResponse']
+        }
+      }
+      /** @description Failed to retrieve user's content count */
+      500: {
+        content: never
+      }
+    }
+  }
   /** Gets project's content root folder */
   ContentFoldersController_getRootFolder: {
     parameters: {
@@ -14249,7 +14292,7 @@ export interface operations {
       }
     }
   }
-  /** [Alpha] Creates a new third-party auth integration */
+  /** Creates a new third-party auth integration */
   ThirdPartyAuthController_createTPAForProject: {
     parameters: {
       path: {