|
| 1 | +import { |
| 2 | + StarbaseApp, |
| 3 | + StarbaseContext, |
| 4 | + StarbaseDBConfiguration, |
| 5 | +} from '../../src/handler' |
| 6 | +import { StarbasePlugin } from '../../src/plugin' |
| 7 | +import { DataSource, QueryResult } from '../../src/types' |
| 8 | + |
| 9 | +const parser = new (require('node-sql-parser').Parser)() |
| 10 | + |
| 11 | +export class SqlMacrosPlugin extends StarbasePlugin { |
| 12 | + config?: StarbaseDBConfiguration |
| 13 | + |
| 14 | + // Prevents SQL statements with `SELECT *` from being executed |
| 15 | + preventSelectStar?: boolean |
| 16 | + |
| 17 | + constructor(opts?: { preventSelectStar: boolean }) { |
| 18 | + super('starbasedb:sql-macros') |
| 19 | + this.preventSelectStar = opts?.preventSelectStar |
| 20 | + } |
| 21 | + |
| 22 | + override async register(app: StarbaseApp) { |
| 23 | + app.use(async (c, next) => { |
| 24 | + this.config = c?.get('config') |
| 25 | + await next() |
| 26 | + }) |
| 27 | + } |
| 28 | + |
| 29 | + override async beforeQuery(opts: { |
| 30 | + sql: string |
| 31 | + params?: unknown[] |
| 32 | + dataSource?: DataSource |
| 33 | + config?: StarbaseDBConfiguration |
| 34 | + }): Promise<{ sql: string; params?: unknown[] }> { |
| 35 | + let { dataSource, sql, params } = opts |
| 36 | + |
| 37 | + // A data source is required for this plugin to operate successfully |
| 38 | + if (!dataSource) { |
| 39 | + return Promise.resolve({ |
| 40 | + sql, |
| 41 | + params, |
| 42 | + }) |
| 43 | + } |
| 44 | + |
| 45 | + sql = await this.replaceExcludeColumns(dataSource, sql, params) |
| 46 | + |
| 47 | + // Prevention of `SELECT *` statements is only enforced on non-admin users |
| 48 | + // Admins should be able to continue running these statements in database |
| 49 | + // tools such as Outerbase Studio. |
| 50 | + if (this.preventSelectStar && this.config?.role !== 'admin') { |
| 51 | + sql = this.checkSelectStar(sql, params) |
| 52 | + } |
| 53 | + |
| 54 | + return Promise.resolve({ |
| 55 | + sql, |
| 56 | + params, |
| 57 | + }) |
| 58 | + } |
| 59 | + |
| 60 | + private checkSelectStar(sql: string, params?: unknown[]): string { |
| 61 | + try { |
| 62 | + const ast = parser.astify(sql)[0] |
| 63 | + |
| 64 | + // Only check SELECT statements |
| 65 | + if (ast.type === 'select') { |
| 66 | + const hasSelectStar = ast.columns.some( |
| 67 | + (col: any) => |
| 68 | + col.expr.type === 'star' || |
| 69 | + (col.expr.type === 'column_ref' && |
| 70 | + col.expr.column === '*') |
| 71 | + ) |
| 72 | + |
| 73 | + if (hasSelectStar) { |
| 74 | + throw new Error( |
| 75 | + 'SELECT * is not allowed. Please specify explicit columns.' |
| 76 | + ) |
| 77 | + } |
| 78 | + } |
| 79 | + |
| 80 | + return sql |
| 81 | + } catch (error) { |
| 82 | + // If the error is our SELECT * error, rethrow it |
| 83 | + if ( |
| 84 | + error instanceof Error && |
| 85 | + error.message.includes('SELECT * is not allowed') |
| 86 | + ) { |
| 87 | + throw error |
| 88 | + } |
| 89 | + // For parsing errors or other issues, return original SQL |
| 90 | + return sql |
| 91 | + } |
| 92 | + } |
| 93 | + |
| 94 | + private async replaceExcludeColumns( |
| 95 | + dataSource: DataSource, |
| 96 | + sql: string, |
| 97 | + params?: unknown[] |
| 98 | + ): Promise<string> { |
| 99 | + // Only currently works for internal data source (Durable Object SQLite) |
| 100 | + if (dataSource.source !== 'internal') { |
| 101 | + return sql |
| 102 | + } |
| 103 | + |
| 104 | + // Special handling for pragma queries |
| 105 | + if (sql.toLowerCase().includes('pragma_table_info')) { |
| 106 | + return sql |
| 107 | + } |
| 108 | + |
| 109 | + try { |
| 110 | + // Add semicolon if missing |
| 111 | + const normalizedSql = sql.trim().endsWith(';') ? sql : `${sql};` |
| 112 | + |
| 113 | + // We allow users to write it `$_exclude` but convert it to `__exclude` so it can be |
| 114 | + // parsed with the AST library without throwing an error. |
| 115 | + const preparedSql = normalizedSql.replaceAll( |
| 116 | + '$_exclude', |
| 117 | + '__exclude' |
| 118 | + ) |
| 119 | + const normalizedQuery = parser.astify(preparedSql)[0] |
| 120 | + |
| 121 | + // Only process SELECT statements |
| 122 | + if (normalizedQuery.type !== 'select') { |
| 123 | + return sql |
| 124 | + } |
| 125 | + |
| 126 | + // Find any columns using `__exclude` |
| 127 | + const columns = normalizedQuery.columns |
| 128 | + const excludeFnIdx = columns.findIndex( |
| 129 | + (col: any) => |
| 130 | + col.expr && |
| 131 | + col.expr.type === 'function' && |
| 132 | + col.expr.name === '__exclude' |
| 133 | + ) |
| 134 | + |
| 135 | + if (excludeFnIdx === -1) { |
| 136 | + return sql |
| 137 | + } |
| 138 | + |
| 139 | + // Get the table name from the FROM clause |
| 140 | + const tableName = normalizedQuery.from[0].table |
| 141 | + let excludedColumns: string[] = [] |
| 142 | + |
| 143 | + try { |
| 144 | + const excludeExpr = normalizedQuery.columns[excludeFnIdx].expr |
| 145 | + |
| 146 | + // Handle both array and single argument cases |
| 147 | + const args = excludeExpr.args.value |
| 148 | + |
| 149 | + // Extract column name(s) from arguments |
| 150 | + excludedColumns = Array.isArray(args) |
| 151 | + ? args.map((arg: any) => arg.column) |
| 152 | + : [args.column] |
| 153 | + } catch (error: any) { |
| 154 | + console.error('Error processing exclude arguments:', error) |
| 155 | + console.error(error.stack) |
| 156 | + return sql |
| 157 | + } |
| 158 | + |
| 159 | + // Query database for all columns in this table |
| 160 | + // This only works for the internal SQLite data source |
| 161 | + const schemaQuery = ` |
| 162 | + SELECT name as column_name |
| 163 | + FROM pragma_table_info('${tableName}') |
| 164 | + ` |
| 165 | + |
| 166 | + const allColumns = (await dataSource?.rpc.executeQuery({ |
| 167 | + sql: schemaQuery, |
| 168 | + })) as QueryResult[] |
| 169 | + |
| 170 | + const includedColumns = allColumns |
| 171 | + .map((row: any) => row.column_name) |
| 172 | + .filter((col: string) => { |
| 173 | + const shouldInclude = !excludedColumns.includes( |
| 174 | + col.toLowerCase() |
| 175 | + ) |
| 176 | + return shouldInclude |
| 177 | + }) |
| 178 | + |
| 179 | + // Replace the __exclude function with explicit columns |
| 180 | + normalizedQuery.columns.splice( |
| 181 | + excludeFnIdx, |
| 182 | + 1, |
| 183 | + ...includedColumns.map((col: string) => ({ |
| 184 | + expr: { type: 'column_ref', table: null, column: col }, |
| 185 | + as: null, |
| 186 | + })) |
| 187 | + ) |
| 188 | + |
| 189 | + // Convert back to SQL and remove trailing semicolon to maintain original format |
| 190 | + return parser.sqlify(normalizedQuery).replace(/;$/, '') |
| 191 | + } catch (error) { |
| 192 | + console.error('SQL parsing error:', error) |
| 193 | + return sql |
| 194 | + } |
| 195 | + } |
| 196 | +} |
0 commit comments