table-row-query.ts266 lines · main
| 1 | import { |
| 2 | ident, |
| 3 | joinSqlFragments, |
| 4 | keyword, |
| 5 | literal, |
| 6 | safeSql, |
| 7 | type SafeSqlFragment, |
| 8 | } from '../pg-format' |
| 9 | import { PGForeignTable } from '../pg-meta-foreign-tables' |
| 10 | import { PGMaterializedView } from '../pg-meta-materialized-views' |
| 11 | import { PGTable } from '../pg-meta-tables' |
| 12 | import { PGView } from '../pg-meta-views' |
| 13 | import { Query } from './Query' |
| 14 | import { Filter, Sort } from './types' |
| 15 | |
| 16 | // Constants |
| 17 | export const MAX_CHARACTERS = 10 * 1024 // 10KB |
| 18 | // Max array size |
| 19 | export const MAX_ARRAY_SIZE = 50 |
| 20 | |
| 21 | export type TableLikeEntity = PGTable | PGView | PGForeignTable | PGMaterializedView |
| 22 | |
| 23 | export interface BuildTableRowsQueryArgs { |
| 24 | table: TableLikeEntity |
| 25 | filters?: Filter[] |
| 26 | sorts?: Sort[] |
| 27 | limit?: number |
| 28 | page?: number |
| 29 | maxCharacters?: number |
| 30 | maxArraySize?: number |
| 31 | /** |
| 32 | * Columns that should not be used for default sorting |
| 33 | */ |
| 34 | sortExcludedColumns?: string[] |
| 35 | } |
| 36 | |
| 37 | // Text and JSON types that should be truncated |
| 38 | export const TEXT_TYPES = [ |
| 39 | safeSql`text`, |
| 40 | safeSql`varchar`, |
| 41 | safeSql`char`, |
| 42 | safeSql`character varying`, |
| 43 | safeSql`character`, |
| 44 | ] |
| 45 | export const JSON_TYPES = [safeSql`json`, safeSql`jsonb`] |
| 46 | const JSON_SET = new Set(JSON_TYPES) |
| 47 | |
| 48 | // Additional PostgreSQL types that can hold large values and should be truncated |
| 49 | export const ADDITIONAL_LARGE_TYPES = [ |
| 50 | // Standard PostgreSQL types |
| 51 | safeSql`bytea`, // Binary data |
| 52 | safeSql`xml`, // XML data |
| 53 | safeSql`hstore`, // Key-value store |
| 54 | safeSql`clob`, // Character large object |
| 55 | |
| 56 | // Extension-specific types |
| 57 | // pgvector extension (for AI/ML/RAG applications) |
| 58 | safeSql`vector`, // Vector type used for embeddings |
| 59 | |
| 60 | // PostGIS extension types |
| 61 | safeSql`geometry`, // Spatial data type |
| 62 | safeSql`geography`, // Spatial data type |
| 63 | |
| 64 | // Full-text search types |
| 65 | safeSql`tsvector`, // Text search vector |
| 66 | safeSql`tsquery`, // Text search query |
| 67 | |
| 68 | // Range types |
| 69 | safeSql`daterange`, // Date range |
| 70 | safeSql`tsrange`, // Timestamp range |
| 71 | safeSql`tstzrange`, // Timestamp with timezone range |
| 72 | safeSql`numrange`, // Numeric range |
| 73 | safeSql`int4range`, // Integer range |
| 74 | safeSql`int8range`, // Bigint range |
| 75 | |
| 76 | // Other extension types |
| 77 | safeSql`cube`, // Multi-dimensional cube |
| 78 | safeSql`ltree`, // Label tree |
| 79 | safeSql`lquery`, // Label tree query |
| 80 | safeSql`jsonpath`, // JSON path expressions |
| 81 | safeSql`citext`, // Case-insensitive text |
| 82 | ] |
| 83 | |
| 84 | export const LARGE_COLUMNS_TYPES = [...TEXT_TYPES, ...JSON_TYPES, ...ADDITIONAL_LARGE_TYPES] |
| 85 | const LARGE_COLUMNS_TYPES_SET = new Set(LARGE_COLUMNS_TYPES) |
| 86 | |
| 87 | // Threshold count for applying default sort |
| 88 | export const THRESHOLD_COUNT = 100000 |
| 89 | |
| 90 | // Return the primary key columns if exists, otherwise return the first column to use as a default sort |
| 91 | export const getDefaultOrderByColumns = ( |
| 92 | table: Pick<PGTable, 'primary_keys' | 'columns'>, |
| 93 | { excludedColumns = [] }: { excludedColumns?: string[] } = {} |
| 94 | ) => { |
| 95 | const primaryKeyColumns = table.primary_keys?.map((pk) => pk.name) |
| 96 | if ( |
| 97 | primaryKeyColumns && |
| 98 | primaryKeyColumns.length > 0 && |
| 99 | !primaryKeyColumns.every((col) => excludedColumns.includes(col)) |
| 100 | ) { |
| 101 | return primaryKeyColumns |
| 102 | } |
| 103 | if (table.columns && table.columns.length > 0) { |
| 104 | const eligibleColumnsForSorting = table.columns.filter( |
| 105 | (x) => !x.data_type.includes('json') && !excludedColumns.includes(x.name) |
| 106 | ) |
| 107 | if (eligibleColumnsForSorting.length > 0) return [eligibleColumnsForSorting[0].name] |
| 108 | else return [] |
| 109 | } |
| 110 | return [] |
| 111 | } |
| 112 | |
| 113 | /** |
| 114 | * Determines if a column type should be truncated based on its format and dataType |
| 115 | * Be aware if the logic in RowEditor.utils.ts -> isValueTruncated needs to be revised |
| 116 | * if we're updating the truncation logic, as it'll affect whether the Table Editor displays |
| 117 | * the data as truncated or not |
| 118 | */ |
| 119 | export const shouldTruncateColumn = (columnFormat: string): boolean => |
| 120 | (LARGE_COLUMNS_TYPES_SET as Set<string>).has(columnFormat.toLowerCase()) |
| 121 | |
| 122 | export const DEFAULT_PAGE_SIZE = 100 |
| 123 | |
| 124 | export function getPagination(page?: number, size: number = DEFAULT_PAGE_SIZE) { |
| 125 | const limit = size |
| 126 | const from = page ? page * limit : 0 |
| 127 | const to = page ? from + size - 1 : size - 1 |
| 128 | |
| 129 | return { from, to } |
| 130 | } |
| 131 | |
| 132 | export const getTableRowsSql = ({ |
| 133 | table, |
| 134 | filters = [], |
| 135 | sorts = [], |
| 136 | page, |
| 137 | limit, |
| 138 | maxCharacters = MAX_CHARACTERS, |
| 139 | maxArraySize = MAX_ARRAY_SIZE, |
| 140 | sortExcludedColumns = [], |
| 141 | }: BuildTableRowsQueryArgs): SafeSqlFragment => { |
| 142 | if (!table || !table.columns) return safeSql`` |
| 143 | |
| 144 | const query = new Query() |
| 145 | |
| 146 | // Properly escape the table name and schema |
| 147 | let queryChains = query.from(table.name, table.schema).select() |
| 148 | |
| 149 | filters.forEach((x) => { |
| 150 | const col = table.columns?.find((y) => y.name === x.column) |
| 151 | const isStringTypeColumn = !!col ? (TEXT_TYPES as string[]).includes(col.format) : true |
| 152 | queryChains = queryChains.filter( |
| 153 | x.column, |
| 154 | x.operator, |
| 155 | !isStringTypeColumn && x.value === '' ? null : x.value |
| 156 | ) |
| 157 | }) |
| 158 | |
| 159 | // If sorts is empty and table row count is within threshold, use the primary key as the default sort. |
| 160 | // Only apply for selections over a Table, not View, MaterializedViews, ... |
| 161 | const liveRowCount = (table as PGTable).live_rows_estimate || 0 |
| 162 | if (sorts.length === 0 && liveRowCount <= THRESHOLD_COUNT && table.columns.length > 0) { |
| 163 | const defaultOrderByColumns = getDefaultOrderByColumns(table as PGTable, { |
| 164 | excludedColumns: sortExcludedColumns, |
| 165 | }) |
| 166 | if (defaultOrderByColumns.length > 0) { |
| 167 | defaultOrderByColumns.forEach((col) => { |
| 168 | queryChains = queryChains.order(table.name, col) |
| 169 | }) |
| 170 | } |
| 171 | } else { |
| 172 | sorts.forEach((x) => { |
| 173 | queryChains = queryChains.order(x.table, x.column, x.ascending, x.nullsFirst) |
| 174 | }) |
| 175 | } |
| 176 | |
| 177 | // getPagination is expecting to start from 0 |
| 178 | const { from, to } = getPagination((page ?? 1) - 1, limit) |
| 179 | |
| 180 | // To have efficient query, we use CTE optimization, to first reduce the number of rows and order them in the right place |
| 181 | // filtering, applying limits and order by, then we can apply selection with some conditional logic to truncate large columns |
| 182 | // allowing postgres to only truncate the columns within the subset that we'll return instead of attemting to do it on |
| 183 | // all the rows within the table |
| 184 | const baseSelectQuery = safeSql`with _base_query as (${queryChains.range(from, to).toSql({ isCTE: false, isFinal: false })})` |
| 185 | |
| 186 | const allColumnNames = table.columns |
| 187 | .sort((a, b) => a.ordinal_position - b.ordinal_position) |
| 188 | .map((column) => ({ name: column.name, format: column.format.toLowerCase() })) |
| 189 | |
| 190 | // Identify columns that might need truncation |
| 191 | const columnsToTruncate = table.columns |
| 192 | .filter((column) => shouldTruncateColumn(column.format)) |
| 193 | .map((column) => column.name) |
| 194 | |
| 195 | // Create select expressions for each column, applying truncation only to needed columns |
| 196 | const selectExpressions: Array<SafeSqlFragment> = allColumnNames.map(({ name: columnName }) => { |
| 197 | const escapedColumnName = ident(columnName) |
| 198 | |
| 199 | if (columnsToTruncate.includes(columnName)) { |
| 200 | return safeSql`case |
| 201 | when octet_length(${escapedColumnName}::text) > ${literal(maxCharacters)} |
| 202 | then left(${escapedColumnName}::text, ${literal(maxCharacters)}) || '...' |
| 203 | else ${escapedColumnName}::text |
| 204 | end as ${escapedColumnName}` |
| 205 | } else { |
| 206 | return escapedColumnName |
| 207 | } |
| 208 | }) |
| 209 | |
| 210 | // Handle array-based columns |
| 211 | const arrayBasedColumns = table.columns |
| 212 | .filter((column) => column.data_type.toLowerCase() === 'array') |
| 213 | // remove the _ prefix for array based format |
| 214 | .map((column) => ({ name: column.name, format: column.format.toLowerCase().slice(1) })) |
| 215 | |
| 216 | // Add array casting for array-based enum columns |
| 217 | arrayBasedColumns.forEach(({ name: columnName, format }) => { |
| 218 | // Find this column in our select expressions |
| 219 | const index = selectExpressions.findIndex( |
| 220 | (expr) => expr === ident(columnName) // if the column is selected without any truncation applied to it |
| 221 | ) |
| 222 | const isJson = (JSON_SET as Set<string>).has(format) |
| 223 | // format comes from pg_attribute (e.g. 'text', 'json') — ident() ensures safe quoting |
| 224 | const arrayTypeCast = isJson ? safeSql`::${keyword(format)}[]` : safeSql`::text[]` |
| 225 | const lastElement: SafeSqlFragment = isJson |
| 226 | ? safeSql`array['{"truncated": true}'::json]` |
| 227 | : safeSql`array['...']` |
| 228 | const col = ident(columnName) |
| 229 | if (index >= 0) { |
| 230 | // We cast to text[] but limit the array size if the total size of the array is too large (same logic than for text fields) |
| 231 | // This returns the first MAX_ARRAY_SIZE elements of the array (adjustable) and adds '...' if truncated |
| 232 | // NOTE: this is not optimal, as the first element in the array could still be very large (more than 10Kb) and in such case |
| 233 | // the trimming might fail. |
| 234 | // Also handle multi-dimentionals array truncation, but won't happen the extra `...` element to it as we can't determine what's |
| 235 | // the right number of items to generate within the array. Studio side, we'll consider any multi-dimentional array as possibly |
| 236 | // truncated. |
| 237 | selectExpressions[index] = safeSql` |
| 238 | case |
| 239 | when octet_length(${col}::text) > ${literal(maxCharacters)} |
| 240 | then |
| 241 | case |
| 242 | when array_ndims(${col}) = 1 |
| 243 | then |
| 244 | (select array_cat(${col}[1:${literal(maxArraySize)}]${arrayTypeCast}, ${lastElement}${arrayTypeCast}))${arrayTypeCast} |
| 245 | else |
| 246 | ${col}[1:${literal(maxArraySize)}]${arrayTypeCast} |
| 247 | end |
| 248 | else ${col}${arrayTypeCast} |
| 249 | end |
| 250 | ` |
| 251 | } |
| 252 | }) |
| 253 | |
| 254 | const selectClause = joinSqlFragments(selectExpressions, ',') |
| 255 | const finalQuery = new Query() |
| 256 | // Now, we apply our selection logic with the tables truncation on the _base_query constructed before |
| 257 | const finalQueryChain = finalQuery.from('_base_query').select(selectClause) |
| 258 | return safeSql`${baseSelectQuery} |
| 259 | ${finalQueryChain.toSql({ isCTE: true, isFinal: true })}` |
| 260 | } |
| 261 | |
| 262 | export default { |
| 263 | shouldTruncateColumn, |
| 264 | getTableRowsSql, |
| 265 | getDefaultOrderByColumns, |
| 266 | } |