lookup-user.ts 12 KB

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