util.ts91 lines · main
1import chalk from 'chalk'
2import type { Code } from 'mdast'
3import { fromMarkdown } from 'mdast-util-from-markdown'
4import { format } from 'sql-formatter'
5
6declare global {
7 interface ReadableStream<R = any> {
8 [Symbol.asyncIterator](): AsyncIterableIterator<R>
9 }
10}
11
12/**
13 * Formats Postgres SQL into a consistent format.
14 *
15 * @returns The formatted SQL.
16 */
17export const formatSql = (sql: string) =>
18 format(sql, { language: 'postgresql', keywordCase: 'lower' })
19
20/**
21 * Collects an `ArrayBuffer` stream into a single decoded string.
22 *
23 * @returns A single string combining all the decoded stream chunks.
24 */
25export async function collectStream<R extends BufferSource>(stream: ReadableStream<R>) {
26 const textDecoderStream = new TextDecoderStream()
27
28 let content = ''
29
30 for await (const chunk of stream.pipeThrough(textDecoderStream)) {
31 const text = chunk.split('0:')[1]
32 content += text.slice(1, text.length - 2)
33 }
34
35 return content.replaceAll('\\n', '\n').replaceAll('\\"', '"')
36}
37
38/**
39 * Parses markdown and extracts all SQL code blocks.
40 *
41 * @returns An array of string content from each SQL code block.
42 */
43export function extractMarkdownSql(markdown: string) {
44 const mdTree = fromMarkdown(markdown)
45
46 return mdTree.children
47 .filter((node): node is Code => node.type === 'code' && node.lang === 'sql')
48 .map(({ value }) => value)
49}
50
51/**
52 * Prints the provided metadata along with any assertion errors.
53 * Works both synchronously and asynchronously.
54 *
55 * Useful for providing extra context for failed tests.
56 */
57export function withMetadata<T extends void | Promise<void>>(
58 metadata: Record<string, string>,
59 fn: () => T
60): T {
61 /**
62 * Prepends metadata to an Error's stack trace.
63 */
64 function modifyError(err: unknown) {
65 if (err instanceof Error && err.stack) {
66 const formattedMetadata = Object.entries(metadata).map(
67 ([key, value]) => `${chalk.bold.dim(key)}:\n\n${chalk.green.dim(value)}`
68 )
69 err.stack = `${formattedMetadata.join('\n\n')}\n\n${err.stack}`
70 }
71
72 return err
73 }
74
75 // Execute the function and handle both
76 // synchronous or asynchronous scenarios
77 try {
78 const maybePromise = fn()
79
80 if (maybePromise instanceof Promise) {
81 return maybePromise.catch((err) => {
82 // Re-throw the error
83 throw modifyError(err)
84 }) as T
85 }
86 return maybePromise
87 } catch (err) {
88 // Re-throw the error
89 throw modifyError(err)
90 }
91}