table-row-query.ts266 lines · main
1import {
2 ident,
3 joinSqlFragments,
4 keyword,
5 literal,
6 safeSql,
7 type SafeSqlFragment,
8} from '../pg-format'
9import { PGForeignTable } from '../pg-meta-foreign-tables'
10import { PGMaterializedView } from '../pg-meta-materialized-views'
11import { PGTable } from '../pg-meta-tables'
12import { PGView } from '../pg-meta-views'
13import { Query } from './Query'
14import { Filter, Sort } from './types'
15
16// Constants
17export const MAX_CHARACTERS = 10 * 1024 // 10KB
18// Max array size
19export const MAX_ARRAY_SIZE = 50
20
21export type TableLikeEntity = PGTable | PGView | PGForeignTable | PGMaterializedView
22
23export 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
38export const TEXT_TYPES = [
39 safeSql`text`,
40 safeSql`varchar`,
41 safeSql`char`,
42 safeSql`character varying`,
43 safeSql`character`,
44]
45export const JSON_TYPES = [safeSql`json`, safeSql`jsonb`]
46const JSON_SET = new Set(JSON_TYPES)
47
48// Additional PostgreSQL types that can hold large values and should be truncated
49export 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
84export const LARGE_COLUMNS_TYPES = [...TEXT_TYPES, ...JSON_TYPES, ...ADDITIONAL_LARGE_TYPES]
85const LARGE_COLUMNS_TYPES_SET = new Set(LARGE_COLUMNS_TYPES)
86
87// Threshold count for applying default sort
88export const THRESHOLD_COUNT = 100000
89
90// Return the primary key columns if exists, otherwise return the first column to use as a default sort
91export 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 */
119export const shouldTruncateColumn = (columnFormat: string): boolean =>
120 (LARGE_COLUMNS_TYPES_SET as Set<string>).has(columnFormat.toLowerCase())
121
122export const DEFAULT_PAGE_SIZE = 100
123
124export 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
132export 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
262export default {
263 shouldTruncateColumn,
264 getTableRowsSql,
265 getDefaultOrderByColumns,
266}