[+] add pg_timetable.current_chain_id and pg_timetable.current_task_id run-time parameters, closes #438

Author: pashagolub

.vscode/launch.json

@@ -11,7 +11,9 @@
             "mode": "auto",
             "program": "${workspaceFolder}",
             "env": {},
-            "args": ["-v"]
+            "args": ["postgresql://scheduler@localhost:5432/timetable",
+                "--clientname=loader", 
+                "--log-level=debug"]
         }
     ]
 }

internal/pgengine/pgengine_test.go

@@ -144,17 +144,19 @@ func TestSchedulerFunctions(t *testing.T) {
 
 	t.Run("Check GetChainElements funсtion", func(t *testing.T) {
 		var chains []pgengine.ChainTask
-		tx, err := pge.StartTransaction(ctx)
+		tx, txid, err := pge.StartTransaction(ctx, 0)
 		assert.NoError(t, err, "Should start transaction")
+		assert.Greater(t, txid, 0, "Should return transaction id")
 		assert.True(t, pge.GetChainElements(ctx, tx, &chains, 0), "Should no error in clean database")
 		assert.Empty(t, chains, "Should be empty in clean database")
 		pge.CommitTransaction(ctx, tx)
 	})
 
 	t.Run("Check GetChainParamValues funсtion", func(t *testing.T) {
 		var paramVals []string
-		tx, err := pge.StartTransaction(ctx)
+		tx, txid, err := pge.StartTransaction(ctx, 0)
 		assert.NoError(t, err, "Should start transaction")
+		assert.Greater(t, txid, 0, "Should return transaction id")
 		assert.True(t, pge.GetChainParamValues(ctx, tx, &paramVals, &pgengine.ChainTask{
 			TaskID:  0,
 			ChainID: 0}), "Should no error in clean database")
@@ -170,8 +172,9 @@ func TestSchedulerFunctions(t *testing.T) {
 	})
 
 	t.Run("Check ExecuteSQLCommand function", func(t *testing.T) {
-		tx, err := pge.StartTransaction(ctx)
+		tx, txid, err := pge.StartTransaction(ctx, 0)
 		assert.NoError(t, err, "Should start transaction")
+		assert.Greater(t, txid, 0, "Should return transaction id")
 		f := func(sql string, params []string) error {
 			_, err := pge.ExecuteSQLCommand(ctx, tx, sql, params)
 			return err

internal/pgengine/transaction.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -31,9 +32,15 @@ type ChainTask struct {
 	Txid          int
 }
 
-// StartTransaction return transaction object and panic in the case of error
-func (pge *PgEngine) StartTransaction(ctx context.Context) (pgx.Tx, error) {
-	return pge.ConfigDb.Begin(ctx)
+// StartTransaction returns transaction object, transaction id and error
+func (pge *PgEngine) StartTransaction(ctx context.Context, chainID int) (tx pgx.Tx, txid int, err error) {
+	tx, err = pge.ConfigDb.Begin(ctx)
+	if err != nil {
+		return
+	}
+	err = pgxscan.Get(ctx, tx, &txid, "SELECT txid_current()")
+	_, err = tx.Exec(ctx, `SELECT set_config('pg_timetable.current_chain_id', $1, true)`, strconv.Itoa(chainID))
+	return
 }
 
 // CommitTransaction commits transaction and log error in the case of error
@@ -128,15 +135,14 @@ func (pge *PgEngine) ExecuteSQLTask(ctx context.Context, tx pgx.Tx, task *ChainT
 		defer pge.FinalizeRemoteDBConnection(ctx, remoteDb)
 	}
 
-	// Set Role
-	if task.RunAs.Status != pgtype.Null && !task.Autonomous {
+	if !task.Autonomous {
 		pge.SetRole(ctx, execTx, task.RunAs)
+		if task.IgnoreError {
+			pge.MustSavepoint(ctx, execTx, fmt.Sprintf("task_%d", task.TaskID))
+		}
 	}
 
-	if task.IgnoreError && !task.Autonomous {
-		pge.MustSavepoint(ctx, execTx, fmt.Sprintf("task_%d", task.TaskID))
-	}
-
+	pge.SetCurrentTaskContext(ctx, execTx, task.TaskID)
 	out, err = pge.ExecuteSQLCommand(ctx, executor, task.Script, paramValues)
 
 	if err != nil && task.IgnoreError && !task.Autonomous {
@@ -223,6 +229,9 @@ func (pge *PgEngine) FinalizeRemoteDBConnection(ctx context.Context, remoteDb Pg
 
 // SetRole - set the current user identifier of the current session
 func (pge *PgEngine) SetRole(ctx context.Context, tx pgx.Tx, runUID pgtype.Varchar) {
+	if runUID.Status == pgtype.Null {
+		return
+	}
 	l := log.GetLogger(ctx)
 	l.Info("Setting Role to ", runUID.String)
 	_, err := tx.Exec(ctx, fmt.Sprintf("SET ROLE %v", runUID.String))
@@ -238,6 +247,16 @@ func (pge *PgEngine) ResetRole(ctx context.Context, tx pgx.Tx) {
 	const sqlResetRole = `RESET ROLE`
 	_, err := tx.Exec(ctx, sqlResetRole)
 	if err != nil {
-		l.WithError(err).Error("Error in ReSetting role", err)
+		l.WithError(err).Error("Failed to set a role", err)
+	}
+}
+
+// SetCurrentTaskContext - set the working transaction "pg_timetable.current_task_id" run-time parameter
+func (pge *PgEngine) SetCurrentTaskContext(ctx context.Context, tx pgx.Tx, taskID int) {
+	l := log.GetLogger(ctx)
+	l.Debug("Setting current task context to ", taskID)
+	_, err := tx.Exec(ctx, "SELECT set_config('pg_timetable.current_task_id', $1, true)", strconv.Itoa(taskID))
+	if err != nil {
+		l.WithError(err).Error("Failed to set current task context", err)
 	}
 }

internal/scheduler/chain.go

@@ -7,7 +7,6 @@ import (
 
 	"github.com/cybertec-postgresql/pg_timetable/internal/log"
 	"github.com/cybertec-postgresql/pg_timetable/internal/pgengine"
-	"github.com/georgysavva/scany/pgxscan"
 	pgx "github.com/jackc/pgx/v4"
 )
 
@@ -178,10 +177,7 @@ func (sch *Scheduler) executeChain(ctx context.Context, chain Chain) {
 
 	chainL := sch.l.WithField("chain", chain.ChainID)
 
-	tx, err := sch.pgengine.StartTransaction(ctx)
-	if err == nil {
-		err = pgxscan.Get(ctx, tx, &txid, "SELECT txid_current()")
-	}
+	tx, txid, err := sch.pgengine.StartTransaction(ctx, chain.ChainID)
 	if err != nil {
 		chainL.WithError(err).Error("Cannot start transaction")
 		return

samples/Chain.sql

@@ -18,6 +18,7 @@ BEGIN
         PRIMARY KEY (chain_log)
     );
 
+    -- Let's create a new chain and add tasks to it later
     INSERT INTO timetable.chain (
         chain_id, 
         chain_name, 
@@ -46,12 +47,27 @@ BEGIN
     VALUES (v_chain_id, 2, 'INSERT INTO timetable.chain_log (EVENT, time) VALUES ($1, CURRENT_TIMESTAMP)', TRUE)
     RETURNING task_id INTO v_task_id;
 
-
     INSERT INTO timetable.parameter(task_id, order_id, value)
     VALUES 
         -- Parameter for HEAD (parent) task
         (v_parent_id, 1, '["Added"]' :: jsonb),
         -- Parameter for the next task
         (v_task_id, 1, '["Updated"]' :: jsonb);
+
+    -- Add one more task swowing IDs for all tasks within the chain
+    INSERT INTO timetable.task (chain_id, task_order, command, ignore_error)
+    VALUES (v_chain_id, 3, 
+    $CMD$
+        DO $BODY$ 
+        DECLARE tasks TEXT;
+        BEGIN 
+            SELECT array_agg(task_id ORDER BY task_order) FROM timetable.task
+            INTO tasks
+            WHERE chain_id = current_setting('pg_timetable.current_chain_id')::bigint;
+            RAISE NOTICE 'Task IDs in chain: %', tasks;
+        END; 
+        $BODY$
+    $CMD$, TRUE);
+   
 END;
 $$ LANGUAGE plpgsql;