pagination.ts

import { DynamoDBClient as V3, ListTablesCommand, DeleteTableCommand, CreateTableCommand } from '@aws-sdk/client-dynamodb';
import { Entity, QueryResponse, EntityItem } from 'electrodb';
import { v4 as uuid } from 'uuid';

const client = new V3({
  endpoint: 'http://localhost:8000',
  region: 'us-east-1',
});

const table = "pagination_test";

const dynamodb = client;

export function createTableManager(tableName: string) {
    return {
        async exists() {
          let tables = await dynamodb.send(new ListTablesCommand({}));
          return !!tables.TableNames?.includes(tableName);
        },
        async drop() {
            return dynamodb.send(new DeleteTableCommand({TableName: tableName}));
        },
        async create() {
            return dynamodb.send(new CreateTableCommand({
              "KeySchema":[
                  {
                      "AttributeName":"pk",
                      "KeyType":"HASH"
                  },
                  {
                      "AttributeName":"sk",
                      "KeyType":"RANGE"
                  }
              ],
              "AttributeDefinitions":[
                  {
                      "AttributeName":"pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"sk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi1pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi1sk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi2pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi2sk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi3pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi3sk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi4pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi4sk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi5pk",
                      "AttributeType":"S"
                  },
                  {
                      "AttributeName":"gsi5sk",
                      "AttributeType":"S"
                  }
              ],
              "GlobalSecondaryIndexes":[
                  {
                      "IndexName":"gsi1pk-gsi1sk-index",
                      "KeySchema":[
                          {
                              "AttributeName":"gsi1pk",
                              "KeyType":"HASH"
                          },
                          {
                              "AttributeName":"gsi1sk",
                              "KeyType":"RANGE"
                          }
                      ],
                      "Projection":{
                          "ProjectionType":"ALL"
                      },
                      "ProvisionedThroughput": {       // Only specified if using provisioned mode
                        "ReadCapacityUnits": 1,
                        "WriteCapacityUnits": 1
                      },
                  },
                  {
                      "IndexName":"gsi2pk-gsi2sk-index",
                      "KeySchema":[
                          {
                              "AttributeName":"gsi2pk",
                              "KeyType":"HASH"
                          },
                          {
                              "AttributeName":"gsi2sk",
                              "KeyType":"RANGE"
                          }
                      ],
                      "Projection":{
                          "ProjectionType":"ALL"
                      },
                      "ProvisionedThroughput": {       // Only specified if using provisioned mode
                        "ReadCapacityUnits": 1,
                        "WriteCapacityUnits": 1
                      },
                  },
                  {
                      "IndexName":"gsi3pk-gsi3sk-index",
                      "KeySchema":[
                          {
                              "AttributeName":"gsi3pk",
                              "KeyType":"HASH"
                          },
                          {
                              "AttributeName":"gsi3sk",
                              "KeyType":"RANGE"
                          }
                      ],
                      "Projection":{
                          "ProjectionType":"ALL"
                      },
                      "ProvisionedThroughput": {       // Only specified if using provisioned mode
                        "ReadCapacityUnits": 1,
                        "WriteCapacityUnits": 1
                      },
                  },
                  {
                      "IndexName":"gsi4pk-gsi4sk-index",
                      "KeySchema":[
                          {
                              "AttributeName":"gsi4pk",
                              "KeyType":"HASH"
                          },
                          {
                              "AttributeName":"gsi4sk",
                              "KeyType":"RANGE"
                          }
                      ],
                      "Projection":{
                          "ProjectionType":"ALL"
                      },
                      "ProvisionedThroughput": {       // Only specified if using provisioned mode
                        "ReadCapacityUnits": 1,
                        "WriteCapacityUnits": 1
                      },
                  },
                  {
                      "IndexName":"gsi5pk-gsi5sk-index",
                      "KeySchema":[
                          {
                              "AttributeName":"gsi5pk",
                              "KeyType":"HASH"
                          },
                          {
                              "AttributeName":"gsi5sk",
                              "KeyType":"RANGE"
                          }
                      ],
                      "Projection":{
                          "ProjectionType":"ALL"
                      },
                      "ProvisionedThroughput": {       // Only specified if using provisioned mode
                        "ReadCapacityUnits": 1,
                        "WriteCapacityUnits": 1
                      },
                  }
              ],
              "ProvisionedThroughput": {       // Only specified if using provisioned mode
                "ReadCapacityUnits": 1,
                "WriteCapacityUnits": 1
              },
              TableName: tableName,
              BillingMode: ''
            }));
        }
    }
}
async function initializeTable() {
  const tableManager = createTableManager(table);
  const exists = await tableManager.exists();
  if (!exists) {
    await tableManager.create();
  }
}

