From da2d4908144fbf840adc883e0647dc944f13ca33 Mon Sep 17 00:00:00 2001
From: Adw8 <adwaitgharpure@gmail.com>
Date: Mon, 27 Jan 2025 18:47:36 +0530
Subject: [PATCH] Add method to initialize DB

---
 .gitignore                 |  2 ++
 src/app/api/flux/route.ts  | 15 +++++++--------
 src/data-source.ts         | 16 ++++++++++------
 src/utils/verifyPayment.ts |  7 +++++--
 4 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index 28028d2..6768f48 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,5 @@ node_modules
 .env.development
 .env.production
 .env.test
+
+database.sqlite
diff --git a/src/app/api/flux/route.ts b/src/app/api/flux/route.ts
index c7c674e..17f0f89 100644
--- a/src/app/api/flux/route.ts
+++ b/src/app/api/flux/route.ts
@@ -1,6 +1,7 @@
 import { NextRequest, NextResponse } from 'next/server'
 import { fal } from "@fal-ai/client"
 import { FLUX_MODELS } from '../../../services/fluxService'
+import { initializeDataSource } from '../../../data-source'
 
 import { verifyPayment, isSignatureUsed, markSignatureAsUsed } from '../../../utils/verifyPayment'
 
@@ -19,6 +20,8 @@ const IMAGE_HEIGHT: number = 1024
 
 export async function POST(req: NextRequest): Promise<NextResponse> {
   try {
+    await initializeDataSource();
+
     const { prompt, modelId, transactionSignature } = await req.json()
 
     if (!prompt || !modelId) {
@@ -35,13 +38,6 @@ export async function POST(req: NextRequest): Promise<NextResponse> {
       )
     }
 
-    if (await isSignatureUsed(transactionSignature)) {
-      return NextResponse.json(
-        { error: 'Transaction signature has already been used' },
-        { status: 400 }
-      )
-    }
-
     const model = FLUX_MODELS.find((model) => model.modelId === modelId)
     if (!model) {
       return NextResponse.json(
@@ -54,7 +50,10 @@ export async function POST(req: NextRequest): Promise<NextResponse> {
     const isPaymentVerified = await verifyPayment(transactionSignature, expectedAmount)
 
     if (!isPaymentVerified) {
-      throw new Error('Payment verification failed');
+      return NextResponse.json(
+        { error: 'Payment verification failed or transaction signature has already been used' },
+        { status: 400 }
+      )
     }
 
     await markSignatureAsUsed(transactionSignature)
diff --git a/src/data-source.ts b/src/data-source.ts
index ffc689a..6f5f10c 100644
--- a/src/data-source.ts
+++ b/src/data-source.ts
@@ -11,10 +11,14 @@ export const AppDataSource = new DataSource({
   subscribers: [],
 });
 
-AppDataSource.initialize()
-  .then(() => {
-    console.log('Data Source has been initialized!');
-  })
-  .catch((err) => {
+export async function initializeDataSource() {
+  try {
+    if (!AppDataSource.isInitialized){
+      await AppDataSource.initialize();
+      console.log('Data Source has been initialized!');
+    };
+  } catch (err) {
     console.error('Error during Data Source initialization:', err);
-  });
+    throw err;
+  }
+}
diff --git a/src/utils/verifyPayment.ts b/src/utils/verifyPayment.ts
index 659ce93..8bd827f 100644
--- a/src/utils/verifyPayment.ts
+++ b/src/utils/verifyPayment.ts
@@ -19,7 +19,10 @@ const connection = new Connection(
 export async function isSignatureUsed(transactionSignature: string): Promise<boolean> {
   const paymentRepository = AppDataSource.getRepository(Payment);
   const payment = await paymentRepository.findOneBy({ transactionSignature });
-  return !!payment;
+  if (payment) {
+    return true;
+  }
+  return false;
 }
 
 export async function markSignatureAsUsed(transactionSignature: string): Promise<void> {
@@ -39,7 +42,7 @@ export async function verifyPayment(
       return false;
     }
 
-    const transaction = await connection.getParsedTransaction(transactionSignature, 'finalized');
+    const transaction = await connection.getParsedTransaction(transactionSignature, 'confirmed');
     if (!transaction) {
       throw new Error('Transaction not found');
     }