Merge pull request #78 from MatteoGioioso/connection-filtering

Connection filtering

README.md

@@ -5,32 +5,47 @@
 ![GitHub](https://img.shields.io/github/license/MatteoGioioso/serverless-pg) \
 [!["Buy Me A Coffee"](https://www.buymeacoffee.com/assets/img/custom_images/orange_img.png)](https://www.buymeacoffee.com/matteogioioso)
 
-
 Serverless-postgres is a wrapper for **[node-pg](https://github.com/brianc/node-postgres)** Node.js module.
 It is heavily inspired by Jeremy Daly's **[serverless-mysql](https://github.com/jeremydaly/serverless-mysql)** package.
 
 ### Why I need this module?
-In a serverless application a function can scale almost "infinitely" by creating separate container instances 
-for each concurrent user. 
+
+In a serverless application a function can scale almost "infinitely" by creating separate container instances
+for each concurrent user.
 Each container can correspond to a database connection which, for performance purposes, is left opened for further
 re-utilization. If we have a sudden spike of concurrent traffic, the available connections can be quickly maxed out
 by other competing functions.
 If we reach the max connections limit, Postgres will automatically reject any frontend trying to connect to its backend.
 This can cause heavy disruption in your application.
 
 ### What does it do?
+
 Serverless-postgres adds a connection management component specifically for FaaS based applications.
 By calling the method `.clean()` at the end of your functions, the module will constantly monitor the status of all
-the processes running in the PostgreSQL backend and then, based on the configuration provided, 
+the processes running in the PostgreSQL backend and then, based on the configuration provided,
 will garbage collect the "zombie" connections.
 If the client fails to connect with `"sorry, too many clients already"` error, the module will retry
 using trusted backoff algorithms.
 
-> **NOTE:** This module *should* work with any PostgreSQL server. 
-It has been tested with AWS's RDS Postgres, Aurora Postgres, and Aurora Serverless.
+> **NOTE:** This module *should* work with any PostgreSQL server.
+> It has been tested with AWS's RDS Postgres, Aurora Postgres, and Aurora Serverless.
 
 Feel free to request additional features and contribute =)
 
+## Changelog
+
+- **Default connections filtering (>= v2)**: this feature leverage postgres `application_name` to differentiate
+  clients created by this library and others, this will avoid terminating connections belonging to long-running
+  process, batch jobs, ect...
+  By default, we set the same `application_name` parameter for all the serverless clients, if you wish you can change it
+  by just specifying it in the client config:
+  ```javascript
+  const client = new ServerlessClient({
+    ...
+    application_name: 'my_client',
+  });
+  ```
+
 ## Install
 
 ```bash
@@ -45,23 +60,23 @@ Declare the ServerlessClient outside the lambda handler
 const ServerlessClient = require('serverless-postgres')
 
 const client = new ServerlessClient({
-    user: process.env.DB_USER,
-    host: process.env.DB_HOST,
-    database: process.env.DB_NAME,
-    password: process.env.DB_PASSWORD,
-    port: process.env.DB_PORT,
-    debug: true,
-    delayMs: 3000,
+  user: process.env.DB_USER,
+  host: process.env.DB_HOST,
+  database: process.env.DB_NAME,
+  password: process.env.DB_PASSWORD,
+  port: process.env.DB_PORT,
+  debug: true,
+  delayMs: 3000,
 });
 
-const handler = async(event, context) => {
-    await client.connect();
-    const result = await client.query(`SELECT 1+1 AS result`);
-    await client.clean();
-    return {
-      body: JSON.stringify({message: result.rows[0]}),
-      statusCode: 200
-    }
+const handler = async (event, context) => {
+  await client.connect();
+  const result = await client.query(`SELECT 1+1 AS result`);
+  await client.clean();
+  return {
+    body: JSON.stringify({ message: result.rows[0] }),
+    statusCode: 200
+  }
 }
 
 ```
@@ -72,24 +87,22 @@ You can set the configuration dynamically if your secret is stored in a vault
 const ServerlessClient = require('serverless-postgres')
 
 const client = new ServerlessClient({
-    host: process.env.DB_HOST,
-    database: process.env.DB_NAME,
-    port: process.env.DB_PORT,
+  host: process.env.DB_HOST,
+  database: process.env.DB_NAME,
+  port: process.env.DB_PORT,
 });
 
-const handler = async(event, context) => {
-    const {user, password} = await getCredentials('my-secret')
-    client.setConfig({
-      user, password
-    })
-    await client.connect();
-    // ...rest of the code
+const handler = async (event, context) => {
+  const { user, password } = await getCredentials('my-secret')
+  client.setConfig({
+    user, password
+  })
+  await client.connect();
+  // ...rest of the code
 }
 
 ```
 
-
-
 ## Configuration Options
 
 | Property | Type | Description | Default |
@@ -112,9 +125,8 @@ const handler = async(event, context) => {
 | allowCredentialsDiffing | `Boolean` | If you are using dynamic credentials, such as IAM, you can set this parameter to `true` and the client will be refreshed | `false` |
 | library | `Function` | Custom postgres library | `require('pg')` |
 
-
 ## Note
 
-- `Serverless-postgres` depends on `pg` package and usually you **do not need to install it on your own**. 
+- `Serverless-postgres` depends on `pg` package and usually you **do not need to install it on your own**.
   As some users have observed, if you have installed it on your own, and it is a different version,
   this package might misbehave.

__tests__/index.test.js

@@ -1,4 +1,5 @@
 const ServerlessClient = require("../index");
+const AWSXRay = require("aws-xray-sdk");
 
 jest.setTimeout(30000);
 
@@ -368,6 +369,7 @@ describe("Serverless client", function () {
         database: "postgres2",
         password: "postgres2",
         port: 22001,
+        application_name: 'serverless_client'
       })
 
       // Switch again database to the previous one

package.json

@@ -14,14 +14,15 @@
     "pg": "^8.5.1"
   },
   "devDependencies": {
+    "@babel/helper-compilation-targets": "^7.17.7",
+    "@babel/helper-environment-visitor": "^7.16.7",
+    "@babel/helper-hoist-variables": "^7.16.7",
     "aws-xray-sdk": "^3.3.3",
     "husky": "^4.3.8",
     "jest": "^26.6.3",
     "prettier": "^1.19.1",
     "semantic-release": "^19.0.2",
-    "@babel/helper-compilation-targets": "^7.17.7",
-    "@babel/helper-environment-visitor": "^7.16.7",
-    "@babel/helper-hoist-variables": "^7.16.7"
+    "pg-query-stream": "^4.2.3"
   },
   "prettier": {
     "printWidth": 100

src/index.js

@@ -3,7 +3,6 @@
  * This module wrap node-postgres package, more detail regarding it can be found here:
  * https://github.com/brianc/node-postgres
  * @author Matteo Gioioso <[email protected]>
- * @version 1.3.0
  * @license MIT
  */
 
@@ -46,13 +45,15 @@ ServerlessClient.prototype._getIdleProcessesListOrderByDate = async function ()
       WHERE datname = $1
         AND state = 'idle'
         AND usename = $2
+        AND application_name = $4
       ORDER BY state_change
       LIMIT $3;`
 
   const values = [
     this._client.database,
     this._client.user,
-    this._strategy.maxIdleConnectionsToKill
+    this._strategy.maxIdleConnectionsToKill,
+    this._application_name
   ]
 
   try {
@@ -78,6 +79,7 @@ ServerlessClient.prototype._getIdleProcessesListByMinimumTimeout = async functio
           WHERE usename = $1
             AND datname = $2
             AND state = 'idle'
+            AND application_name = $5
       )
       SELECT pid
       FROM processes
@@ -88,7 +90,8 @@ ServerlessClient.prototype._getIdleProcessesListByMinimumTimeout = async functio
     this._client.user,
     this._client.database,
     this._strategy.minConnIdleTimeSec,
-    this._strategy.maxIdleConnectionsToKill
+    this._strategy.maxIdleConnectionsToKill,
+    this._application_name
   ]
 
   try {
@@ -117,9 +120,10 @@ ServerlessClient.prototype._getProcessesCount = async function () {
         SELECT COUNT(pid)
         FROM pg_stat_activity
         WHERE datname = $1
-          AND usename = $2;`
+          AND usename = $2
+          AND application_name = $3;`
 
-    const values = [this._client.database, this._client.user]
+    const values = [this._client.database, this._client.user, this._application_name]
 
     try {
       const result = await this._client.query(query, values);
@@ -147,14 +151,15 @@ ServerlessClient.prototype._killProcesses = async function (processesList) {
       SELECT pg_terminate_backend(pid)
       FROM pg_stat_activity
       WHERE pid = ANY ($1)
-        AND state = 'idle';`
+        AND state = 'idle'
+        AND application_name = $2;`
 
-  const values = [pids]
+  const values = [pids, this._application_name]
 
   try {
     return await this._client.query(query, values)
   } catch (e) {
-    this._logger("Swallowed internal error", e.message)
+    this._logger("Swallowed internal error: ", e.message)
     // Swallow the error, if this produce an error there is no need to error the function
 
     return {
@@ -377,6 +382,9 @@ ServerlessClient.prototype.setConfig = function (config) {
     queryRetries: 0
   }
 
+  this._application_name = this._config.application_name || "serverless_client"
+  this._config.application_name = this._application_name
+
   // Prevent diffing also if client is null
   if (this._multipleCredentials.allowCredentialsDiffing && this._client !== null) {
     this._diffCredentials(prevConfig, config)