/* Users Entity */
const users = new Entity(
  {
    model: {
      entity: "user",
      service: "taskapp",
      version: "1"
    },
    attributes: {
      team: {
        type: "string"
      },
      user: {
        type: "string"
      },
      role: {
        type: ["dev", "senior", "staff", "principal"] as const,
        set: (title: string) => {
          // save as index for comparison
          return [
            "dev",
            "senior",
            "staff",
            "principal"
          ].indexOf(title);
        },
        get: (index: number) => {
          return [
            "dev",
            "senior",
            "staff",
            "principal"
          ][index] || "other";
        }
      },
      manager: {
        type: "set",
        items: ["frank", "jane", "joe", "sally"] as const,
      },
      firstName: {
        type: "string"
      },
      lastName: {
        type: "string"
      },
      fullName: {
        type: "string",
        // never set value to the database
        set: () => undefined,
        // calculate full name on retrieval
        get: (_, {firstName, lastName}) => {
          return `${firstName ?? ""} ${lastName ?? ""}`.trim();
        }
      },
      profile: {
        type: "map",
        properties: {
          photo: {
            type: "string"
          },
          bio: {
            type: "string"
          },
          location: {
            type: "string"
          }
        }
      },
      pinned: {
        type: "any"
      },
      following: {
        type: "set",
        items: "string"
      },
      followers: {
        type: "set",
        items: "string"
      },
      createdAt: {
        type: "number",
        default: () => Date.now(),
        readOnly: true
      },
      updatedAt: {
        type: "number",
        watch: "*",
        set: () => Date.now(),
        readOnly: true
      }
    },
    indexes: {
      members: {
        collection: "organization",
        pk: {
          composite: ["team"],
          field: "pk"
        },
        sk: {
          composite: ["user"],
          field: "sk"
        }
      },
      user: {
        collection: "assignments",
        index: "gsi1pk-gsi1sk-index",
        pk: {
          composite: ["user"],
          field: "gsi1pk"
        },
        sk: {
          field: "gsi1sk",
          composite: []
        }
      }
    }
  },
  { table, client }
);
/* Write queries to generate parameters on the right */

function generateUsers(count: number, team: string) {
  return new Array(count).fill(0).map(() => ({
    team,
    user: uuid(),
    role: 'staff' as const,
    lastName: uuid(),
    firstName: uuid(),
    // manager: ['frank', 'jane'] as const, // This fails if Entity is instantiated without a client
    profile: {
      bio: uuid(),
      photo: uuid(),
      location: uuid(),
    },
    // interact with DynamoDB sets like arrays
    following: [uuid(),]
  }));
}

type UserItem = EntityItem<typeof users>;
type UserQueryResponse = QueryResponse<typeof users>;

async function getTeamMembers(team: string) {
  let stores: UserItem[] = [];
  let cursor = null;
  let pages = 0;
  do {
    pages++;
    const results: UserQueryResponse = await users.query
      .members({ team })
      .go({ cursor });
    stores = [...stores, ...results.data];
    cursor = results.cursor;
  } while(cursor !== null);
  console.log('pages:', pages);
  return stores;
}

const run = async () => {
try {
  // await initializeTable();
  const team = uuid();
  const generated = generateUsers(50000, team);
  const lookup = new Map<string, UserItem>();
  for (const item of generated) {
    lookup.set(item.user, item);
  }
  console.log('generated count:', generated.length);

  await users.put(generated).go({concurrency: 100});
  const members = await getTeamMembers(team);
  console.log('member count:', members.length);
  const missing = members.filter(m => !lookup.has(m.user));
  console.log('missing count:', missing.length);
} catch(err) {
  console.log('err %o', err);
}
};

initializeTable().then(run).catch(console.log);