From 301bbc29ccf91b71fe00ce73aee598b7807398e0 Mon Sep 17 00:00:00 2001
From: Frank Paulo Filho <me@frankpf.com>
Date: Sat, 10 Feb 2018 20:26:36 -0200
Subject: [PATCH] Output verbose schema information in "Fields" interfaces

---
 src/schemaInterfaces.ts |  4 +++-
 src/schemaPostgres.ts   | 20 ++++++++++++++++++--
 src/typescript.ts       | 19 ++++++++++++++++---
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/src/schemaInterfaces.ts b/src/schemaInterfaces.ts
index 1c33d6f..60c45db 100644
--- a/src/schemaInterfaces.ts
+++ b/src/schemaInterfaces.ts
@@ -3,7 +3,9 @@ import Options from './options'
 export interface ColumnDefinition {
     udtName: string,
     nullable: boolean,
-    tsType?: string
+    tsType?: string,
+    primaryKey?: boolean,
+    unique?: boolean
 }
 
 export interface TableDefinition {
diff --git a/src/schemaPostgres.ts b/src/schemaPostgres.ts
index 2acceb9..9272db1 100644
--- a/src/schemaPostgres.ts
+++ b/src/schemaPostgres.ts
@@ -119,7 +119,7 @@ export class PostgresDatabase implements Database {
 
     public async getTableDefinition (tableName: string, tableSchema: string) {
         let tableDefinition: TableDefinition = {}
-        type T = { column_name: string, udt_name: string, is_nullable: string }
+        type T = { column_name: string, udt_name: string, is_nullable: string, constraint_type: 'UNIQUE' | 'PRIMARY KEY' }
         await this.db.each<T>(
             'SELECT column_name, udt_name, is_nullable ' +
             'FROM information_schema.columns ' +
@@ -128,9 +128,25 @@ export class PostgresDatabase implements Database {
             (schemaItem: T) => {
                 tableDefinition[schemaItem.column_name] = {
                     udtName: schemaItem.udt_name,
-                    nullable: schemaItem.is_nullable === 'YES'
+                    nullable: schemaItem.is_nullable === 'YES',
+                    primaryKey: false,
+                    unique: false
                 }
             })
+
+        // Fill in PRIMARY KEY and UNIQUE constraint details
+        await this.db.each<T>(
+            'SELECT kcu.column_name, tc.constraint_type ' +
+            'FROM information_schema.table_constraints AS tc ' +
+            'JOIN information_schema.key_column_usage AS kcu ' +
+            'ON tc.constraint_name = kcu.constraint_name ' +
+            'WHERE tc.table_name = $1 and tc.table_schema = $2',
+            [tableName, tableSchema],
+            (schemaItem: T) => {
+                tableDefinition[schemaItem.column_name].unique = schemaItem.constraint_type === 'UNIQUE'
+                tableDefinition[schemaItem.column_name].primaryKey = schemaItem.constraint_type === 'PRIMARY KEY'
+            }
+            )
         return tableDefinition
     }
 
diff --git a/src/typescript.ts b/src/typescript.ts
index 0add9e0..6a8a60e 100644
--- a/src/typescript.ts
+++ b/src/typescript.ts
@@ -54,10 +54,23 @@ export function generateTableTypes (tableNameRaw: string, tableDefinition: Table
     const tableName = options.transformTypeName(tableNameRaw)
     let fields = ''
     Object.keys(tableDefinition).forEach((columnNameRaw) => {
-        let type = tableDefinition[columnNameRaw].tsType
-        let nullable = tableDefinition[columnNameRaw].nullable ? '| null' : ''
         const columnName = options.transformColumnName(columnNameRaw)
-        fields += `export type ${normalizeName(columnName, options)} = ${type}${nullable};\n`
+        fields += `export type ${normalizeName(columnName, options)} = {`
+
+        const { tsType, nullable, primaryKey, unique } = tableDefinition[columnNameRaw]
+
+        // Mapped TS type
+        fields += `type: ${tsType}`
+        fields += nullable ? '| null' : ''
+        fields += `,`
+
+        // Primary key constraint
+        fields += primaryKey !== undefined ? `primaryKey: ${primaryKey},` : ''
+
+        // Unique constraint
+        fields += unique !== undefined ? `unique: ${unique},` : ''
+
+        fields += '};\n'
     })
 
     return `