From 0d1396b8dec859439252ebec6251a889659ddc19 Mon Sep 17 00:00:00 2001
From: bros <arnaubennassar5@gmail.com>
Date: Fri, 2 Feb 2024 15:43:37 +0000
Subject: [PATCH] Fix lost data

---
 db/migrations/state/validium-001.sql | 32 ++++++++++++++++++++++++++++
 state/pgstatestorage.go              | 20 +++++++++++++++++
 state/pgstatestorage_test.go         | 27 +++++++++++++++++++++++
 3 files changed, 79 insertions(+)
 create mode 100644 db/migrations/state/validium-001.sql

diff --git a/db/migrations/state/validium-001.sql b/db/migrations/state/validium-001.sql
new file mode 100644
index 0000000000..74b867d31d
--- /dev/null
+++ b/db/migrations/state/validium-001.sql
@@ -0,0 +1,32 @@
+-- +migrate Up
+
+CREATE TABLE IF NOT EXISTS state.batch_data_backup
+(
+    batch_num  BIGINT,
+    data       BYTEA,
+    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT CURRENT_TIMESTAMP,
+    PRIMARY KEY (batch_num, created_at)
+);
+
+-- +migrate StatementBegin
+CREATE OR REPLACE FUNCTION backup_batch() RETURNS trigger AS $$
+    BEGIN
+        INSERT  INTO state.batch_data_backup (batch_num,     data) 
+                VALUES                       (OLD.batch_num, OLD.raw_txs_data)
+                ON CONFLICT (batch_num, created_at) DO UPDATE SET
+                    data = EXCLUDED.data;
+        RETURN OLD;
+    END;
+$$
+LANGUAGE plpgsql;
+-- +migrate StatementEnd
+
+CREATE TRIGGER backup_batch
+   BEFORE DELETE ON state.batch FOR EACH ROW
+   EXECUTE PROCEDURE backup_batch();
+
+-- +migrate Down
+
+DROP TRIGGER IF EXISTS backup_batch ON state.batch;
+DROP FUNCTION IF EXISTS backup_batch();
+DROP TABLE IF EXISTS state.batch_data_backup;
diff --git a/state/pgstatestorage.go b/state/pgstatestorage.go
index 6f38db4462..2bcca13335 100644
--- a/state/pgstatestorage.go
+++ b/state/pgstatestorage.go
@@ -2912,6 +2912,26 @@ func (p *PostgresStorage) GetBatchL2DataByNumber(ctx context.Context, batchNumbe
 	var batchL2Data []byte
 	err := q.QueryRow(ctx, getBatchL2DataByBatchNumber, batchNumber).Scan(&batchL2Data)
 
+	if errors.Is(err, pgx.ErrNoRows) {
+		return p.GetBatchL2DataByNumberFromBackup(ctx, batchNumber, dbTx)
+	} else if err != nil {
+		return nil, err
+	}
+	return batchL2Data, nil
+}
+
+// GetBatchL2DataByNumberFromBackup returns the batch L2 data of the given batch number from the backup table
+func (p *PostgresStorage) GetBatchL2DataByNumberFromBackup(ctx context.Context, batchNumber uint64, dbTx pgx.Tx) ([]byte, error) {
+	getBatchL2DataByBatchNumber := `
+		SELECT data FROM state.batch_data_backup
+		WHERE batch_num = $1
+		ORDER BY created_at DESC
+		LIMIT 1
+	`
+	q := p.getExecQuerier(dbTx)
+	var batchL2Data []byte
+	err := q.QueryRow(ctx, getBatchL2DataByBatchNumber, batchNumber).Scan(&batchL2Data)
+
 	if errors.Is(err, pgx.ErrNoRows) {
 		return nil, ErrNotFound
 	} else if err != nil {
diff --git a/state/pgstatestorage_test.go b/state/pgstatestorage_test.go
index d873d1dbc8..386c9d8007 100644
--- a/state/pgstatestorage_test.go
+++ b/state/pgstatestorage_test.go
@@ -993,4 +993,31 @@ func TestGetBatchL2DataByNumber(t *testing.T) {
 	actualData, err := testState.GetBatchL2DataByNumber(ctx, batchNum, tx)
 	require.NoError(t, err)
 	assert.Equal(t, expectedData, actualData)
+
+	// Force backup
+	_, err = tx.Exec(ctx, "DELETE FROM state.batch")
+	require.NoError(t, err)
+
+	// Get batch 4 from backup
+	batchNum = 4
+	data, err = testState.GetBatchL2DataByNumber(ctx, batchNum, tx)
+	require.NoError(t, err)
+	assert.Nil(t, data)
+
+	// Get batch 5 from backup
+	batchNum = 5
+	actualData, err = testState.GetBatchL2DataByNumber(ctx, batchNum, tx)
+	require.NoError(t, err)
+	assert.Equal(t, expectedData, actualData)
+
+	// Update batch 5 and get it from backup
+	expectedData = []byte("new foo bar")
+	_, err = tx.Exec(ctx, openBatchSQL, batchNum, expectedData)
+	require.NoError(t, err)
+	_, err = tx.Exec(ctx, "DELETE FROM state.batch")
+	require.NoError(t, err)
+	actualData, err = testState.GetBatchL2DataByNumber(ctx, batchNum, tx)
+	require.NoError(t, err)
+	assert.Equal(t, expectedData, actualData)
+
 }