add encode entity key plugin (#970)

* add encode entity key plugin

* add bs58 types

* version fix

Author: bryzettler

packages/account-postgres-sink-service/package.json

@@ -74,6 +74,7 @@
   "devDependencies": {
     "@types/async-retry": "^1.4.8",
     "@types/bn.js": "^5.1.1",
+    "@types/bs58": "^4.0.1",
     "@types/cron": "^2.4.0",
     "@types/lodash": "^4.14.195",
     "@types/node": "^18.11.11",

packages/account-postgres-sink-service/src/plugins/encodeEntityKey.ts

@@ -0,0 +1,67 @@
+import { DataTypes } from "sequelize";
+import { IPlugin } from "../types";
+import bs58 from "bs58";
+import { camelize } from "inflection";
+
+export const EncodeEntityKeyPlugin = ((): IPlugin => {
+  const name = "EncodeEntityKey";
+  const init = async (config: { [key: string]: any }) => {
+    const updateOnDuplicateFields = ["encoded_entity_key"];
+    const addFields = (schema: { [key: string]: any }, accountName: string) => {
+      schema[accountName] = {
+        ...schema[accountName],
+        encoded_entity_key: DataTypes.TEXT,
+      };
+    };
+
+    const addIndexes = (
+      schema: { [key: string]: any },
+      accountName: string
+    ) => {
+      schema[accountName] = {
+        ...schema[accountName],
+        indexes: [
+          {
+            fields: ["encoded_entity_key"],
+            name: `idx_encoded_entity_key`,
+            unique: true,
+          },
+        ],
+      };
+    };
+
+    const processAccount = async (account: { [key: string]: any }) => {
+      const entityKey = account[camelize(config.field || "entity_key", true)];
+      const keySerializationRaw = account[camelize("key_serialization", true)];
+      const keySerialization =
+        typeof keySerializationRaw === "string"
+          ? keySerializationRaw.trim().toLowerCase()
+          : String(keySerializationRaw).trim().toLowerCase();
+      let encodedEntityKey: string | null = null;
+      if (entityKey && keySerialization) {
+        if (keySerialization === "utf8") {
+          encodedEntityKey = Buffer.from(entityKey, "utf8").toString("utf8");
+        } else if (keySerialization === "b58" || keySerialization === "bs58") {
+          encodedEntityKey = bs58.encode(entityKey);
+        }
+      }
+
+      return {
+        ...account,
+        encoded_entity_key: encodedEntityKey,
+      };
+    };
+
+    return {
+      updateOnDuplicateFields,
+      addFields,
+      addIndexes,
+      processAccount,
+    };
+  };
+
+  return {
+    name,
+    init,
+  };
+})();

packages/account-postgres-sink-service/src/plugins/extractHexLocation.ts

@@ -94,7 +94,7 @@ export const ExtractHexLocationPlugin = ((): IPlugin => {
     const mapbox = MapboxService.getInstance();
     const processAccount = async (account: { [key: string]: any }) => {
       let reverseGeod: ReverseGeoCache | null = null;
-      const location = account[config.field || "location"];
+      const location = account[camelize(config.field || "location", true)];
       if (location) {
         reverseGeod = await ReverseGeoCache.findByPk(location.toString(), {
           attributes: updateOnDuplicateFields,

packages/account-postgres-sink-service/src/plugins/index.ts

@@ -1,8 +1,9 @@
 import { IConfig, IInitedPlugin, IPluginConfig } from "../types";
 import { truthy } from "../utils/truthy";
 import { ExtractHexLocationPlugin } from "./extractHexLocation";
+import { EncodeEntityKeyPlugin } from "./encodeEntityKey";
 
-export const Plugins = [ExtractHexLocationPlugin];
+export const Plugins = [ExtractHexLocationPlugin, EncodeEntityKeyPlugin];
 
 export const initPlugins = async (pluginConfigs: IPluginConfig[] = []) =>
   (

packages/account-postgres-sink-service/src/types.ts

@@ -1,6 +1,6 @@
 import { Transaction } from "sequelize";
 
-type Plugins = "ExtractHexLocation";
+type Plugins = "ExtractHexLocation" | "EncodeEntityKey";
 type Crons = "refresh-accounts" | "integrity-check";
 
 export interface IPluginConfig {
@@ -34,6 +34,7 @@ export interface IConfig {
 export interface IInitedPlugin {
   updateOnDuplicateFields?: string[];
   addFields?: (schema: { [key: string]: any }, accountName: string) => void;
+  addIndexes?: (schema: { [key: string]: any }, accountName: string) => void;
   processAccount: (account: any, t?: Transaction) => Promise<any>;
 }
 

packages/account-postgres-sink-service/src/utils/defineIdlModels.ts

@@ -10,6 +10,7 @@ import {
   IdlTypeDef,
   IdlTypeDefTyStruct,
 } from "@coral-xyz/anchor/dist/cjs/idl";
+import { omit, pick } from "lodash";
 
 const TypeMap = new Map<string, any>([
   ["string", DataTypes.STRING],
@@ -57,6 +58,11 @@ const determineType = (type: string | object): any => {
   return DataTypes.JSONB;
 };
 
+const shouldEnableCreatedAt = (schema: any, accName: string) => {
+  const accSchema = schema[accName];
+  return !accSchema || (!accSchema.createdAt && !accSchema.created_at);
+};
+
 export const defineIdlModels = async ({
   idl,
   accounts,
@@ -84,9 +90,10 @@ export const defineIdlModels = async ({
         }
       }
 
-      (await initPlugins(accConfig?.plugins)).map(
-        (plugin) => plugin?.addFields && plugin.addFields(schema, acc.name)
-      );
+      (await initPlugins(accConfig?.plugins)).map((plugin) => {
+        if (plugin?.addFields) plugin.addFields(schema, acc.name);
+        if (plugin?.addIndexes) plugin.addIndexes(schema, acc.name);
+      });
 
       if (accConfig.schema) {
         await sequelize.createSchema(accConfig.schema, {});
@@ -99,7 +106,7 @@ export const defineIdlModels = async ({
             type: DataTypes.STRING,
             primaryKey: true,
           },
-          ...schema[acc.name],
+          ...omit(schema[acc.name] || {}, ["indexes"]),
           refreshed_at: {
             type: DataTypes.DATE,
           },
@@ -109,9 +116,8 @@ export const defineIdlModels = async ({
           updatedAt: false,
           schema: underscore(accConfig.schema || "public"),
           tableName: underscore(accConfig.table || acc.name),
-          createdAt:
-            !schema[acc.name] ||
-            (!schema[acc.name].createdAt && !schema[acc.name].created_at),
+          createdAt: shouldEnableCreatedAt(schema, acc.name),
+          ...pick(schema[acc.name], "indexes"),
         }
       );
 

packages/account-postgres-sink-service/yarn.deploy.lock

@@ -222,6 +222,7 @@ __metadata:
     "@triton-one/yellowstone-grpc": ^0.4.0
     "@types/async-retry": ^1.4.8
     "@types/bn.js": ^5.1.1
+    "@types/bs58": ^4.0.1
     "@types/cron": ^2.4.0
     "@types/deep-equal": ^1.0.4
     "@types/lodash": ^4.14.195
@@ -715,6 +716,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/bs58@npm:^4.0.1":
+  version: 4.0.1
+  resolution: "@types/bs58@npm:4.0.1"
+  dependencies:
+    base-x: "npm:^3.0.6"
+  checksum: 5063fed6bb3816f28eef3e689a05f8f85d03ce7c2e3ceefe77c349ed147023999eb024e863bf3fa81ba2bded690e5296efe4d34864c875097bd261614e37f786
+  languageName: node
+  linkType: hard
+
 "@types/connect@npm:^3.4.33":
   version: 3.4.35
   resolution: "@types/connect@npm:3.4.35"
@@ -1176,7 +1186,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"base-x@npm:^3.0.2":
+"base-x@npm:^3.0.2, base-x@npm:^3.0.6":
   version: 3.0.9
   resolution: "base-x@npm:3.0.9"
   dependencies:

yarn.lock

@@ -749,6 +749,7 @@ __metadata:
     "@triton-one/yellowstone-grpc": ^0.4.0
     "@types/async-retry": ^1.4.8
     "@types/bn.js": ^5.1.1
+    "@types/bs58": ^4.0.1
     "@types/cron": ^2.4.0
     "@types/deep-equal": ^1.0.4
     "@types/lodash": ^4.14.195