From 107d4f1c86cc8d0dae9d930637519df4241a2327 Mon Sep 17 00:00:00 2001 From: Joshen Lim Date: Thu, 18 Sep 2025 20:36:17 +0800 Subject: [PATCH] Attempt --- apps/studio/data/sql/execute-sql.ts | 23 ++++-- apps/studio/lib/ai/org-ai-details.ts | 85 ++++++++++++++++++++- apps/studio/pages/api/ai/sql/generate-v4.ts | 10 +-- 3 files changed, 103 insertions(+), 15 deletions(-) diff --git a/apps/studio/data/sql/execute-sql.ts b/apps/studio/data/sql/execute-sql.ts index 12ac175c4e8..36f3391d5d6 100644 --- a/apps/studio/data/sql/execute-sql.ts +++ b/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( { @@ -70,7 +81,7 @@ export async function executeSql( query: { key: queryKey - ?.filter((seg) => typeof seg === 'string' || typeof seg === 'number') + ?.filter((seg: any) => typeof seg === 'string' || typeof seg === 'number') .join('-') ?? '', }, }, diff --git a/apps/studio/lib/ai/org-ai-details.ts b/apps/studio/lib/ai/org-ai-details.ts index 4c04eaa2fe2..8757cfb9a35 100644 --- a/apps/studio/lib/ai/org-ai-details.ts +++ b/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 +}): Promise { + 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 +}) { + 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, diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index da0df8aaae8..eff64418df0 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/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()),