fix: add timeout to psql activity command and close connection (#660)

Author: agneum

engine/internal/retrieval/engine/postgres/logical/activity.go

@@ -43,6 +43,8 @@ func dbSourceActivity(ctx context.Context, dbCfg Connection) ([]activity.PGEvent
 		return nil, fmt.Errorf("failed to connect to DB: %w", err)
 	}
 
+	defer func() { _ = querier.Close(context.Background()) }()
+
 	rows, err := querier.Query(ctx, statQuery)
 	if err != nil {
 		return nil, fmt.Errorf("failed to perform query to get DB activity: %w", err)
@@ -77,7 +79,10 @@ func pgContainerActivity(ctx context.Context, docker *client.Client, containerID
 		return []activity.PGEvent{}, nil
 	}
 
-	activityCmd := []string{"psql", "-U", db.User(), "-d", db.Name(), "--record-separator='" + customRecordSeparator + "'", "-XAtc", statQuery}
+	activityCmd := []string{"psql", "-U", db.User(), "-d", db.Name(),
+		"--set=connect_timeout=5",
+		"-c", "SET statement_timeout = '10s';",
+		"--record-separator='" + customRecordSeparator + "'", "-XAtc", statQuery}
 
 	log.Msg("Running activity command: ", activityCmd)
 

engine/internal/retrieval/engine/postgres/tools/tools.go

@@ -624,6 +624,10 @@ func processAttachResponse(ctx context.Context, reader io.Reader) ([]byte, error
 		break
 
 	case <-ctx.Done():
+		if closer, ok := reader.(io.Closer); ok {
+			_ = closer.Close()
+		}
+
 		return nil, ctx.Err()
 	}
 

engine/test/_cleanup.sh

@@ -3,7 +3,7 @@ set -euxo pipefail
 
 DLE_TEST_MOUNT_DIR="/var/lib/test/dblab_mount"
 DLE_TEST_POOL_NAME="test_dblab_pool"
-TMP_DATA_DIR="/tmp/dle_test/logical_generic"
+TMP_DATA_DIR="/tmp/dle_test"
 ZFS_FILE="$(pwd)/zfs_file"
 
 # Stop and remove test Docker containers