feat: add auto pk selection (#1071)

Author: fpnuseis

README.md

@@ -250,7 +250,7 @@ const paginateConfig: PaginateConfig<CatEntity> {
    * Default: None
    * Description: TypeORM partial selection. Limit selection further by using `select` query param.
    * https://typeorm.io/select-query-builder#partial-selection
-   * Note: You must include the primary key in the selection.
+   * Note: if you do not contain the primary key in the select array, primary key will be added automatically.
    */
   select: ['id', 'name', 'color'],
 

src/helper.ts

@@ -153,6 +153,30 @@ export function includesAllPrimaryKeyColumns(qb: SelectQueryBuilder<unknown>, pr
         .every((column) => propertyPath.includes(column))
 }
 
+export function getPrimaryKeyColumns(qb: SelectQueryBuilder<any>, entityName?: string): string[] {
+    return qb.expressionMap.mainAlias?.metadata?.primaryColumns.map((column) =>
+        entityName ? `${entityName}.${column.propertyName}` : column.propertyName
+    )
+}
+
+export function getMissingPrimaryKeyColumns(qb: SelectQueryBuilder<any>, transformedCols: string[]): string[] {
+    if (!transformedCols || transformedCols.length === 0) return []
+
+    const mainEntityPrimaryKeys = getPrimaryKeyColumns(qb)
+    const missingPrimaryKeys: string[] = []
+
+    for (const pk of mainEntityPrimaryKeys) {
+        const columnProperties = getPropertiesByColumnName(pk)
+        const pkAlias = fixColumnAlias(columnProperties, qb.alias, false)
+
+        if (!transformedCols.includes(pkAlias)) {
+            missingPrimaryKeys.push(pkAlias)
+        }
+    }
+
+    return missingPrimaryKeys
+}
+
 export function hasColumnWithPropertyPath(
     qb: SelectQueryBuilder<unknown>,
     columnProperties: ColumnProperties

src/paginate.spec.ts

@@ -2515,7 +2515,7 @@ describe('paginate', () => {
         )
     })
 
-    it("should return all columns if select doesn't contain all primary columns", async () => {
+    it("should return primary columns if select doesn't contain all primary columns", async () => {
         const config: PaginateConfig<CatEntity> = {
             sortableColumns: ['id', 'name'],
             select: ['name'],
@@ -2526,7 +2526,10 @@ describe('paginate', () => {
 
         const result = await paginate<CatEntity>(query, catRepo, config)
 
-        expect(result.data).toStrictEqual(cats)
+        result.data.forEach((cat, index) => {
+            expect(cat.id).toBe(cats[index].id)
+            expect(cat.name).toBe(cats[index].name)
+        })
         expect(result.meta.select).toStrictEqual(undefined)
         expect(result.links.current).toBe('?page=1&limit=20&sortBy=id:ASC')
     })

src/paginate.ts

@@ -12,10 +12,10 @@ import {
     createRelationSchema,
     extractVirtualProperty,
     fixColumnAlias,
+    getMissingPrimaryKeyColumns,
     getPaddedExpr,
     getPropertiesByColumnName,
     getQueryUrlComponents,
-    includesAllPrimaryKeyColumns,
     isDateColumnType,
     isEntityKey,
     isFindOperator,
@@ -566,20 +566,26 @@ export async function paginate<T extends ObjectLiteral>(
 
     // When we partial select the columns (main or relation) we must add the primary key column otherwise
     // typeorm will not be able to map the result.
-    let selectParams =
+    // so, we check if the selected columns are a subset of the primary key columns
+    // and add the missing ones.
+    const selectParams =
         config.select && query.select && !config.ignoreSelectInQueryParam
             ? config.select.filter((column) => query.select.includes(column))
             : config.select
-    if (!includesAllPrimaryKeyColumns(queryBuilder, query.select)) {
-        selectParams = config.select
-    }
-    if (selectParams?.length > 0 && includesAllPrimaryKeyColumns(queryBuilder, selectParams)) {
-        const cols: string[] = selectParams.reduce((cols, currentCol) => {
+
+    if (selectParams?.length > 0) {
+        let cols: string[] = selectParams.reduce((cols, currentCol) => {
             const columnProperties = getPropertiesByColumnName(currentCol)
             const isRelation = checkIsRelation(queryBuilder, columnProperties.propertyPath)
             cols.push(fixColumnAlias(columnProperties, queryBuilder.alias, isRelation))
             return cols
         }, [])
+
+        const missingPrimaryKeys = getMissingPrimaryKeyColumns(queryBuilder, cols)
+        if (missingPrimaryKeys.length > 0) {
+            cols = cols.concat(missingPrimaryKeys)
+        }
+
         queryBuilder.select(cols)
     }