[WIP] Feature: Refactor api storage

packages/api/src/graphql/client.ts

@@ -17,6 +17,7 @@ import { ownership } from "./ownership";
 import { permission } from "./permission";
 import { proof } from "./proof";
 import { user } from "./user";
+import { transaction } from "./transaction";
 
 export interface IApiClient<T = any> {
   api: ApiClient<T>;
@@ -30,6 +31,8 @@ export interface IApiClient<T = any> {
   permission: ReturnType<typeof permission>;
   merkle: ReturnType<typeof merkle>;
   proof: ReturnType<typeof proof>;
+
+  transaction: ReturnType<typeof transaction>;
 }
 
 export class ApiClient<T = any> {
@@ -87,7 +90,6 @@ export class ApiClient<T = any> {
         },
       };
     });
-
 
     this.#client = new ApolloClient({
       link: ApolloLink.from([removeTypenameLink, authLink, httpLink]),
@@ -108,6 +110,7 @@ export class ApiClient<T = any> {
       permission: permission(api.apollo),
       merkle: merkle(api.apollo),
       proof: proof(api.apollo),
+      transaction: transaction(api.apollo),
     };
   }
 }

packages/api/src/graphql/database.ts

@@ -19,16 +19,8 @@ const DATABASE_CHANGE_OWNER = gql`
 `;
 
 const DATABASE_CREATE = gql`
-  mutation DbCreate(
-    $databaseName: String!
-    $merkleHeight: Int!
-    $publicKey: String!
-  ) {
-    dbCreate(
-      databaseName: $databaseName
-      merkleHeight: $merkleHeight
-      publicKey: $publicKey
-    )
+  mutation DbCreate($databaseName: String!, $merkleHeight: Int!) {
+    dbCreate(databaseName: $databaseName, merkleHeight: $merkleHeight)
   }
 `;
 
