Skip to content

Commit

Permalink
feat(logger): add clearState() method (#2408)
Browse files Browse the repository at this point in the history
Co-authored-by: Andrea Amorosi <[email protected]>
  • Loading branch information
shdq and dreamorosi authored Jun 14, 2024
1 parent 76303c5 commit f55e2d0
Show file tree
Hide file tree
Showing 5 changed files with 888 additions and 138 deletions.
136 changes: 102 additions & 34 deletions packages/logger/src/Logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,10 @@ class Logger extends Utility implements LoggerInterface {
* Standard attributes managed by Powertools that will be logged in all log items.
*/
private powertoolsLogData: PowertoolsLogData = <PowertoolsLogData>{};
/**
* Temporary log attributes that can be appended with `appendKeys()` method.
*/
private temporaryLogAttributes: LogAttributes = {};
/**
* Buffer used to store logs until the logger is initialized.
*
Expand All @@ -183,6 +187,13 @@ class Logger extends Utility implements LoggerInterface {
* Flag used to determine if the logger is initialized.
*/
#isInitialized = false;
/**
* Map used to hold the list of keys and their type.
*
* Because keys of different types can be overwritten, we keep a list of keys that were added and their last
* type. We then use this map at log preparation time to pick the last one.
*/
#keys: Map<string, 'temp' | 'persistent'> = new Map();

/**
* Log level used by the current instance of Logger.
Expand Down Expand Up @@ -234,23 +245,40 @@ class Logger extends Utility implements LoggerInterface {
}

/**
* It adds the given attributes (key-value pairs) to all log items generated by this Logger instance.
* It adds the given persistent attributes (key-value pairs) to all log items generated by this Logger instance.
*
* @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys()} instead.
*
* @param {LogAttributes} attributes
* @returns {void}
*/
public addPersistentLogAttributes(attributes?: LogAttributes): void {
merge(this.persistentLogAttributes, attributes);
public addPersistentLogAttributes(attributes: LogAttributes): void {
this.appendPersistentKeys(attributes);
}

/**
* Alias for addPersistentLogAttributes.
* It adds the given temporary attributes (key-value pairs) to all log items generated by this Logger instance.
*
* @param {LogAttributes} attributes
* @returns {void}
*/
public appendKeys(attributes?: LogAttributes): void {
this.addPersistentLogAttributes(attributes);
public appendKeys(attributes: LogAttributes): void {
for (const attributeKey of Object.keys(attributes)) {
this.#keys.set(attributeKey, 'temp');
}
merge(this.temporaryLogAttributes, attributes);
}

/**
* It adds the given persistent attributes (key-value pairs) to all log items generated by this Logger instance.
*
* @param attributes - The attributes to add to all log items.
*/
public appendPersistentKeys(attributes: LogAttributes): void {
for (const attributeKey of Object.keys(attributes)) {
this.#keys.set(attributeKey, 'persistent');
}
merge(this.persistentLogAttributes, attributes);
}

/**
Expand All @@ -274,6 +302,7 @@ class Logger extends Utility implements LoggerInterface {
customConfigService: this.getCustomConfigService(),
environment: this.powertoolsLogData.environment,
persistentLogAttributes: this.persistentLogAttributes,
temporaryLogAttributes: this.temporaryLogAttributes,
},
options
)
Expand Down Expand Up @@ -417,13 +446,6 @@ class Logger extends Utility implements LoggerInterface {
context,
callback
) {
let initialPersistentAttributes = {};
if (options && options.clearState === true) {
initialPersistentAttributes = {
...loggerRef.getPersistentLogAttributes(),
};
}

Logger.injectLambdaContextBefore(loggerRef, event, context, options);

let result: unknown;
Expand All @@ -432,25 +454,25 @@ class Logger extends Utility implements LoggerInterface {
} catch (error) {
throw error;
} finally {
Logger.injectLambdaContextAfterOrOnError(
loggerRef,
initialPersistentAttributes,
options
);
if (options?.clearState) loggerRef.resetState();
}

return result;
};
};
}

/**
* @deprecated This method is deprecated and will be removed in the future major versions. Use {@link resetState()} instead.
*/
/* istanbul ignore next */
public static injectLambdaContextAfterOrOnError(
logger: Logger,
initialPersistentAttributes: LogAttributes,
_persistentAttributes: LogAttributes,
options?: InjectLambdaContextOptions
): void {
if (options && options.clearState === true) {
logger.setPersistentLogAttributes(initialPersistentAttributes);
logger.resetState();
}
}

Expand Down Expand Up @@ -493,27 +515,53 @@ class Logger extends Utility implements LoggerInterface {
}

/**
* Alias for removePersistentLogAttributes.
* It removes temporary attributes based on provided keys to all log items generated by this Logger instance.
*
* @param {string[]} keys
* @returns {void}
*/
public removeKeys(keys: string[]): void {
this.removePersistentLogAttributes(keys);
for (const key of keys) {
this.temporaryLogAttributes[key] = undefined;

if (this.persistentLogAttributes[key]) {
this.#keys.set(key, 'persistent');
} else {
this.#keys.delete(key);
}
}
}

/**
* It removes attributes based on provided keys to all log items generated by this Logger instance.
* It removes persistent attributes based on provided keys to all log items generated by this Logger instance.
*
* @param {string[]} keys
* @returns {void}
*/
public removePersistentLogAttributes(keys: string[]): void {
for (const key of keys) {
if (this.persistentLogAttributes && key in this.persistentLogAttributes) {
delete this.persistentLogAttributes[key];
this.persistentLogAttributes[key] = undefined;

if (this.temporaryLogAttributes[key]) {
this.#keys.set(key, 'temp');
} else {
this.#keys.delete(key);
}
}
}

/**
* It resets the state, by removing all temporary log attributes added with `appendKeys()` method.
*/
public resetState(): void {
for (const key of Object.keys(this.temporaryLogAttributes)) {
if (this.persistentLogAttributes[key]) {
this.#keys.set(key, 'persistent');
} else {
this.#keys.delete(key);
}
}
this.temporaryLogAttributes = {};
}

/**
Expand All @@ -537,6 +585,8 @@ class Logger extends Utility implements LoggerInterface {
* It sets the given attributes (key-value pairs) to all log items generated by this Logger instance.
* Note: this replaces the pre-existing value.
*
* @deprecated This method is deprecated and will be removed in the future major versions, please use {@link appendPersistentKeys()} instead.
*
* @param {LogAttributes} attributes
* @returns {void}
*/
Expand Down Expand Up @@ -665,10 +715,18 @@ class Logger extends Utility implements LoggerInterface {
...this.getPowertoolsLogData(),
};

// gradually merge additional attributes starting from customer-provided persistent attributes
let additionalLogAttributes = { ...this.getPersistentLogAttributes() };
const additionalAttributes: LogAttributes = {};
// gradually add additional attributes picking only the last added for each key
for (const [key, type] of this.#keys) {
if (type === 'persistent') {
additionalAttributes[key] = this.persistentLogAttributes[key];
} else {
additionalAttributes[key] = this.temporaryLogAttributes[key];
}
}

// if the main input is not a string, then it's an object with additional attributes, so we merge it
additionalLogAttributes = merge(additionalLogAttributes, otherInput);
merge(additionalAttributes, otherInput);
// then we merge the extra input attributes (if any)
for (const item of extraInput) {
const attributes: LogAttributes =
Expand All @@ -678,12 +736,12 @@ class Logger extends Utility implements LoggerInterface {
? { extra: item }
: item;

additionalLogAttributes = merge(additionalLogAttributes, attributes);
merge(additionalAttributes, attributes);
}

return this.getLogFormatter().formatAttributes(
unformattedBaseAttributes,
additionalLogAttributes
additionalAttributes
);
}

Expand Down Expand Up @@ -1023,13 +1081,23 @@ class Logger extends Utility implements LoggerInterface {
serviceName,
sampleRateValue,
logFormatter,
persistentLogAttributes,
persistentKeys,
persistentLogAttributes, // deprecated in favor of persistentKeys
environment,
} = options;

if (persistentLogAttributes && persistentKeys) {
this.warn(
'Both persistentLogAttributes and persistentKeys options were provided. Using persistentKeys as persistentLogAttributes is deprecated and will be removed in future releases'
);
}

// configurations that affect log content
this.setPowertoolsLogData(serviceName, environment);
this.addPersistentLogAttributes(persistentLogAttributes);
this.setPowertoolsLogData(
serviceName,
environment,
persistentKeys || persistentLogAttributes
);

// configurations that affect Logger behavior
this.setLogEvent();
Expand Down Expand Up @@ -1070,7 +1138,7 @@ class Logger extends Utility implements LoggerInterface {
this.getEnvVarsService().getServiceName() ||
this.getDefaultServiceName(),
});
this.addPersistentLogAttributes(persistentLogAttributes);
this.appendPersistentKeys(persistentLogAttributes);
}
}

Expand Down
16 changes: 3 additions & 13 deletions packages/logger/src/middleware/middy.ts
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
import { Logger } from '../Logger.js';
import type { LogAttributes } from '../types/Log.js';
import type { InjectLambdaContextOptions } from '../types/Logger.js';
import { LOGGER_KEY } from '@aws-lambda-powertools/commons';
import type {
Expand Down Expand Up @@ -37,7 +36,6 @@ const injectLambdaContext = (
options?: InjectLambdaContextOptions
): MiddlewareLikeObj => {
const loggers = target instanceof Array ? target : [target];
const persistentAttributes: LogAttributes[] = [];
const isClearState = options && options.clearState === true;

/**
Expand All @@ -55,12 +53,8 @@ const injectLambdaContext = (
const injectLambdaContextBefore = async (
request: MiddyLikeRequest
): Promise<void> => {
loggers.forEach((logger: Logger, index: number) => {
loggers.forEach((logger: Logger) => {
if (isClearState) {
persistentAttributes[index] = {
...logger.getPersistentLogAttributes(),
};

setCleanupFunction(request);
}
Logger.injectLambdaContextBefore(
Expand All @@ -74,12 +68,8 @@ const injectLambdaContext = (

const injectLambdaContextAfterOrOnError = async (): Promise<void> => {
if (isClearState) {
loggers.forEach((logger: Logger, index: number) => {
Logger.injectLambdaContextAfterOrOnError(
logger,
persistentAttributes[index],
options
);
loggers.forEach((logger: Logger) => {
logger.resetState();
});
}
};
Expand Down
29 changes: 27 additions & 2 deletions packages/logger/src/types/Logger.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,41 @@ type InjectLambdaContextOptions = {
clearState?: boolean;
};

type ConstructorOptions = {
type BaseConstructorOptions = {
logLevel?: LogLevel;
serviceName?: string;
sampleRateValue?: number;
logFormatter?: LogFormatterInterface;
customConfigService?: ConfigServiceInterface;
persistentLogAttributes?: LogAttributes;
environment?: Environment;
};

type PersistentKeysOption = {
persistentKeys?: LogAttributes;
persistentLogAttributes?: never;
};

type DeprecatedOption = {
persistentLogAttributes?: LogAttributes;
persistentKeys?: never;
};

/**
* Options for the Logger class constructor.
*
* @type {Object} ConstructorOptions
* @property {LogLevel} [logLevel] - The log level.
* @property {string} [serviceName] - The service name.
* @property {number} [sampleRateValue] - The sample rate value.
* @property {LogFormatterInterface} [logFormatter] - The custom log formatter.
* @property {ConfigServiceInterface} [customConfigService] - The custom config service.
* @property {Environment} [environment] - The environment.
* @property {LogAttributes} [persistentKeys] - The keys that will be persisted in all log items.
* @property {LogAttributes} [persistentLogAttributes] - **Deprecated!** Use `persistentKeys`.
*/
type ConstructorOptions = BaseConstructorOptions &
(PersistentKeysOption | DeprecatedOption);

type LambdaFunctionContext = Pick<
Context,
| 'functionName'
Expand Down
Loading

0 comments on commit f55e2d0

Please sign in to comment.