import { createLogger } from '@sim/logger' import { keepPreviousData, useMutation, useQuery, useQueryClient } from '@tanstack/react-query' import { requestJson } from '@/lib/api/client/request' import { type ContractBodyInput, removeWorkspaceEnvironmentContract, savePersonalEnvironmentContract, upsertWorkspaceEnvironmentContract, } from '@/lib/api/contracts' import type { WorkspaceEnvironmentData } from '@/lib/environment/api' import { fetchPersonalEnvironment, fetchWorkspaceEnvironment } from '@/lib/environment/api' const logger = createLogger('EnvironmentQueries') /** * Query key factories for environment variable queries */ export const PERSONAL_ENVIRONMENT_STALE_TIME = 60 * 1000 export const WORKSPACE_ENVIRONMENT_STALE_TIME = 60 * 1000 export const environmentKeys = { all: ['environment'] as const, personal: () => [...environmentKeys.all, 'personal'] as const, workspaces: () => [...environmentKeys.all, 'workspace'] as const, workspace: (workspaceId: string) => [...environmentKeys.workspaces(), workspaceId] as const, } /** * Hook to fetch personal environment variables */ export function usePersonalEnvironment() { return useQuery({ queryKey: environmentKeys.personal(), queryFn: ({ signal }) => fetchPersonalEnvironment(signal), staleTime: PERSONAL_ENVIRONMENT_STALE_TIME, }) } /** * Hook to fetch workspace environment variables */ export function useWorkspaceEnvironment( workspaceId: string, options?: { select?: (data: WorkspaceEnvironmentData) => TData } ) { return useQuery({ queryKey: environmentKeys.workspace(workspaceId), queryFn: ({ signal }) => fetchWorkspaceEnvironment(workspaceId, signal), enabled: !!workspaceId, staleTime: WORKSPACE_ENVIRONMENT_STALE_TIME, placeholderData: keepPreviousData, ...options, }) } /** * Save personal environment variables mutation */ type SavePersonalEnvironmentParams = ContractBodyInput export function useSavePersonalEnvironment() { const queryClient = useQueryClient() return useMutation({ mutationFn: async ({ variables }: SavePersonalEnvironmentParams) => { await requestJson(savePersonalEnvironmentContract, { body: { variables } }) logger.info('Saved personal environment variables') }, onSettled: async () => { await Promise.all([ queryClient.invalidateQueries({ queryKey: environmentKeys.personal() }), queryClient.invalidateQueries({ queryKey: environmentKeys.workspaces() }), ]) }, }) } /** * Upsert workspace environment variables mutation */ type UpsertWorkspaceEnvironmentParams = { workspaceId: string } & ContractBodyInput< typeof upsertWorkspaceEnvironmentContract > export function useUpsertWorkspaceEnvironment() { const queryClient = useQueryClient() return useMutation({ mutationFn: async ({ workspaceId, variables }: UpsertWorkspaceEnvironmentParams) => { const data = await requestJson(upsertWorkspaceEnvironmentContract, { params: { id: workspaceId }, body: { variables }, }) logger.info(`Upserted workspace environment variables for workspace: ${workspaceId}`) return data }, onSettled: (_data, _error, variables) => queryClient.invalidateQueries({ queryKey: environmentKeys.workspace(variables.workspaceId), }), }) } /** * Remove workspace environment variables mutation */ type RemoveWorkspaceEnvironmentParams = { workspaceId: string } & ContractBodyInput< typeof removeWorkspaceEnvironmentContract > export function useRemoveWorkspaceEnvironment() { const queryClient = useQueryClient() return useMutation({ mutationFn: async ({ workspaceId, keys }: RemoveWorkspaceEnvironmentParams) => { const data = await requestJson(removeWorkspaceEnvironmentContract, { params: { id: workspaceId }, body: { keys }, }) logger.info(`Removed ${keys.length} workspace environment keys for workspace: ${workspaceId}`) return data }, onSettled: (_data, _error, variables) => queryClient.invalidateQueries({ queryKey: environmentKeys.workspace(variables.workspaceId), }), }) }