Clean up

apps/studio/data/organizations/fetchers.ts

@@ -0,0 +1,49 @@
+// [Joshen] These are being placed separately as they're also being used in the API
+// which we should avoid mixing client side and server side logic (main problem was importing of react query)
+
+import { components } from 'api-types'
+import { get, handleError } from 'data/fetchers'
+import { MANAGED_BY, ManagedBy } from 'lib/constants'
+import { Organization } from 'types'
+
+export type OrganizationBase = components['schemas']['OrganizationResponse']
+
+function getManagedBy(org: OrganizationBase): ManagedBy {
+  switch (org.billing_partner) {
+    case 'vercel_marketplace':
+      return MANAGED_BY.VERCEL_MARKETPLACE
+    // TODO(ignacio): Uncomment this when we've deployed the AWS Marketplace new slug
+    // case 'aws_marketplace':
+    // return MANAGED_BY.AWS_MARKETPLACE
+    case 'aws':
+      return MANAGED_BY.AWS_MARKETPLACE
+    default:
+      return MANAGED_BY.SUPABASE
+  }
+}
+
+export function castOrganizationResponseToOrganization(org: OrganizationBase): Organization {
+  return {
+    ...org,
+    billing_email: org.billing_email ?? 'Unknown',
+    managed_by: getManagedBy(org),
+    partner_id: org.slug.startsWith('vercel_') ? org.slug.replace('vercel_', '') : undefined,
+  }
+}
+
+export async function getOrganizations({
+  signal,
+  headers,
+}: {
+  signal?: AbortSignal
+  headers?: Record<string, string>
+}): Promise<Organization[]> {
+  const { data, error } = await get('/platform/organizations', { signal, headers })
+
+  if (error) handleError(error)
+  if (!Array.isArray(data)) return []
+
+  return data
+    .map(castOrganizationResponseToOrganization)
+    .sort((a, b) => a.name.localeCompare(b.name))
+}

apps/studio/data/organizations/organization-create-mutation.ts

@@ -4,8 +4,8 @@ import { toast } from 'sonner'
 import { handleError, post } from 'data/fetchers'
 import { permissionKeys } from 'data/permissions/keys'
 import type { ResponseError } from 'types'
+import { castOrganizationResponseToOrganization } from './fetchers'
 import { organizationKeys } from './keys'
