feat: Output nullability overriding with aliases (adelsz#377)

* feat: Alias nullability support You can now use alias hints for nullability overriding in the output columns. For example: `"name!"` for a non-nullable type, or `"name?"` for a nullable type * Revert "feat: Alias nullability support" This reverts commit b1a6797. * Simpler alias nullability hints implementation * Fix Node.js 14 compatibility * Apply suggestions from code review Co-authored-by: Adel Salakh <[email protected]> * Fix tag.ts cleanup * Fix cleanup & linter issues Co-authored-by: Adel Salakh <[email protected]>
logan12358 · May 6, 2022 · dfb0b66 · dfb0b66
1 parent d790e9f
commit dfb0b66
Show file tree

Hide file tree

Showing 10 changed files with 148 additions and 18 deletions.
diff --git a/docs-new/docs/sql-file.md b/docs-new/docs/sql-file.md
@@ -28,6 +28,7 @@ PgTyped has a number of requirements for SQL file contents:
 4. Queries can contain parameters. Parameters should start with a colon, ex. `:paramName`.
 5. Annotations can include param expansions if needed using the `@param` tag.
 6. Parameters can be forced to be not nullable using an exclamation mark `:paramName!`.
+7. Nullability on output columns for output columns can be specified using column aliases such as `AS "name?"` or `AS "name!"`.
 
 ## Parameter expansions
 
@@ -155,7 +156,7 @@ insertUsers.run(parameters, connection);
 INSERT INTO users (name, age) VALUES ($1, $2), ($3, $4) RETURNING id;
 ```
 
-### Enforcing non-nullability
+### Enforcing non-nullability for parameters
 
 Sometimes you might want to force pgTyped to use a non-nullable type for a nullable parameter.
 This can be done using the exclamation mark modifier `:paramName!`.
@@ -180,6 +181,32 @@ export interface IGetAllCommentsStrictParams {
 }
 ```
 
+## Enforcing (non)-nullability on output columns
+Sometimes you might want to force pgTyped to use a (non-)nullable type for an output column as Postgres limits how much
+information can be automatically discovered. This can be the case with materialized views and function calls for
+example.   
+You can enforce nullability on output columns by aliasing the column by using `AS "name?"` to make it nullable or
+`AS "name!"` to enforce non-nullability.
+
+#### Example:
+
+```sql title="Query definition:"
+/* @name GetTotalUserScores */
+SELECT coalesce(sum(score), 0) AS "total_score!" FROM users WHERE id = :id!;
+/* @name GetUsersAndTheirNames */
+SELECT id, name AS "name?" FROM users LEFT JOIN names USING (id);
+```
+
+```ts title="Resulting code:"
+export interface IGetTotalUserScoresResult {
+  total_score: number;
+}
+export interface IGetUsersAndTheirNamesResult {
+  id: number;
+  name: string | null;
+}
+```
+
 :::note
 We will be adding more annotation tags and expansion types in the future.  
 If you have an idea for a new expansion type, or a new annotation tag, please submit an issue for that so we can consider it.

diff --git a/packages/cli/package-lock.json b/packages/cli/package-lock.json
diff --git a/packages/cli/src/generator.test.ts b/packages/cli/src/generator.test.ts
@@ -1,14 +1,14 @@
 import * as queryModule from '@pgtyped/query';
 import { parseSQLFile, parseTypeScriptFile } from '@pgtyped/query';
 import { IQueryTypes } from '@pgtyped/query/lib/actions';
+import { ParsedConfig } from './config';
 import {
   escapeComment,
   generateInterface,
   queryToTypeDeclarations,
 } from './generator';
 import { ProcessingMode } from './index';
 import { DefaultTypeMapping, TypeAllocator } from './types';
-import { ParsedConfig } from './config';
 
 const getTypesMocked = jest.spyOn(queryModule, 'getTypes').mockName('getTypes');
 
@@ -469,6 +469,78 @@ export interface IGetNotificationsResult {
   type: PayloadType;
 }
 
+/** 'GetNotifications' query type */
+export interface IGetNotificationsQuery {
+  params: IGetNotificationsParams;
+  result: IGetNotificationsResult;
+}\n\n`;
+      expect(result).toEqual(expected);
+    });
+
+    test(`Columns with nullability hints (${mode})`, async () => {
+      const queryStringSQL = `
+    /* @name GetNotifications */
+    SELECT payload as "payload!", type as "type?" FROM notifications WHERE id = :userId;
+    `;
+      const queryStringTS = `
+      const getNotifications = sql\`SELECT payload as "payload!", type FROM notifications WHERE id = $userId\`;
+      `;
+      const queryString =
+        mode === ProcessingMode.SQL ? queryStringSQL : queryStringTS;
+      const mockTypes: IQueryTypes = {
+        returnTypes: [
+          {
+            returnName: 'payload!',
+            columnName: 'payload!',
+            type: 'json',
+          },
+          {
+            returnName: 'type?',
+            columnName: 'type?',
+            type: { name: 'PayloadType', enumValues: ['message', 'dynamite'] },
+            nullable: false,
+          },
+        ],
+        paramMetadata: {
+          params: ['uuid'],
+          mapping: [
+            {
+              name: 'userId',
+              type: queryModule.ParamTransform.Scalar,
+              required: false,
+              assignedIndex: 1,
+            },
+          ],
+        },
+      };
+      getTypesMocked.mockResolvedValue(mockTypes);
+      const types = new TypeAllocator(DefaultTypeMapping);
+      // Test out imports
+      types.use({ name: 'PreparedQuery', from: '@pgtyped/query' });
+      const result = await queryToTypeDeclarations(
+        parsedQuery(mode, queryString),
+        null,
+        types,
+        {} as ParsedConfig,
+      );
+      const expectedTypes = `import { PreparedQuery } from '@pgtyped/query';
+
+export type PayloadType = 'dynamite' | 'message';
+
+export type Json = null | boolean | number | string | Json[] | { [key: string]: Json };\n`;
+
+      expect(types.declaration()).toEqual(expectedTypes);
+      const expected = `/** 'GetNotifications' parameters type */
+export interface IGetNotificationsParams {
+  userId: string | null | void;
+}
+
+/** 'GetNotifications' return type */
+export interface IGetNotificationsResult {
+  payload: Json;
+  type: PayloadType | null;
+}
+
 /** 'GetNotifications' query type */
 export interface IGetNotificationsQuery {
   params: IGetNotificationsParams;

diff --git a/packages/cli/src/generator.ts b/packages/cli/src/generator.ts
@@ -4,19 +4,19 @@ import {
   parseSQLFile,
   parseTypeScriptFile,
   prettyPrintEvents,
-  processTSQueryAST,
   processSQLQueryIR,
+  processTSQueryAST,
   queryASTToIR,
   SQLQueryAST,
   SQLQueryIR,
   TSQueryAST,
 } from '@pgtyped/query';
 import { camelCase } from 'camel-case';
 import { pascalCase } from 'pascal-case';
+import path from 'path';
+import { ParsedConfig } from './config';
 import { ProcessingMode } from './index';
 import { DefaultTypeMapping, TypeAllocator } from './types';
-import { ParsedConfig } from './config';
-import path from 'path';
 
 export interface IField {
   fieldName: string;
@@ -102,10 +102,21 @@ export async function queryToTypeDeclarations(
 
   returnTypes.forEach(({ returnName, type, nullable, comment }) => {
     let tsTypeName = types.use(type);
-    if (nullable || nullable == null) {
+
+    const lastCharacter = returnName[returnName.length - 1]; // Checking for type hints
+    const addNullability = lastCharacter === '?';
+    const removeNullability = lastCharacter === '!';
+    if (
+      (addNullability || nullable || nullable == null) &&
+      !removeNullability
+    ) {
       tsTypeName += ' | null';
     }
 
+    if (addNullability || removeNullability) {
+      returnName = returnName.slice(0, -1);
+    }
+
     returnFieldTypes.push({
       fieldName: config.camelCaseColumnNames
         ? camelCase(returnName)

diff --git a/packages/example/package-lock.json b/packages/example/package-lock.json
diff --git a/packages/example/src/books/books.queries.ts b/packages/example/src/books/books.queries.ts
@@ -211,7 +211,7 @@ export interface IAggregateEmailsAndTestParams {
 /** 'AggregateEmailsAndTest' return type */
 export interface IAggregateEmailsAndTestResult {
   agetest: boolean | null;
-  emails: stringArray | null;
+  emails: stringArray;
 }
 
 /** 'AggregateEmailsAndTest' query type */
@@ -220,12 +220,12 @@ export interface IAggregateEmailsAndTestQuery {
   result: IAggregateEmailsAndTestResult;
 }
 
-const aggregateEmailsAndTestIR: any = {"usedParamSet":{"testAges":true},"params":[{"name":"testAges","required":false,"transform":{"type":"scalar"},"locs":[{"a":52,"b":60}]}],"statement":"SELECT array_agg(email) as emails, array_agg(age) = :testAges as ageTest FROM users"};
+const aggregateEmailsAndTestIR: any = {"usedParamSet":{"testAges":true},"params":[{"name":"testAges","required":false,"transform":{"type":"scalar"},"locs":[{"a":55,"b":63}]}],"statement":"SELECT array_agg(email) as \"emails!\", array_agg(age) = :testAges as ageTest FROM users"};
 
 /**
  * Query generated from SQL:
  * ```
- * SELECT array_agg(email) as emails, array_agg(age) = :testAges as ageTest FROM users
+ * SELECT array_agg(email) as "emails!", array_agg(age) = :testAges as ageTest FROM users
  * ```
  */
 export const aggregateEmailsAndTest = new PreparedQuery<IAggregateEmailsAndTestParams,IAggregateEmailsAndTestResult>(aggregateEmailsAndTestIR);

diff --git a/packages/example/src/books/books.sql b/packages/example/src/books/books.sql
@@ -47,4 +47,4 @@ INNER JOIN authors a ON a.id = b.author_id
 WHERE a.first_name || ' ' || a.last_name = :authorName!;
 
 /* @name AggregateEmailsAndTest */
-SELECT array_agg(email) as emails, array_agg(age) = :testAges as ageTest FROM users;
+SELECT array_agg(email) as "emails!", array_agg(age) = :testAges as ageTest FROM users;
diff --git a/packages/query/src/tag.ts b/packages/query/src/tag.ts
@@ -1,12 +1,31 @@
-import { processTSQueryAST } from './preprocessor-ts';
-import { processSQLQueryIR } from './preprocessor-sql';
 import { QueryIR } from './loader/sql';
 import { parseTSQuery, TSQueryAST } from './loader/typescript';
+import { processSQLQueryIR } from './preprocessor-sql';
+import { processTSQueryAST } from './preprocessor-ts';
 
 export interface IDatabaseConnection {
   query: (query: string, bindings: any[]) => Promise<{ rows: any[] }>;
 }
 
+/** Check for column modifier suffixes (exclamation and question marks). */
+function isHintedColumn(columnName: string): boolean {
+  const lastCharacter = columnName[columnName.length - 1];
+  return lastCharacter === '!' || lastCharacter === '?';
+}
+
+function mapQueryResultRows(rows: any[]): any[] {
+  for (const row of rows) {
+    for (const columnName in row) {
+      if (isHintedColumn(columnName)) {
+        const newColumnNameWithoutSuffix = columnName.slice(0, -1);
+        row[newColumnNameWithoutSuffix] = row[columnName];
+        delete row[columnName];
+      }
+    }
+  }
+  return rows;
+}
+
 /* Used for SQL-in-TS */
 export class TaggedQuery<TTypePair extends { params: any; result: any }> {
   public run: (
@@ -24,7 +43,7 @@ export class TaggedQuery<TTypePair extends { params: any; result: any }> {
         params as any,
       );
       const result = await connection.query(processedQuery, bindings);
-      return result.rows;
+      return mapQueryResultRows(result.rows);
     };
   }
 }
@@ -58,7 +77,7 @@ export class PreparedQuery<TParamType, TResultType> {
         params as any,
       );
       const result = await connection.query(processedQuery, bindings);
-      return result.rows;
+      return mapQueryResultRows(result.rows);
     };
   }
 }

diff --git a/packages/wire/package-lock.json b/packages/wire/package-lock.json
diff --git a/tslint.json b/tslint.json
@@ -14,7 +14,8 @@
     "array-type": false,
     "max-line-length": false,
     "max-classes-per-file": false,
-    "no-unused-variable": true
+    "no-unused-variable": true,
+    "forin": false
   },
   "rulesDirectory": [],
   "linterOptions": {