Attempt

apps/studio/data/sql/execute-sql.ts

@@ -1,12 +1,23 @@
 import { DEFAULT_PLATFORM_APPLICATION_NAME } from '@supabase/pg-meta/src/constants'
 import { handleError as handleErrorFetchers, post } from 'data/fetchers'
 import { MB } from 'lib/constants'
-import {
-  ROLE_IMPERSONATION_NO_RESULTS,
-  ROLE_IMPERSONATION_SQL_LINE_COUNT,
-} from 'lib/role-impersonation'
 import { ResponseError } from 'types'
-import { ExecuteSqlVariables } from './execute-sql-query'
+
+// [Joshen] If this works, we'll need to clean up
+const ROLE_IMPERSONATION_SQL_LINE_COUNT = 11
+const ROLE_IMPERSONATION_NO_RESULTS = 'ROLE_IMPERSONATION_NO_RESULTS'
+
+type ExecuteSqlVariables = {
+  projectRef?: string
+  connectionString?: string | null
+  sql: string
+  queryKey?: any
+  handleError?: (error: ResponseError) => { result: any }
+  isRoleImpersonationEnabled?: boolean
+  isStatementTimeoutDisabled?: boolean
+  autoLimit?: number
+  contextualInvalidation?: boolean
+}
 
 export async function executeSql<T = any>(
   {
@@ -70,7 +81,7 @@ export async function executeSql<T = any>(
         query: {
           key:
             queryKey
-              ?.filter((seg) => typeof seg === 'string' || typeof seg === 'number')
+              ?.filter((seg: any) => typeof seg === 'string' || typeof seg === 'number')
               .join('-') ?? '',
         },
       },

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

@@ -1,6 +1,85 @@
-import { getOrganizations } from 'data/organizations/organizations-query'
-import { getProjects } from 'data/projects/projects-query'
-import { getAiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
+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'
+  }
+}
 
 export const getOrgAIDetails = async ({
   orgSlug,

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

@@ -6,13 +6,8 @@ import { z } from 'zod/v4'
 
 import { IS_PLATFORM } from 'common'
 import { executeSql } from 'data/sql/execute-sql'
-import { AiOptInLevel } from 'hooks/misc/useOrgOptedIntoAi'
 import { getModel } from 'lib/ai/model'
-import { getOrgAIDetails } from 'lib/ai/org-ai-details'
-import { getTools } from 'lib/ai/tools'
-import apiWrapper from 'lib/api/apiWrapper'
-import { queryPgMetaSelfHosted } from 'lib/self-hosted'
-
+import { AiOptInLevel, getOrgAIDetails } from 'lib/ai/org-ai-details'
 import {
   CHAT_PROMPT,
   EDGE_FUNCTION_PROMPT,
@@ -21,6 +16,9 @@ import {
   RLS_PROMPT,
   SECURITY_PROMPT,
 } from 'lib/ai/prompts'
+import { getTools } from 'lib/ai/tools'
+import apiWrapper from 'lib/api/apiWrapper'
+import { queryPgMetaSelfHosted } from 'lib/self-hosted'
 
 const requestBodySchema = z.object({
   messages: z.array(z.any()),