-import { castOrganizationResponseToOrganization } from './organizations-query'
 import type { CustomerAddress, CustomerTaxId } from './types'
 
 export type OrganizationCreateVariables = {

apps/studio/data/organizations/organizations-query.ts

@@ -1,53 +1,9 @@
 import { QueryClient, useQuery, UseQueryOptions } from '@tanstack/react-query'
 
-import { components } from 'api-types'
-import { get, handleError } from 'data/fetchers'
 import { useProfile } from 'lib/profile'
-import type { Organization, ResponseError } from 'types'
+import type { ResponseError } from 'types'
+import { getOrganizations } from './fetchers'
 import { organizationKeys } from './keys'
-import { MANAGED_BY, ManagedBy } from 'lib/constants/infrastructure'
-
-export type OrganizationBase = components['schemas']['OrganizationResponse']
-
-export function castOrganizationResponseToOrganization(org: OrganizationBase): Organization {
-  return {
-    ...org,
-    billing_email: org.billing_email ?? 'Unknown',
-    managed_by: getManagedBy(org),
-    partner_id: org.slug.startsWith('vercel_') ? org.slug.replace('vercel_', '') : undefined,
-  }
-}
-
-function getManagedBy(org: OrganizationBase): ManagedBy {
-  switch (org.billing_partner) {
-    case 'vercel_marketplace':
-      return MANAGED_BY.VERCEL_MARKETPLACE
-    // TODO(ignacio): Uncomment this when we've deployed the AWS Marketplace new slug
-    // case 'aws_marketplace':
-    // return MANAGED_BY.AWS_MARKETPLACE
-    case 'aws':
-      return MANAGED_BY.AWS_MARKETPLACE
-    default:
-      return MANAGED_BY.SUPABASE
-  }
-}
-
-export async function getOrganizations({
-  signal,
-  headers,
-}: {
-  signal?: AbortSignal
-  headers?: Record<string, string>
-}): Promise<Organization[]> {
-  const { data, error } = await get('/platform/organizations', { signal, headers })
-
-  if (error) handleError(error)
-  if (!Array.isArray(data)) return []
-
-  return data
-    .map(castOrganizationResponseToOrganization)
-    .sort((a, b) => a.name.localeCompare(b.name))
-}
 
 export type OrganizationsData = Awaited<ReturnType<typeof getOrganizations>>
 export type OrganizationsError = ResponseError

apps/studio/data/projects/fetchers.ts

@@ -0,0 +1,20 @@
+// [Joshen] These are being placed separately as they're also being used in the API
+// which we should avoid mixing client side and server side logic (main problem was importing of react query)
+
+import { get, handleError } from 'data/fetchers'
+
+export async function getProjects({
+  signal,
+  headers,
+}: {
+  signal?: AbortSignal
+  headers?: Record<string, string>
+}) {
+  const { data, error } = await get('/platform/projects', { signal, headers })
+
+  if (error) handleError(error)
+  // The /platform/projects endpoint has a v2 which is activated by passing a {version: '2'} header. The v1 API returns
+  // all projects while the v2 returns paginated list of projects. Wrapping the v1 API response into a
+  // { projects: ProjectInfo[] } is intentional to be forward compatible with the structure of v2 for easier migration.
+  return { projects: data }
+}

apps/studio/data/projects/projects-query.ts

@@ -2,9 +2,9 @@ import { QueryClient, useQuery, useQueryClient, UseQueryOptions } from '@tanstac
 import { useCallback, useRef } from 'react'
 
 import type { components } from 'data/api'
-import { get, handleError } from 'data/fetchers'
 import { useProfile } from 'lib/profile'
 import type { ResponseError } from 'types'
+import { getProjects } from './fetchers'
 import { projectKeys } from './keys'
 import type { Project } from './project-detail-query'
 
@@ -15,22 +15,6 @@ export type ProjectsVariables = {
 type PaginatedProjectsResponse = components['schemas']['ListProjectsPaginatedResponse']
 export type ProjectInfo = PaginatedProjectsResponse['projects'][number]
 
-export async function getProjects({
-  signal,
-  headers,
-}: {
-  signal?: AbortSignal
-  headers?: Record<string, string>
-}) {
-  const { data, error } = await get('/platform/projects', { signal, headers })
-
-  if (error) handleError(error)
-  // The /platform/projects endpoint has a v2 which is activated by passing a {version: '2'} header. The v1 API returns
-  // all projects while the v2 returns paginated list of projects. Wrapping the v1 API response into a
-  // { projects: ProjectInfo[] } is intentional to be forward compatible with the structure of v2 for easier migration.
-  return { projects: data }
-}
-
 export type ProjectsData = Awaited<ReturnType<typeof getProjects>>
 export type ProjectsError = ResponseError
 

apps/studio/data/subscriptions/org-subscription-confirm-pending-create.ts

@@ -1,12 +1,12 @@
 import { useMutation, UseMutationOptions, useQueryClient } from '@tanstack/react-query'
 import { toast } from 'sonner'
 
+import type { components } from 'api-types'
 import { handleError, post } from 'data/fetchers'
-import type { ResponseError } from 'types'
+import { castOrganizationResponseToOrganization } from 'data/organizations/fetchers'
 import { organizationKeys } from 'data/organizations/keys'
 import { permissionKeys } from 'data/permissions/keys'
-import { castOrganizationResponseToOrganization } from 'data/organizations/organizations-query'
-import type { components } from 'api-types'
+import type { ResponseError } from 'types'
 
 export type PendingSubscriptionCreateVariables = {
   payment_intent_id: string

apps/studio/hooks/forms/useAIOptInForm.ts

@@ -8,9 +8,9 @@ import * as z from 'zod'
 import { LOCAL_STORAGE_KEYS } from 'common'
 import { useOrganizationUpdateMutation } from 'data/organizations/organization-update-mutation'
 import { invalidateOrganizationsQuery } from 'data/organizations/organizations-query'
+import { getAiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { useAsyncCheckPermissions } from 'hooks/misc/useCheckPermissions'
 import { useLocalStorageQuery } from 'hooks/misc/useLocalStorage'
-import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
 import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { OPT_IN_TAGS } from 'lib/constants'
 import type { ResponseError } from 'types'

apps/studio/hooks/misc/aiOptInLevel.ts

@@ -0,0 +1,27 @@
+import { OPT_IN_TAGS } from 'lib/constants'
+import { z } from 'zod'
+
+export const aiOptInLevelSchema = z.enum([
+  'disabled',
+  'schema',
+  'schema_and_log',
+  'schema_and_log_and_data',
+])
+
+export type AiOptInLevel = z.infer<typeof aiOptInLevelSchema>
+
+export const getAiOptInLevel = (tags: string[] | undefined): AiOptInLevel => {
+  const hasSql = tags?.includes(OPT_IN_TAGS.AI_SQL)
+  const hasData = tags?.includes(OPT_IN_TAGS.AI_DATA)
+  const hasLog = tags?.includes(OPT_IN_TAGS.AI_LOG)
+
+  if (hasData) {
+    return 'schema_and_log_and_data'
+  } else if (hasLog) {
+    return 'schema_and_log'
+  } else if (hasSql) {
+    return 'schema'
+  } else {
+    return 'disabled'
+  }
+}

apps/studio/hooks/misc/useOrgOptedIntoAi.ts

@@ -1,36 +1,10 @@
-import { z } from 'zod'
-
 import { subscriptionHasHipaaAddon } from 'components/interfaces/Billing/Subscription/Subscription.utils'
 import { useProjectSettingsV2Query } from 'data/config/project-settings-v2-query'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { useSelectedOrganizationQuery } from 'hooks/misc/useSelectedOrganization'
 import { useSelectedProjectQuery } from 'hooks/misc/useSelectedProject'
-import { IS_PLATFORM, OPT_IN_TAGS } from 'lib/constants'
-
-export const aiOptInLevelSchema = z.enum([
-  'disabled',
-  'schema',
-  'schema_and_log',
-  'schema_and_log_and_data',
-])
-
-export type AiOptInLevel = z.infer<typeof aiOptInLevelSchema>
-
-export const getAiOptInLevel = (tags: string[] | undefined): AiOptInLevel => {
-  const hasSql = tags?.includes(OPT_IN_TAGS.AI_SQL)
-  const hasData = tags?.includes(OPT_IN_TAGS.AI_DATA)
-  const hasLog = tags?.includes(OPT_IN_TAGS.AI_LOG)
-
-  if (hasData) {
-    return 'schema_and_log_and_data'
-  } else if (hasLog) {
-    return 'schema_and_log'
-  } else if (hasSql) {
-    return 'schema'
-  } else {
-    return 'disabled'
-  }
-}
+import { IS_PLATFORM } from 'lib/constants'
+import { AiOptInLevel, getAiOptInLevel } from './aiOptInLevel'
 
 /**
  * Determines if the organization has opted into *any* level of AI features (schema or schema_and_log or schema_and_log_and_data).

apps/studio/lib/ai/org-ai-details.ts

@@ -1,85 +1,6 @@
-import { components } from 'api-types'
-import { get, handleError } from 'data/fetchers'
-import { MANAGED_BY, ManagedBy, OPT_IN_TAGS } from 'lib/constants'
-import { Organization } from 'types'
-
-// [Joshen] If this works, we need to clean up
-
-export type AiOptInLevel = 'disabled' | 'schema' | 'schema_and_log' | 'schema_and_log_and_data'
-
-type OrganizationBase = components['schemas']['OrganizationResponse']
-
-function getManagedBy(org: OrganizationBase): ManagedBy {
-  switch (org.billing_partner) {
-    case 'vercel_marketplace':
-      return MANAGED_BY.VERCEL_MARKETPLACE
-    // TODO(ignacio): Uncomment this when we've deployed the AWS Marketplace new slug
-    // case 'aws_marketplace':
-    // return MANAGED_BY.AWS_MARKETPLACE
-    case 'aws':
-      return MANAGED_BY.AWS_MARKETPLACE
-    default:
-      return MANAGED_BY.SUPABASE
-  }
-}
-
-function castOrganizationResponseToOrganization(org: OrganizationBase): Organization {
-  return {
-    ...org,
-    billing_email: org.billing_email ?? 'Unknown',
-    managed_by: getManagedBy(org),
-    partner_id: org.slug.startsWith('vercel_') ? org.slug.replace('vercel_', '') : undefined,
-  }
-}
-
-async function getOrganizations({
-  signal,
-  headers,
-}: {
-  signal?: AbortSignal
-  headers?: Record<string, string>
-}): Promise<Organization[]> {
-  const { data, error } = await get('/platform/organizations', { signal, headers })
-
-  if (error) handleError(error)
-  if (!Array.isArray(data)) return []
-
-  return data
-    .map(castOrganizationResponseToOrganization)
-    .sort((a, b) => a.name.localeCompare(b.name))
-}
-
-async function getProjects({
-  signal,
-  headers,
-}: {
-  signal?: AbortSignal
-  headers?: Record<string, string>
-}) {
-  const { data, error } = await get('/platform/projects', { signal, headers })
-
-  if (error) handleError(error)
-  // The /platform/projects endpoint has a v2 which is activated by passing a {version: '2'} header. The v1 API returns
-  // all projects while the v2 returns paginated list of projects. Wrapping the v1 API response into a
-  // { projects: ProjectInfo[] } is intentional to be forward compatible with the structure of v2 for easier migration.
-  return { projects: data }
-}
-
-const getAiOptInLevel = (tags: string[] | undefined): AiOptInLevel => {
-  const hasSql = tags?.includes(OPT_IN_TAGS.AI_SQL)
-  const hasData = tags?.includes(OPT_IN_TAGS.AI_DATA)
-  const hasLog = tags?.includes(OPT_IN_TAGS.AI_LOG)
-
-  if (hasData) {
-    return 'schema_and_log_and_data'
-  } else if (hasLog) {
-    return 'schema_and_log'
-  } else if (hasSql) {
-    return 'schema'
-  } else {
-    return 'disabled'
-  }
-}
+import { getOrganizations } from 'data/organizations/fetchers'
+import { getProjects } from 'data/projects/fetchers'
+import { getAiOptInLevel } from 'hooks/misc/aiOptInLevel'
 
 export const getOrgAIDetails = async ({
   orgSlug,

apps/studio/lib/ai/tool-filter.ts

@@ -1,6 +1,6 @@
 import { Tool, ToolSet } from 'ai'
+import { AiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { z } from 'zod'
-import { AiOptInLevel } from './org-ai-details'
 
 // Add the DatabaseExtension type import
 export type DatabaseExtension = {

apps/studio/lib/ai/tools/index.ts

@@ -1,13 +1,11 @@
 import { ToolSet } from 'ai'
 
 import { IS_PLATFORM } from 'common'
-import { AiOptInLevel } from '../org-ai-details'
+import { AiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { filterToolsByOptInLevel } from '../tool-filter'
+import { getFallbackTools } from './fallback-tools'
 import { getMcpTools } from './mcp-tools'
 import { getRenderingTools } from './rendering-tools'
-
-import { getFallbackTools } from './fallback-tools'
-
 import { getSchemaTools } from './schema-tools'
 
 export const getTools = async ({

apps/studio/lib/ai/tools/mcp-tools.ts

@@ -1,4 +1,4 @@
-import { AiOptInLevel } from '../org-ai-details'
+import { AiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { createSupabaseMCPClient } from '../supabase-mcp'
 import { filterToolsByOptInLevel, toolSetValidationSchema } from '../tool-filter'
 

apps/studio/pages/api/ai/code/complete.ts

@@ -3,7 +3,7 @@ import { generateText, ModelMessage, stepCountIs } from 'ai'
 import { IS_PLATFORM } from 'common'
 import { source } from 'common-tags'
 import { executeSql } from 'data/sql/execute-sql'
-import { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+import { AiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { getModel } from 'lib/ai/model'
 import { getOrgAIDetails } from 'lib/ai/org-ai-details'
 import {

apps/studio/pages/api/ai/sql/generate-v4.ts

@@ -6,8 +6,9 @@ import { z } from 'zod/v4'
 
 import { IS_PLATFORM } from 'common'
 import { executeSql } from 'data/sql/execute-sql'
+import { AiOptInLevel } from 'hooks/misc/aiOptInLevel'
 import { getModel } from 'lib/ai/model'
-import { AiOptInLevel, getOrgAIDetails } from 'lib/ai/org-ai-details'
+import { getOrgAIDetails } from 'lib/ai/org-ai-details'
 import {
   CHAT_PROMPT,
   EDGE_FUNCTION_PROMPT,

apps/studio/types/base.ts

@@ -1,5 +1,5 @@
 import { PermissionAction } from '@supabase/shared-types/out/constants'
-import { OrganizationBase } from 'data/organizations/organizations-query'
+import { OrganizationBase } from 'data/organizations/fetchers'
 import { PlanId } from 'data/subscriptions/types'
 import jsonLogic from 'json-logic-js'
 import { ManagedBy } from 'lib/constants/infrastructure'