lookup-user.ts 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386
  1. import { Database, and, eq, sql } from "../src/drizzle/index.js"
  2. import { AuthTable } from "../src/schema/auth.sql.js"
  3. import { UserTable } from "../src/schema/user.sql.js"
  4. import {
  5. BillingTable,
  6. PaymentTable,
  7. SubscriptionTable,
  8. BlackPlans,
  9. UsageTable,
  10. LiteTable,
  11. } from "../src/schema/billing.sql.js"
  12. import { WorkspaceTable } from "../src/schema/workspace.sql.js"
  13. import { KeyTable } from "../src/schema/key.sql.js"
  14. import { BlackData } from "../src/black.js"
  15. import { centsToMicroCents } from "../src/util/price.js"
  16. import { getWeekBounds } from "../src/util/date.js"
  17. import { ModelTable } from "../src/schema/model.sql.js"
  18. // get input from command line
  19. const identifier = process.argv[2]
  20. const verbose = process.argv[process.argv.length - 1] === "-v"
  21. if (!identifier) {
  22. console.error("Usage: bun lookup-user.ts <email|workspaceID|apiKey> [-v]")
  23. process.exit(1)
  24. }
  25. // loop up by workspace ID
  26. if (identifier.startsWith("wrk_")) {
  27. await printWorkspace(identifier)
  28. }
  29. // lookup by API key ID
  30. else if (identifier.startsWith("key_")) {
  31. const key = await Database.use((tx) =>
  32. tx
  33. .select()
  34. .from(KeyTable)
  35. .where(eq(KeyTable.id, identifier))
  36. .then((rows) => rows[0]),
  37. )
  38. if (!key) {
  39. console.error("API key not found")
  40. process.exit(1)
  41. }
  42. await printWorkspace(key.workspaceID)
  43. }
  44. // lookup by API key value
  45. else if (identifier.startsWith("sk-")) {
  46. const key = await Database.use((tx) =>
  47. tx
  48. .select()
  49. .from(KeyTable)
  50. .where(eq(KeyTable.key, identifier))
  51. .then((rows) => rows[0]),
  52. )
  53. if (!key) {
  54. console.error("API key not found")
  55. process.exit(1)
  56. }
  57. await printWorkspace(key.workspaceID)
  58. }
  59. // lookup by email
  60. else {
  61. const authData = await Database.use(async (tx) =>
  62. tx.select().from(AuthTable).where(eq(AuthTable.subject, identifier)),
  63. )
  64. if (authData.length === 0) {
  65. console.error("Email not found")
  66. process.exit(1)
  67. }
  68. if (authData.length > 1) console.warn("Multiple users found for email", identifier)
  69. // Get all auth records for email
  70. const accountID = authData[0].accountID
  71. await printTable("Auth", (tx) => tx.select().from(AuthTable).where(eq(AuthTable.accountID, accountID)))
  72. // Get all workspaces for this account
  73. const users = await printTable("Workspaces", (tx) =>
  74. tx
  75. .select({
  76. userID: UserTable.id,
  77. workspaceID: UserTable.workspaceID,
  78. workspaceName: WorkspaceTable.name,
  79. role: UserTable.role,
  80. black: SubscriptionTable.timeCreated,
  81. lite: LiteTable.timeCreated,
  82. })
  83. .from(UserTable)
  84. .rightJoin(WorkspaceTable, eq(WorkspaceTable.id, UserTable.workspaceID))
  85. .leftJoin(SubscriptionTable, eq(SubscriptionTable.userID, UserTable.id))
  86. .leftJoin(LiteTable, eq(LiteTable.userID, UserTable.id))
  87. .where(eq(UserTable.accountID, accountID))
  88. .then((rows) =>
  89. rows.map((row) => ({
  90. userID: row.userID,
  91. workspaceID: row.workspaceID,
  92. workspaceName: row.workspaceName,
  93. role: row.role,
  94. black: formatDate(row.black),
  95. lite: formatDate(row.lite),
  96. })),
  97. ),
  98. )
  99. for (const user of users) {
  100. await printWorkspace(user.workspaceID)
  101. }
  102. }
  103. async function printWorkspace(workspaceID: string) {
  104. const workspace = await Database.use((tx) =>
  105. tx
  106. .select()
  107. .from(WorkspaceTable)
  108. .where(eq(WorkspaceTable.id, workspaceID))
  109. .then((rows) => rows[0]),
  110. )
  111. printHeader(`Workspace "${workspace.name}" (${workspace.id})`)
  112. await printTable("Users", (tx) =>
  113. tx
  114. .select({
  115. authEmail: AuthTable.subject,
  116. inviteEmail: UserTable.email,
  117. role: UserTable.role,
  118. timeSeen: UserTable.timeSeen,
  119. monthlyLimit: UserTable.monthlyLimit,
  120. monthlyUsage: UserTable.monthlyUsage,
  121. timeDeleted: UserTable.timeDeleted,
  122. fixedUsage: SubscriptionTable.fixedUsage,
  123. rollingUsage: SubscriptionTable.rollingUsage,
  124. timeFixedUpdated: SubscriptionTable.timeFixedUpdated,
  125. timeRollingUpdated: SubscriptionTable.timeRollingUpdated,
  126. timeSubscriptionCreated: SubscriptionTable.timeCreated,
  127. subscription: BillingTable.subscription,
  128. })
  129. .from(UserTable)
  130. .innerJoin(BillingTable, eq(BillingTable.workspaceID, workspace.id))
  131. .leftJoin(AuthTable, and(eq(UserTable.accountID, AuthTable.accountID), eq(AuthTable.provider, "email")))
  132. .leftJoin(SubscriptionTable, eq(SubscriptionTable.userID, UserTable.id))
  133. .where(eq(UserTable.workspaceID, workspace.id))
  134. .then((rows) =>
  135. rows.map((row) => {
  136. const subStatus = getSubscriptionStatus(row)
  137. return {
  138. email: (row.timeDeleted ? "❌ " : "") + (row.authEmail ?? row.inviteEmail),
  139. role: row.role,
  140. timeSeen: formatDate(row.timeSeen),
  141. monthly: formatMonthlyUsage(row.monthlyUsage, row.monthlyLimit),
  142. subscribed: formatDate(row.timeSubscriptionCreated),
  143. subWeekly: subStatus.weekly,
  144. subRolling: subStatus.rolling,
  145. rateLimited: subStatus.rateLimited,
  146. retryIn: subStatus.retryIn,
  147. }
  148. }),
  149. ),
  150. )
  151. await printTable("Billing", (tx) =>
  152. tx
  153. .select({
  154. balance: BillingTable.balance,
  155. customerID: BillingTable.customerID,
  156. reload: BillingTable.reload,
  157. blackSubscriptionID: BillingTable.subscriptionID,
  158. blackSubscription: {
  159. plan: BillingTable.subscriptionPlan,
  160. booked: BillingTable.timeSubscriptionBooked,
  161. enrichment: BillingTable.subscription,
  162. },
  163. timeBlackSubscriptionSelected: BillingTable.timeSubscriptionSelected,
  164. liteSubscriptionID: BillingTable.liteSubscriptionID,
  165. })
  166. .from(BillingTable)
  167. .where(eq(BillingTable.workspaceID, workspace.id))
  168. .then(
  169. (rows) =>
  170. rows.map((row) => ({
  171. balance: `$${(row.balance / 100000000).toFixed(2)}`,
  172. reload: row.reload ? "yes" : "no",
  173. customerID: row.customerID,
  174. GO: row.liteSubscriptionID,
  175. Black: row.blackSubscriptionID
  176. ? [
  177. `Black ${row.blackSubscription.enrichment!.plan}`,
  178. row.blackSubscription.enrichment!.seats > 1
  179. ? `X ${row.blackSubscription.enrichment!.seats} seats`
  180. : "",
  181. row.blackSubscription.enrichment!.coupon
  182. ? `(coupon: ${row.blackSubscription.enrichment!.coupon})`
  183. : "",
  184. `(ref: ${row.blackSubscriptionID})`,
  185. ].join(" ")
  186. : row.blackSubscription.booked
  187. ? `Waitlist ${row.blackSubscription.plan} plan${row.timeBlackSubscriptionSelected ? " (selected)" : ""}`
  188. : undefined,
  189. }))[0],
  190. ),
  191. )
  192. await printTable("Payments", (tx) =>
  193. tx
  194. .select({
  195. amount: PaymentTable.amount,
  196. paymentID: PaymentTable.paymentID,
  197. invoiceID: PaymentTable.invoiceID,
  198. customerID: PaymentTable.customerID,
  199. timeCreated: PaymentTable.timeCreated,
  200. timeRefunded: PaymentTable.timeRefunded,
  201. })
  202. .from(PaymentTable)
  203. .where(eq(PaymentTable.workspaceID, workspace.id))
  204. .orderBy(sql`${PaymentTable.timeCreated} DESC`)
  205. .limit(100)
  206. .then((rows) =>
  207. rows.map((row) => ({
  208. ...row,
  209. amount: `$${(row.amount / 100000000).toFixed(2)}`,
  210. paymentID: row.paymentID
  211. ? `https://dashboard.stripe.com/acct_1RszBH2StuRr0lbX/payments/${row.paymentID}`
  212. : null,
  213. })),
  214. ),
  215. )
  216. if (verbose) {
  217. await printTable("28-Day Usage", (tx) =>
  218. tx
  219. .select({
  220. date: sql<string>`DATE(${UsageTable.timeCreated})`.as("date"),
  221. requests: sql<number>`COUNT(*)`.as("requests"),
  222. inputTokens: sql<number>`SUM(${UsageTable.inputTokens})`.as("input_tokens"),
  223. outputTokens: sql<number>`SUM(${UsageTable.outputTokens})`.as("output_tokens"),
  224. reasoningTokens: sql<number>`SUM(${UsageTable.reasoningTokens})`.as("reasoning_tokens"),
  225. cacheReadTokens: sql<number>`SUM(${UsageTable.cacheReadTokens})`.as("cache_read_tokens"),
  226. cacheWrite5mTokens: sql<number>`SUM(${UsageTable.cacheWrite5mTokens})`.as("cache_write_5m_tokens"),
  227. cacheWrite1hTokens: sql<number>`SUM(${UsageTable.cacheWrite1hTokens})`.as("cache_write_1h_tokens"),
  228. cost: sql<number>`SUM(${UsageTable.cost})`.as("cost"),
  229. })
  230. .from(UsageTable)
  231. .where(
  232. and(
  233. eq(UsageTable.workspaceID, workspace.id),
  234. sql`${UsageTable.timeCreated} >= DATE_SUB(NOW(), INTERVAL 28 DAY)`,
  235. ),
  236. )
  237. .groupBy(sql`DATE(${UsageTable.timeCreated})`)
  238. .orderBy(sql`DATE(${UsageTable.timeCreated}) DESC`)
  239. .then((rows) => {
  240. const totalCost = rows.reduce((sum, r) => sum + Number(r.cost), 0)
  241. const mapped = rows.map((row) => ({
  242. ...row,
  243. cost: `$${(Number(row.cost) / 100000000).toFixed(2)}`,
  244. }))
  245. if (mapped.length > 0) {
  246. mapped.push({
  247. date: "TOTAL",
  248. requests: null as any,
  249. inputTokens: null as any,
  250. outputTokens: null as any,
  251. reasoningTokens: null as any,
  252. cacheReadTokens: null as any,
  253. cacheWrite5mTokens: null as any,
  254. cacheWrite1hTokens: null as any,
  255. cost: `$${(totalCost / 100000000).toFixed(2)}`,
  256. })
  257. }
  258. return mapped
  259. }),
  260. )
  261. await printTable("Disabled Models", (tx) =>
  262. tx
  263. .select({
  264. model: ModelTable.model,
  265. timeCreated: ModelTable.timeCreated,
  266. })
  267. .from(ModelTable)
  268. .where(eq(ModelTable.workspaceID, workspace.id))
  269. .orderBy(sql`${ModelTable.timeCreated} DESC`)
  270. .then((rows) =>
  271. rows.map((row) => ({
  272. model: row.model,
  273. timeCreated: formatDate(row.timeCreated),
  274. })),
  275. ),
  276. )
  277. }
  278. }
  279. function formatMicroCents(value: number | null | undefined) {
  280. if (value === null || value === undefined) return null
  281. return `$${(value / 100000000).toFixed(2)}`
  282. }
  283. function formatDate(value: Date | null | undefined) {
  284. if (!value) return null
  285. return value.toISOString().split("T")[0]
  286. }
  287. function formatMonthlyUsage(usage: number | null | undefined, limit: number | null | undefined) {
  288. const usageText = formatMicroCents(usage) ?? "$0.00"
  289. if (limit === null || limit === undefined) return `${usageText} / no limit`
  290. return `${usageText} / $${limit.toFixed(2)}`
  291. }
  292. function formatRetryTime(seconds: number) {
  293. const days = Math.floor(seconds / 86400)
  294. if (days >= 1) return `${days} day${days > 1 ? "s" : ""}`
  295. const hours = Math.floor(seconds / 3600)
  296. const minutes = Math.ceil((seconds % 3600) / 60)
  297. if (hours >= 1) return `${hours}hr ${minutes}min`
  298. return `${minutes}min`
  299. }
  300. function getSubscriptionStatus(row: {
  301. subscription: {
  302. plan: (typeof BlackPlans)[number]
  303. } | null
  304. timeSubscriptionCreated: Date | null
  305. fixedUsage: number | null
  306. rollingUsage: number | null
  307. timeFixedUpdated: Date | null
  308. timeRollingUpdated: Date | null
  309. }) {
  310. if (!row.timeSubscriptionCreated || !row.subscription) {
  311. return { weekly: null, rolling: null, rateLimited: null, retryIn: null }
  312. }
  313. const black = BlackData.getLimits({ plan: row.subscription.plan })
  314. const now = new Date()
  315. const week = getWeekBounds(now)
  316. const fixedLimit = black.fixedLimit ? centsToMicroCents(black.fixedLimit * 100) : null
  317. const rollingLimit = black.rollingLimit ? centsToMicroCents(black.rollingLimit * 100) : null
  318. const rollingWindowMs = (black.rollingWindow ?? 5) * 3600 * 1000
  319. // Calculate current weekly usage (reset if outside current week)
  320. const currentWeekly =
  321. row.fixedUsage && row.timeFixedUpdated && row.timeFixedUpdated >= week.start ? row.fixedUsage : 0
  322. // Calculate current rolling usage
  323. const windowStart = new Date(now.getTime() - rollingWindowMs)
  324. const currentRolling =
  325. row.rollingUsage && row.timeRollingUpdated && row.timeRollingUpdated >= windowStart ? row.rollingUsage : 0
  326. // Check rate limiting
  327. const isWeeklyLimited = fixedLimit !== null && currentWeekly >= fixedLimit
  328. const isRollingLimited = rollingLimit !== null && currentRolling >= rollingLimit
  329. let retryIn: string | null = null
  330. if (isWeeklyLimited) {
  331. const retryAfter = Math.ceil((week.end.getTime() - now.getTime()) / 1000)
  332. retryIn = formatRetryTime(retryAfter)
  333. } else if (isRollingLimited && row.timeRollingUpdated) {
  334. const retryAfter = Math.ceil((row.timeRollingUpdated.getTime() + rollingWindowMs - now.getTime()) / 1000)
  335. retryIn = formatRetryTime(retryAfter)
  336. }
  337. return {
  338. weekly: fixedLimit !== null ? `${formatMicroCents(currentWeekly)} / $${black.fixedLimit}` : null,
  339. rolling: rollingLimit !== null ? `${formatMicroCents(currentRolling)} / $${black.rollingLimit}` : null,
  340. rateLimited: isWeeklyLimited || isRollingLimited ? "yes" : "no",
  341. retryIn,
  342. }
  343. }
  344. function printHeader(title: string) {
  345. console.log()
  346. console.log("─".repeat(title.length))
  347. console.log(`${title}`)
  348. console.log("─".repeat(title.length))
  349. }
  350. function printTable(title: string, callback: (tx: Database.TxOrDb) => Promise<any>): Promise<any> {
  351. return Database.use(async (tx) => {
  352. const data = await callback(tx)
  353. console.log(`\n== ${title} ==`)
  354. if (data.length === 0) {
  355. console.log("(no data)")
  356. } else {
  357. console.table(data)
  358. }
  359. return data
  360. })
  361. }