From 03a672300f5e8f7ef571982cff6e0bdc9b560cfc Mon Sep 17 00:00:00 2001
From: Ricardo Espinoza <ricardo@n8n.io>
Date: Tue, 10 Nov 2020 17:05:08 -0500
Subject: [PATCH] :zap: Add custom traits to Segment Node (#1145)

---
 .../nodes/Segment/IdentifyDescription.ts      | 32 +++++++++++++++
 .../nodes-base/nodes/Segment/Segment.node.ts  | 40 +++++++++++++++++++
 .../nodes/Segment/TrackDescription.ts         | 32 +++++++++++++++
 3 files changed, 104 insertions(+)

diff --git a/packages/nodes-base/nodes/Segment/IdentifyDescription.ts b/packages/nodes-base/nodes/Segment/IdentifyDescription.ts
index 44b2c60b48..72b78fef35 100644
--- a/packages/nodes-base/nodes/Segment/IdentifyDescription.ts
+++ b/packages/nodes-base/nodes/Segment/IdentifyDescription.ts
@@ -262,6 +262,38 @@ export const identifyFields = [
 							},
 						],
 					},
+					{
+						displayName: 'Custom Traits',
+						name: 'customTraitsUi',
+						placeholder: 'Add Custom Trait',
+						type: 'fixedCollection',
+						default: '',
+						typeOptions: {
+							multipleValues: true,
+						},
+						options: [
+							{
+								name: 'customTraitValues',
+								displayName: 'Custom Traits',
+								values: [
+									{
+										displayName: 'Key',
+										name: 'key',
+										type: 'string',
+										default: '',
+										description: '',
+									},
+									{
+										displayName: 'Value',
+										name: 'value',
+										type: 'string',
+										default: '',
+										description: '',
+									},
+								],
+							},
+						],
+					},
 				],
 			},
 		],
diff --git a/packages/nodes-base/nodes/Segment/Segment.node.ts b/packages/nodes-base/nodes/Segment/Segment.node.ts
index b4ad23e9c1..954193969e 100644
--- a/packages/nodes-base/nodes/Segment/Segment.node.ts
+++ b/packages/nodes-base/nodes/Segment/Segment.node.ts
@@ -38,6 +38,7 @@ import {
 } from './TrackInterface';
 
 import * as uuid from 'uuid/v4';
+import { customerFields } from '../CustomerIo/CustomerDescription';
 
 export class Segment implements INodeType {
 	description: INodeTypeDescription = {
@@ -170,6 +171,7 @@ export class Segment implements INodeType {
 						if (traits.id) {
 							body.traits!.id = traits.id as string;
 						}
+						
 						if (traits.company) {
 							const company = (traits.company as IDataObject).companyUi as IDataObject;
 							if (company) {
@@ -384,6 +386,14 @@ export class Segment implements INodeType {
 						if (traits.id) {
 							body.traits!.id = traits.id as string;
 						}
+						if (traits.customTraitsUi) {
+							const customTraits = (traits.customTraitsUi as IDataObject).customTraitValues as IDataObject[];
+							if (customTraits && customTraits.length !== 0) {
+								for (const customTrait of customTraits) {
+									body.traits![customTrait.key as string] = customTrait.value;
+								}
+							}
+						}
 						if (traits.company) {
 							const company = (traits.company as IDataObject).companyUi as IDataObject;
 							if (company) {
@@ -531,6 +541,17 @@ export class Segment implements INodeType {
 							body.integrations!.salesforce = integrations.salesforce as boolean;
 						}
 					}
+
+					if (Object.keys(traits.company as IDataObject).length === 0) {
+						//@ts-ignore
+						delete body.traits.company;
+					}
+
+					if (Object.keys(traits.address as IDataObject).length === 0) {
+						//@ts-ignore
+						delete body.traits.address;
+					}
+
 					responseData = await segmentApiRequest.call(this, 'POST', '/identify', body);
 				}
 			}
@@ -602,6 +623,14 @@ export class Segment implements INodeType {
 						if (traits.id) {
 							body.traits!.id = traits.id as string;
 						}
+						if (traits.customTraitsUi) {
+							const customTraits = (traits.customTraitsUi as IDataObject).customTraitValues as IDataObject[];
+							if (customTraits && customTraits.length !== 0) {
+								for (const customTrait of customTraits) {
+									body.traits![customTrait.key as string] = customTrait.value;
+								}
+							}
+						}
 						if (traits.company) {
 							const company = (traits.company as IDataObject).companyUi as IDataObject;
 							if (company) {
@@ -760,6 +789,17 @@ export class Segment implements INodeType {
 							body.properties!.value = properties.value as string;
 						}
 					}
+
+					if (Object.keys(traits.company as IDataObject).length === 0) {
+						//@ts-ignore
+						delete body.traits.company;
+					}
+
+					if (Object.keys(traits.address as IDataObject).length === 0) {
+						//@ts-ignore
+						delete body.traits.address;
+					}
+
 					responseData = await segmentApiRequest.call(this, 'POST', '/track', body);
 				}
 				//https://segment.com/docs/connections/sources/catalog/libraries/server/http-api/#page
diff --git a/packages/nodes-base/nodes/Segment/TrackDescription.ts b/packages/nodes-base/nodes/Segment/TrackDescription.ts
index f32b03278c..1e49c5e579 100644
--- a/packages/nodes-base/nodes/Segment/TrackDescription.ts
+++ b/packages/nodes-base/nodes/Segment/TrackDescription.ts
@@ -285,6 +285,38 @@ export const trackFields = [
 							},
 						],
 					},
+					{
+						displayName: 'Custom Traits',
+						name: 'customTraitsUi',
+						placeholder: 'Add Custom Trait',
+						type: 'fixedCollection',
+						default: '',
+						typeOptions: {
+							multipleValues: true,
+						},
+						options: [
+							{
+								name: 'customTraitValues',
+								displayName: 'Custom Traits',
+								values: [
+									{
+										displayName: 'Key',
+										name: 'key',
+										type: 'string',
+										default: '',
+										description: '',
+									},
+									{
+										displayName: 'Value',
+										name: 'value',
+										type: 'string',
+										default: '',
+										description: '',
+									},
+								],
+							},
+						],
+					},
 				],
 			},
 		],