@@ -106,7 +98,7 @@ export const database = <T>(client: TApolloClient<T>) => ({
   >(client, DATABASE_CHANGE_OWNER, (data) => data.dbChangeOwner),
   create: createMutateFunction<
     boolean,
-    { databaseName: string; merkleHeight: number; publicKey: string },
+    { databaseName: string; merkleHeight: number },
     { dbCreate: boolean }
   >(client, DATABASE_CREATE, (data) => data.dbCreate),
   setting: createQueryFunction<

packages/api/src/graphql/transaction.ts

@@ -0,0 +1,28 @@
+import { gql } from "@apollo/client";
+import { createQueryFunction, TApolloClient } from "./common.js";
+import { TDbTransaction, TTransactionRequest } from "./types/transaction.js";
+
+const TRANSACTION_GET = gql`
+  query GetTransaction(
+    $databaseName: String!
+    $transactionType: TransactionType!
+  ) {
+    getTransaction(
+      databaseName: $databaseName
+      transactionType: $transactionType
+    ) {
+      databaseName
+      transactionType
+      tx
+      zkAppPublicKey
+    }
+  }
+`;
+
+export const transaction = <T>(client: TApolloClient<T>) => ({
+  getTransaction: createQueryFunction<
+    TDbTransaction,
+    TTransactionRequest,
+    { getTransaction: TDbTransaction }
+  >(client, TRANSACTION_GET, (data) => data.getTransaction),
+});

packages/api/src/graphql/types/transaction.ts

@@ -0,0 +1,20 @@
+export type TTransactionRequest = {
+  databaseName: string;
+  transactionType: "deploy" | "rollup";
+};
+export type TTransactionType = "deploy" | "rollup";
+export type TTransactionStatus =
+  | "start"
+  | "ready"
+  | "pending"
+  | "failed"
+  | "success"
+  | "unknown";
+export type TDbTransaction = {
+  id: string;
+  databaseName: string;
+  status: TTransactionStatus;
+  transactionType: TTransactionType;
+  tx: string;
+  zkAppPublicKey: string;
+};

packages/broker-service/src/apollo/app/task.ts

@@ -1,5 +1,5 @@
-import GraphQLJSON from 'graphql-type-json';
-import { getNextTaskId } from '../../domain/get-next-task.js';
+import GraphQLJSON from "graphql-type-json";
+import { getNextTaskId } from "../../domain/get-next-task.js";
 
 export const typeDefsTask = `#graphql
 scalar JSON
@@ -16,5 +16,5 @@ export const resolversTask = {
   JSON: GraphQLJSON,
   Query: {
     taskId,
-  }
-};
+  },
+};

packages/broker-service/src/domain/get-next-task.ts

@@ -1,7 +1,12 @@
-import { ModelQueueTask, withTransaction } from "@zkdb/storage";
+import { ModelQueueTask, TransactionManager } from "@zkdb/storage";
+import { proofDb } from "../helper/config.js";
 
 export async function getNextTaskId(): Promise<string | null> {
-  return withTransaction(async (session) => {
+  TransactionManager.addSession({
+    name: "proof",
+    session: proofDb.client.startSession(),
+  });
+  return TransactionManager.withSingleTransaction("proof", async (session) => {
     const modelQueueTask = ModelQueueTask.getInstance();
     const task = await modelQueueTask.getLatestQueuedTaskByDatabase(session);
 
@@ -11,5 +16,5 @@ export async function getNextTaskId(): Promise<string | null> {
     }
 
     return null;
-  }, "proof");
+  });
 }

packages/broker-service/src/helper/config.ts

@@ -1,4 +1,5 @@
 import { ConfigLoader, Joi } from "@orochi-network/framework";
+import { DatabaseEngine } from "@zkdb/storage";
 export const NODE_ENV_VALUES = [
   "development",
   "production",
@@ -45,3 +46,4 @@ const configLoader = new ConfigLoader<TApplicationConfig>(
 );
 
 export const { config } = configLoader;
+export const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);

packages/broker-service/src/index.ts

@@ -1,26 +1,24 @@
-import { DatabaseEngine } from "@zkdb/storage";
+import {
+  DatabaseEngine,
+  ModelQueueTask,
+  TransactionManager,
+} from "@zkdb/storage";
 import { ApolloServer } from "apollo-server-express";
 import express from "express";
 import http from "http";
 import { ResolversApp, TypedefsApp } from "./apollo/index.js";
-import { config } from "./helper/config.js";
+import { config, proofDb } from "./helper/config.js";
 import logger from "./helper/logger.js";
 
 (async () => {
   const app = express();
 
-  // DB service
-  const serviceDb = DatabaseEngine.getInstance(config.MONGODB_URL);
-  // DB proof
-  const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);
-  if (!serviceDb.isConnected()) {
-    await serviceDb.connect();
-  }
-
   if (!proofDb.isConnected()) {
     await proofDb.connect();
   }
 
+  ModelQueueTask.createModel(proofDb);
+
   app.use(express.json());
 
   const httpServer = http.createServer(app);

packages/compile-service/src/helper/model-loader.ts

@@ -0,0 +1,26 @@
+import {
+  DatabaseEngine,
+  ModelDbSetting,
+  ModelDbTransaction,
+  ModelProof,
+  ModelSecureStorage,
+} from "@zkdb/storage";
+import { config } from "./config";
+
+export async function initModelLoader() {
+  const serviceDb = DatabaseEngine.getInstance(config.MONGODB_URL);
+  const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);
+
+  if (!serviceDb.isConnected()) {
+    await serviceDb.connect();
+  }
+
+  if (!proofDb.isConnected()) {
+    await proofDb.connect();
+  }
+
+  ModelDbTransaction.createModel(serviceDb);
+  ModelDbSetting.createModel(serviceDb);
+  ModelProof.createModel(proofDb);
+  ModelSecureStorage.createModel(proofDb);
+}

packages/compile-service/src/index.ts

@@ -2,16 +2,16 @@ import { logger } from "@helper";
 import { EncryptionKey } from "@orochi-network/vault";
 import { UnsignedTransaction, ZkCompileService } from "@service";
 import {
-  DatabaseEngine,
   ModelDbSetting,
   ModelDbTransaction,
   ModelProof,
   ModelSecureStorage,
 } from "@zkdb/storage";
 import { PrivateKey, PublicKey } from "o1js";
+import { setTimeout } from "timers/promises";
 import { config } from "./helper/config";
+import { initModelLoader } from "./helper/model-loader";
 import { RedisQueueService } from "./message-queue";
-import { setTimeout } from 'timers/promises';
 
 export type TransactionType = "deploy" | "rollup";
 
@@ -60,16 +60,7 @@ async function findTransactionWithRetry(
     { url: config.REDIS_URL }
   );
   // Connect to db
-  const serviceDb = DatabaseEngine.getInstance(config.MONGODB_URL);
-  const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);
-
-  if (!serviceDb.isConnected()) {
-    await serviceDb.connect();
-  }
-
-  if (!proofDb.isConnected()) {
-    await proofDb.connect();
-  }
+  await initModelLoader();
 
   const modelTransaction = ModelDbTransaction.getInstance();
   const modelDbSettings = ModelDbSetting.getInstance();

packages/compile-service/src/service/zk-compile.ts

@@ -72,7 +72,7 @@ export class ZkCompileService {
 
       return partialSignedTx.toJSON();
     } catch (error) {
-      logger.error(`Cannot compile & deploy: ${databaseName}`, logger);
+      logger.error(`Cannot compile & deploy: ${databaseName}`);
       await ModelDbSetting.getInstance().updateSetting(databaseName, {
         appPublicKey: undefined,
       });

packages/proof-service/src/domain/create-proof.ts

@@ -8,7 +8,7 @@ import {
   ModelMerkleTree,
   ModelProof,
   ModelQueueTask,
-  withTransaction,
+  TransactionManager,
 } from '@zkdb/storage';
 import { ObjectId } from 'mongodb';
 import {
@@ -119,7 +119,7 @@ export async function createProof(taskId: string) {
       } else {
         const rollupProof = await modelProof.findOne({
           merkleRoot: onChainRootState.toString(),
-          database: task.database
+          database: task.database,
         });
         if (rollupProof) {
           proof = await circuit.updateTransition(
@@ -150,7 +150,7 @@ export async function createProof(taskId: string) {
 
     // TODO: Should we consider both on-chain action and off-chain leaf. Off-chain leaf = On-chain action
 
-    await withTransaction(async (session) => {
+    await TransactionManager.withSingleTransaction('proof', async (session) => {
       await modelProof.saveProof(
         {
           ...proof.toJSON(),
@@ -162,7 +162,7 @@ export async function createProof(taskId: string) {
         { session }
       );
       await queue.markTaskProcessed(task._id, { session });
-    }, 'proof');
+    });
 
     logger.debug('Task processed successfully.');
   } catch (error) {

packages/proof-service/src/helper/model-loader.ts

@@ -0,0 +1,37 @@
+import {
+  DatabaseEngine,
+  ModelDbSetting,
+  ModelDbTransaction,
+  ModelMerkleTree,
+  ModelProof,
+  ModelQueueTask,
+  ModelSecureStorage,
+  TransactionManager,
+} from '@zkdb/storage';
+import { config } from './config.js';
+
+export async function initModelLoader() {
+  // db service
+  const serviceDb = DatabaseEngine.getInstance(config.MONGODB_URL);
+  // db proof
+  const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);
+  console.log('🚀 ~ initModelLoader ~ L:', config.PROOF_MONGODB_URL);
+  if (!serviceDb.isConnected()) {
+    await serviceDb.connect();
+  }
+
+  if (!proofDb.isConnected()) {
+    await proofDb.connect();
+  }
+
+  TransactionManager.addSession({
+    name: 'proof',
+    session: proofDb.client.startSession(),
+  });
+  console.log('🚀 ~ initModelLoader ~ TransactionManager:', TransactionManager);
+
+  ModelDbSetting.createModel(serviceDb);
+  ModelMerkleTree.createModel(serviceDb);
+  ModelQueueTask.createModel(proofDb);
+  ModelProof.createModel(proofDb);
+}

packages/proof-service/src/index.ts

@@ -1,31 +1,20 @@
-import { DatabaseEngine } from '@zkdb/storage';
 import TaskService from './service/task-service.js';
 import logger from './helper/logger.js';
 import { config } from './helper/config.js';
 import { Mina } from 'o1js';
+import { initModelLoader } from './helper/model-loader.js';
 
 (async () => {
+  await initModelLoader();
+
   const network = Mina.Network({
     networkId: config.NETWORK_ID,
     mina: config.MINA_URL,
   });
 
   Mina.setActiveInstance(network);
 
-  // db service
-  const serviceDb = DatabaseEngine.getInstance(config.MONGODB_URL);
-  // db proof
-  const proofDb = DatabaseEngine.getInstance(config.PROOF_MONGODB_URL);
-  if (!serviceDb.isConnected()) {
-    await serviceDb.connect();
-  }
-
-  if (!proofDb.isConnected()) {
-    await proofDb.connect();
-  }
-
   const taskService = new TaskService();
-
   await taskService.fetchAndProcessTasks();
 
   logger.info('Proof service stopped.');

packages/proof-service/src/service/task-service.ts

@@ -19,7 +19,7 @@ class TaskService {
     let retries = 0;
     let delay = this.initialDelay;
 
-    while (retries < this.maxRetries) {
+    while (true) {
       const result = await getNextTaskId();
 
       if (result.type === 'success') {