diff --git a/demo/example/lib/exampleStack.ts b/demo/example/lib/exampleStack.ts
index 6773514..763da2d 100755
--- a/demo/example/lib/exampleStack.ts
+++ b/demo/example/lib/exampleStack.ts
@@ -35,15 +35,7 @@ export class ExampleStack extends Stack {
         id: 'default',
         migrationHandling: {
           type: 'lambda',
-          migrationLambdaFunction: new NodejsFunction(
-            this,
-            'MigrationLambdaFunction',
-            {
-              runtime: Runtime.NODEJS_18_X,
-              handler: 'index.handler',
-              entry: path.join(__dirname, 'runMigrationsLambda/index.ts'),
-            },
-          ),
+          migrationLambdaFunction,
         },
       },
     });
diff --git a/demo/example/lib/runMigrationsLambda/migrations/migration1.ts b/demo/example/lib/runMigrationsLambda/migrations/migration1.ts
index 54b8c8c..a527a8e 100644
--- a/demo/example/lib/runMigrationsLambda/migrations/migration1.ts
+++ b/demo/example/lib/runMigrationsLambda/migrations/migration1.ts
@@ -6,27 +6,57 @@ const client = new DynamoDB({ region: 'eu-west-1' });
 export const migration1: Migration = {
   id: 1,
   up: async (): Promise<{ status: string }> => {
+    console.log('MIGRATION_1_UP');
+
     try {
       const elements = await client.scan({ TableName: 'Dinosaurs' });
-      console.log('ELEMENTS', elements);
+      console.log('ELEMENTS_1_UP', elements);
+
+      const modifiedElements = elements.Items?.map((element) => ({
+        ...element,
+        eyeColor: { S: 'green' },
+      }));
+
+      if (modifiedElements === undefined) {
+        return { status: 'SUCCEEDED' };
+      }
+      await Promise.all(
+        modifiedElements.map((element) =>
+          client.putItem({ TableName: 'Dinosaurs', Item: element }),
+        ),
+      );
     } catch (err) {
       console.log('ERROR', err);
+
+      return Promise.resolve({ status: 'FAILED' });
     }
-    console.log('Some up migration happening here', 1);
+
     // Modify the element and put it back in db
-    return Promise.resolve({
-      status: Math.random() < 0.01 ? 'FAILED' : 'SUCCEEDED',
-    });
+    return { status: 'SUCCEEDED' };
   },
   down: async (): Promise<{ status: string }> => {
-    console.log('Some down migration happening here', 1);
+    console.log('MIGRATION_1_DOWN');
     try {
       const elements = await client.scan({ TableName: 'Dinosaurs' });
       console.log('ELEMENTS down', elements);
+
+      const modifiedElements = elements.Items?.map(
+        ({ eyeColor, ...rest }) => rest,
+      );
+
+      if (modifiedElements === undefined) {
+        return { status: 'SUCCEEDED' };
+      }
+      await Promise.all(
+        modifiedElements.map((element) =>
+          client.putItem({ TableName: 'Dinosaurs', Item: element }),
+        ),
+      );
     } catch (err) {
       console.log('ERROR down', err);
+
+      return Promise.resolve({ status: 'FAILED' });
     }
-    console.log('Some up migration happening here', 1);
 
     return Promise.resolve({ status: 'SUCCEEDED' });
   },
diff --git a/demo/example/lib/runMigrationsLambda/migrations/migration2.ts b/demo/example/lib/runMigrationsLambda/migrations/migration2.ts
index e984f60..46e98c7 100644
--- a/demo/example/lib/runMigrationsLambda/migrations/migration2.ts
+++ b/demo/example/lib/runMigrationsLambda/migrations/migration2.ts
@@ -1,19 +1,71 @@
+import { DynamoDB } from '@aws-sdk/client-dynamodb';
 import { Migration } from 'migration-version-helpers';
 
+const client = new DynamoDB({ region: 'eu-west-1' });
+
 export const migration2: Migration = {
   id: 2,
   up: async (): Promise<{ status: string }> => {
-    console.log('Some up migration happening here', 2);
+    console.log('MIGRATION_2_UP');
+    try {
+      const elements = await client.scan({ TableName: 'Dinosaurs' });
+
+      const modifiedElements = elements.Items?.map(({ eyeColor, ...rest }) => ({
+        ...rest,
+        eyeColor: { L: [eyeColor ?? { S: 'yellow' }] },
+      }));
+
+      if (modifiedElements === undefined) {
+        return { status: 'SUCCEEDED' };
+      }
+      await Promise.all(
+        modifiedElements.map((element) =>
+          client.putItem({ TableName: 'Dinosaurs', Item: element }),
+        ),
+      );
+    } catch (err) {
+      console.log('ERROR', err);
+
+      return Promise.resolve({ status: 'FAILED' });
+    }
+    console.log('Some up migration happening here', 1);
 
-    return Promise.resolve({
-      status: Math.random() > 0.001 ? 'SUCCEEDED' : 'FAILED',
-    });
+    // Modify the element and put it back in db
+    return { status: 'SUCCEEDED' };
   },
   down: async (): Promise<{ status: string }> => {
-    console.log('Some down migration happening here', 2);
+    console.log('MIGRATION_2_DOWN');
+
+    try {
+      const elements = await client.scan({ TableName: 'Dinosaurs' });
+      console.log('ELEMENTS down', elements.Items);
+
+      const modifiedElements = elements.Items?.map((eyeColor, ...rest) => ({
+        ...rest,
+        // @ts-expect-error AttributeValue type is not correct
+        // eslint-disable-next-line @typescript-eslint/strict-boolean-expressions
+        ...(eyeColor.L?.[0]
+          ? {
+              eyeColor: { S: eyeColor.L[0].S },
+            }
+          : {}),
+      }));
+
+      if (modifiedElements === undefined) {
+        return { status: 'SUCCEEDED' };
+      }
+      await Promise.all(
+        modifiedElements.map((element) =>
+          client.putItem({ TableName: 'Dinosaurs', Item: element }),
+        ),
+      );
+    } catch (err) {
+      console.log('ERROR down', err);
+
+      return Promise.resolve({ status: 'FAILED' });
+    }
+    console.log('Some up migration happening here', 1);
 
-    return Promise.resolve({
-      status: Math.random() > 0.5 ? 'SUCCEEDED' : 'FAILED',
-    });
+    return Promise.resolve({ status: 'SUCCEEDED' });
   },
 };