import { organizationQueryKeys } from "@better-auth-ui/core/plugins" import { type DataTag, type QueryClient, queryOptions, skipToken, useQuery } from "@tanstack/react-query" import type { BetterFetchError } from "better-auth/react" import type { InferData, OrganizationAuthClient } from "../../lib/auth-client" import { useSession } from "../auth/session-query" import { useActiveOrganization } from "./active-organization-query" export type HasPermissionData< TAuthClient extends OrganizationAuthClient = OrganizationAuthClient > = InferData export type HasPermissionParams< TAuthClient extends OrganizationAuthClient = OrganizationAuthClient > = Parameters[0] export type HasPermissionOptions< TAuthClient extends OrganizationAuthClient = OrganizationAuthClient > = Omit< ReturnType>, "queryKey" | "queryFn" > export function hasPermissionOptions< TAuthClient extends OrganizationAuthClient >( authClient: TAuthClient, userId: string | undefined, params: HasPermissionParams ) { type TData = HasPermissionData // `hasPermission` is the only org client method without a `query` field — its // params are flat — so the "query" portion is everything except `fetchOptions`. const { fetchOptions, ...query } = params const queryKey = organizationQueryKeys.permissions.has(userId, query) const options = queryOptions( { queryKey, queryFn: ({ signal }) => authClient.organization.hasPermission({ ...query, fetchOptions: { ...fetchOptions, signal, throw: true } }) as Promise } ) return options as typeof options & { queryKey: DataTag } } export const ensureHasPermission = ( queryClient: QueryClient, authClient: TAuthClient, userId: string | undefined, params: HasPermissionParams ) => queryClient.ensureQueryData(hasPermissionOptions(authClient, userId, params)) export const prefetchHasPermission = < TAuthClient extends OrganizationAuthClient >( queryClient: QueryClient, authClient: TAuthClient, userId: string | undefined, params: HasPermissionParams ) => queryClient.prefetchQuery(hasPermissionOptions(authClient, userId, params)) export const fetchHasPermission = ( queryClient: QueryClient, authClient: TAuthClient, userId: string | undefined, params: HasPermissionParams ) => queryClient.fetchQuery(hasPermissionOptions(authClient, userId, params)) export type UseHasPermissionOptions< TAuthClient extends OrganizationAuthClient = OrganizationAuthClient > = HasPermissionOptions & HasPermissionParams export function useHasPermission( authClient: TAuthClient, options: UseHasPermissionOptions, queryClient?: QueryClient ) { const { data: session } = useSession(authClient, undefined, queryClient) const userId = session?.user.id const { fetchOptions, permissions, organizationId: optionsOrganizationId, ...queryOptions } = options const { data: activeOrganization } = useActiveOrganization( authClient, { enabled: !optionsOrganizationId }, queryClient ) const organizationId = optionsOrganizationId ?? activeOrganization?.id const baseOptions = hasPermissionOptions(authClient, userId, { fetchOptions, organizationId, permissions }) return useQuery( { ...queryOptions, ...baseOptions, queryFn: userId && organizationId ? baseOptions.queryFn : skipToken }, queryClient ) }