Feat/enable pg pool connection

README.md

@@ -14,10 +14,10 @@ Additional information regarding the adapter and getting started is provided bel
 
 PostgreSQL Prometheus Adapter supports:
 
-* PostgreSQL 14
-* PostgreSQL 13
-* PostgreSQL 12
-* PostgreSQL 11
+- PostgreSQL 14
+- PostgreSQL 13
+- PostgreSQL 12
+- PostgreSQL 11
 
 ## Building
 
@@ -54,7 +54,7 @@ export DATABASE_URL=...
 ```shell
 Default: None
 Description: Database connection parameters
-Ex: “user=<> password=<> host=<> port=<> database=<>”
+Ex: “user=<> password=<> host=<> port=<> database=<> pool_min_conns=[] pool_max_conns=[]”
 ```
 
 #### Adapter parameters
@@ -80,6 +80,7 @@ Flags:
       --pg-threads=1                   Writer DB threads to run 1-10
       --parser-threads=5               parser threads to run per DB writer 1-10
 ```
+
 :point_right: Note: pg_commit_secs and pg_commit_rows controls when data rows will be flushed to database. First one to reach threshold will trigger the flush.
 
 ### Container
@@ -93,7 +94,7 @@ podman run --rm \
   -e DATABASE_URL="user=testuser password=test123 host=192.168.12.36 port=5432 database=testdb" \
   --detach \
   crunchydata/postgresql-prometheus-adapterl:latest
-  ```
+```
 
 #### Stop container
 
@@ -117,6 +118,7 @@ pg_commit_rows=20000           Write data to database every N Rows
 pg_threads=1                   Writer DB threads to run 1-10
 parser_threads=5               parser threads to run per DB writer 1-10
 ```
+
 :point_right: Note: pg_commit_secs and pg_commit_rows controls when data rows will be flushed to database. First one to reach threshold will trigger the flush.
 
 ## Prometheus Configuration
@@ -125,18 +127,18 @@ Add the following to your prometheus.yml:
 
 ```yaml
 remote_write:
-    - url: "http://<ip address>:9201/write"
+  - url: "http://<ip address>:9201/write"
 remote_read:
-    - url: "http://<ip address>:9201/read"
- ```
+  - url: "http://<ip address>:9201/read"
+```
 
 ## Maintainers
 
 The PostgreSQL Prometheus Adapter is maintained by the team at [Crunchy Data](https://www.crunchydata.com/).
 
 ## Contributing to the Project
 
-Want to contribute to the PostgreSQL Prometheus Adapter? Great! Please use GitHub to submit an issue for the PostgreSQL Prometheus Adapter project.  If you would like to work the issue, please add that information in the issue so that we can confirm we are not already working no need to duplicate efforts.
+Want to contribute to the PostgreSQL Prometheus Adapter? Great! Please use GitHub to submit an issue for the PostgreSQL Prometheus Adapter project. If you would like to work the issue, please add that information in the issue so that we can confirm we are not already working no need to duplicate efforts.
 
 ## License
 

pkg/postgresql/client.go

@@ -16,6 +16,7 @@ import (
 	"github.com/go-kit/kit/log"
 	"github.com/go-kit/kit/log/level"
 	"github.com/jackc/pgx/v4"
+	"github.com/jackc/pgx/v4/pgxpool"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/prompb"
 )
@@ -40,7 +41,7 @@ var vMetricIDMap tMetricIDMap
 
 // PGWriter - Threaded writer
 type PGWriter struct {
-	DB          *pgx.Conn
+	DB          *pgxpool.Pool
 	id          int
 	KeepRunning bool
 	Running     bool
@@ -123,11 +124,13 @@ func (c *PGWriter) RunPGWriter(l log.Logger, tid int, commitSecs int, commitRows
 	period := commitSecs * 1000
 	var err error
 	var parser [20]PGParser
-	c.DB, err = pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
+	c.DB, err = pgxpool.Connect(context.Background(), os.Getenv("DATABASE_URL"))
 	if err != nil {
 		level.Error(c.logger).Log("err", err)
 		os.Exit(1)
 	}
+	defer c.DB.Close()
+
 	if c.id == 0 {
 		c.setupPgPrometheus()
 		_ = c.setupPgPartitions(partitionScheme, time.Now())
@@ -165,6 +168,7 @@ func (c *PGWriter) PGWriterSave() {
 	var copyCount, lblCount, rowCount int64
 	var err error
 	begin := time.Now()
+	vMetricIDMapMutex.Lock()
 	lblCount = int64(len(c.labelRows))
 	c.PGWriterMutex.Lock()
 	if lblCount > 0 {
@@ -181,6 +185,7 @@ func (c *PGWriter) PGWriterSave() {
 	rowCount = int64(len(c.valueRows))
 	c.valueRows = nil
 	c.PGWriterMutex.Unlock()
+	vMetricIDMapMutex.Unlock()
 	if err != nil {
 		level.Error(c.logger).Log("msg", "COPY failed for metric_values", "err", err)
 	}
@@ -212,7 +217,7 @@ func Pop() *model.Samples {
 // Client - struct to hold critical values
 type Client struct {
 	logger log.Logger
-	DB     *pgx.Conn
+	DB     *pgxpool.Pool
 	cfg    *Config
 }
 
@@ -222,7 +227,7 @@ func NewClient(logger log.Logger, cfg *Config) *Client {
 		logger = log.NewNopLogger()
 	}
 
-	conn1, err := pgx.Connect(context.Background(), os.Getenv("DATABASE_URL"))
+	conn1, err := pgxpool.Connect(context.Background(), os.Getenv("DATABASE_URL"))
 	if err != nil {
 		fmt.Fprintln(os.Stderr, "Error: Unable to connect to database using DATABASE_URL=", os.Getenv("DATABASE_URL"))
 		os.Exit(1)
@@ -369,9 +374,7 @@ func createOrderedKeys(m *map[string]string) []string {
 // Close - Close database connections
 func (c *Client) Close() {
 	if c.DB != nil {
-		if err1 := c.DB.Close(context.Background()); err1 != nil {
-			level.Error(c.logger).Log("msg", err1.Error())
-		}
+		c.DB.Close()
 	}
 }
 

start.sh

@@ -2,7 +2,7 @@
 
 if [[ "${DATABASE_URL}" == "" ]]; then
   echo 'Missing DATABASE_URL'
-  echo 'example -e DATABASE_URL="user=<db user> password=<db user password> host=<db host> port=<db port> database=<db name>"'
+  echo 'example -e DATABASE_URL="user=<db user> password=<db user password> host=<db host> port=<db port> database=<db name> pool_min_conns=[minimum pool connection] pool_max_conns=[maximum pool connection]"'
   exit 1
 fi
 
@@ -46,4 +46,3 @@ echo /postgresql-prometheus-adapter \
   --pg-commit-rows=${pg_commit_rows} \
   --pg-threads=${pg_threads} \
   --parser-threads=${parser_threads}
-