Allow to decorate an attribute with multiple indexes

docs/docs/guide/query.mdx

@@ -152,6 +152,14 @@ This comparison function will check if the given key is between the values that
 UserManager.query().sortKey('sortKey').between('bla', 'cla'); // Resulting in: `sortKey BETWEEN 'bla' AND 'cla'`
 ```
 
+## Query.indexName(name)
+
+The name of a secondary index to query. This index can be any local secondary index or global secondary index. The `name` parameter is a string, narrowed down to entity index names. It uses DynamoDB's `IndexName`.
+
+```ts  
+UserManager.query().partitionKey('partitionKey').eq('1').indexName('LSI_1_NAME') // Resulting in: `IndexName: LSI_1_NAME`
+```
+
 ## Query.sort(order)
 
 This method sorts the items you receive by using sort key. It uses DynamoDB's `ScanIndexForward`.

lib/decorators/helpers/decorateAttribute.ts

@@ -1,26 +1,40 @@
 import { IndexDecoratorOptions, PrefixSuffixOptions } from '@lib/decorators/types';
 import Dynamode from '@lib/dynamode/index';
-import { AttributeMetadata, AttributeRole, AttributeType } from '@lib/dynamode/storage/types';
+import { AttributeIndexRole, AttributeRole, AttributeType } from '@lib/dynamode/storage/types';
 
 export function decorateAttribute(
   type: AttributeType,
-  role: AttributeRole,
+  role: Exclude<AttributeRole, 'index'> | AttributeIndexRole,
   options?: PrefixSuffixOptions | IndexDecoratorOptions,
 ): (Entity: any, propertyName: string) => void {
   return (Entity: any, propertyName: string) => {
     const entityName = Entity.constructor.name;
     const prefix = options && 'prefix' in options ? options.prefix : undefined;
     const suffix = options && 'suffix' in options ? options.suffix : undefined;
-    const indexName = options && 'indexName' in options ? options.indexName : undefined;
-    const attributeMetadata: AttributeMetadata = {
+
+    if (role === 'gsiPartitionKey' || role === 'gsiSortKey' || role === 'lsiSortKey') {
+      const indexName = options && 'indexName' in options ? options.indexName : undefined;
+
+      if (!indexName) {
+        throw new Error(`Index name is required for ${role} attribute`);
+      }
+
+      return Dynamode.storage.registerAttribute(entityName, propertyName, {
+        propertyName,
+        type,
+        role: 'index',
+        indexes: [{ name: indexName, role }],
+        prefix,
+        suffix,
+      });
+    }
+
+    Dynamode.storage.registerAttribute(entityName, propertyName, {
       propertyName,
       type,
       role,
-      indexName,
       prefix,
       suffix,
-    };
-
-    Dynamode.storage.registerAttribute(entityName, propertyName, attributeMetadata);
+    });
   };
 }

lib/dynamode/storage/helpers/validator.ts

@@ -20,8 +20,20 @@ export function validateMetadataAttribute({
     throw new ValidationError(`Attribute "${name}" is decorated with a wrong role in "${entityName}" Entity.`);
   }
 
-  if (attribute.indexName !== indexName) {
-    throw new ValidationError(`Attribute "${name}" is decorated with a wrong index in "${entityName}" Entity.`);
+  if (!indexName && attribute.role === 'index') {
+    throw new ValidationError(`Index for attribute "${name}" should be added to "${entityName}" Entity metadata.`);
+  }
+
+  if (indexName && attribute.role !== 'index') {
+    throw new ValidationError(
+      `Attribute "${name}" should be decorated with index "${indexName}" in "${entityName}" Entity.`,
+    );
+  }
+
+  if (indexName && attribute.role === 'index' && !attribute.indexes.some((index) => index.name === indexName)) {
+    throw new ValidationError(
+      `Attribute "${name}" is not decorated with index "${indexName}" in "${entityName}" Entity.`,
+    );
   }
 
   if (!DYNAMODE_ALLOWED_KEY_TYPES.includes(attribute.type)) {
@@ -38,30 +50,45 @@ export function validateDecoratedAttribute({
   const roleValidationMap: Record<AttributeRole, (v: ValidateDecoratedAttribute) => boolean> = {
     partitionKey: ({ name, metadata }) => metadata.partitionKey !== name,
     sortKey: ({ name, metadata }) => metadata.sortKey !== name,
-    gsiPartitionKey: ({ attribute, name, metadata }) =>
-      !!attribute.indexName && metadata.indexes?.[attribute.indexName]?.partitionKey !== name,
-    gsiSortKey: ({ attribute, name, metadata }) =>
-      !!attribute.indexName && metadata.indexes?.[attribute.indexName]?.sortKey !== name,
-    lsiSortKey: ({ attribute, name, metadata }) =>
-      !!attribute.indexName && metadata.indexes?.[attribute.indexName]?.sortKey !== name,
+    index: ({ attribute, name, metadata }) => {
+      if (!('indexes' in attribute) || !attribute.indexes.length) {
+        return true;
+      }
+
+      return attribute.indexes.some((index) => {
+        switch (index.role) {
+          case 'gsiPartitionKey':
+            return metadata.indexes?.[index.name]?.partitionKey !== name;
+          case 'gsiSortKey':
+          case 'lsiSortKey':
+            return metadata.indexes?.[index.name]?.sortKey !== name;
+          default:
+            return true;
+        }
+      });
+    },
     date: () => false,
     attribute: () => false,
     dynamodeEntity: () => false,
   };
 
   const validateAttributeRole = roleValidationMap[attribute.role];
   if (validateAttributeRole({ attribute, name, metadata, entityName })) {
-    throw new ValidationError(`Attribute "${name}" is decorated with a wrong role in "${entityName}" Entity.`);
+    throw new ValidationError(
+      `Attribute "${name}" is decorated with a wrong role in "${entityName}" Entity. This could mean two things:\n1. The attribute is not defined in the table metadata.\n2. The attribute is defined in the table metadata but wrong decorator was used.\n`,
+    );
   }
 }
 
 export function validateMetadataUniqueness(entityName: string, metadata: Metadata<typeof Entity>): void {
+  const allIndexes = Object.values(metadata.indexes ?? {}).flatMap((index) => [index.partitionKey, index.sortKey]);
+
   const metadataKeys = [
     metadata.partitionKey,
     metadata.sortKey,
     metadata.createdAt,
     metadata.updatedAt,
-    ...Object.values(metadata.indexes ?? {}).flatMap((index) => [index.partitionKey, index.sortKey]),
+    ...new Set(allIndexes),
   ].filter((attribute) => !!attribute);
 
   if (metadataKeys.length !== new Set(metadataKeys).size) {

lib/dynamode/storage/index.ts

@@ -56,10 +56,16 @@ export default class DynamodeStorage {
       this.entities[entityName] = { attributes: {} } as EntityMetadata;
     }
 
-    if (this.entities[entityName].attributes[propertyName]) {
+    const attributeMetadata = this.entities[entityName].attributes[propertyName];
+
+    if (attributeMetadata && (attributeMetadata.role !== 'index' || value.role !== 'index')) {
       throw new DynamodeStorageError(`Attribute "${propertyName}" was already decorated in entity "${entityName}"`);
     }
 
+    if (attributeMetadata && attributeMetadata.role === 'index' && value.role === 'index') {
+      return void attributeMetadata.indexes.push(...value.indexes);
+    }
+
     this.entities[entityName].attributes[propertyName] = value;
   }
 
@@ -169,19 +175,19 @@ export default class DynamodeStorage {
         validateMetadataAttribute({
           name: index.partitionKey,
           attributes,
-          role: 'gsiPartitionKey',
+          role: 'index',
           indexName,
           entityName,
         });
 
         if (index.sortKey) {
-          validateMetadataAttribute({ name: index.sortKey, attributes, role: 'gsiSortKey', indexName, entityName });
+          validateMetadataAttribute({ name: index.sortKey, attributes, role: 'index', indexName, entityName });
         }
       }
 
       // Validate LSI
       if (index.sortKey && !index.partitionKey) {
-        validateMetadataAttribute({ name: index.sortKey, attributes, role: 'lsiSortKey', indexName, entityName });
+        validateMetadataAttribute({ name: index.sortKey, attributes, role: 'index', indexName, entityName });
       }
     });
   }

lib/dynamode/storage/types.ts

@@ -14,25 +14,30 @@ export type AttributeType =
   | MapConstructor
   | Uint8ArrayConstructor;
 
-export type AttributeRole =
-  | 'partitionKey'
-  | 'sortKey'
-  | 'gsiPartitionKey'
-  | 'gsiSortKey'
-  | 'lsiSortKey'
-  | 'date'
-  | 'attribute'
-  | 'dynamodeEntity';
-
-export type AttributeMetadata = {
+export type AttributeRole = 'partitionKey' | 'sortKey' | 'index' | 'date' | 'attribute' | 'dynamodeEntity';
+export type AttributeIndexRole = 'gsiPartitionKey' | 'gsiSortKey' | 'lsiSortKey';
+
+type BaseAttributeMetadata = {
   propertyName: string;
   type: AttributeType;
-  role: AttributeRole;
-  indexName?: string;
   prefix?: string;
   suffix?: string;
 };
 
+export type NonIndexAttributeMetadata = BaseAttributeMetadata & {
+  role: Exclude<AttributeRole, 'index'>;
+  indexName?: never;
+};
+
+export type IndexMetadata = { name: string; role: AttributeIndexRole };
+
+export type IndexAttributeMetadata = BaseAttributeMetadata & {
+  role: 'index';
+  indexes: IndexMetadata[];
+};
+
+export type AttributeMetadata = NonIndexAttributeMetadata | IndexAttributeMetadata;
+
 export type AttributesMetadata = {
   [attributeName: string]: AttributeMetadata;
 };

lib/query/index.ts

@@ -1,16 +1,27 @@
 import { QueryCommandOutput, QueryInput } from '@aws-sdk/client-dynamodb';
 import Dynamode from '@lib/dynamode/index';
+import { AttributeMetadata, IndexMetadata } from '@lib/dynamode/storage/types';
 import Entity from '@lib/entity';
 import { convertAttributeValuesToEntity, convertAttributeValuesToLastKey } from '@lib/entity/helpers/converters';
 import { EntityKey, EntityValue } from '@lib/entity/types';
 import type { QueryRunOptions, QueryRunOutput } from '@lib/query/types';
 import RetrieverBase from '@lib/retriever';
 import { Metadata, TablePartitionKeys, TableSortKeys } from '@lib/table/types';
-import { AttributeValues, BASE_OPERATOR, ExpressionBuilder, isNotEmptyString, Operators, timeout } from '@lib/utils';
+import {
+  AttributeValues,
+  BASE_OPERATOR,
+  ExpressionBuilder,
+  isNotEmptyString,
+  Operators,
+  timeout,
+  ValidationError,
+} from '@lib/utils';
 
 export default class Query<M extends Metadata<E>, E extends typeof Entity> extends RetrieverBase<M, E> {
   protected declare input: QueryInput;
   protected keyOperators: Operators = [];
+  protected partitionKeyMetadata?: AttributeMetadata;
+  protected sortKeyMetadata?: AttributeMetadata;
 
   constructor(entity: E) {
     super(entity);
@@ -20,6 +31,7 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
   public run(options: QueryRunOptions & { return: 'output' }): Promise<QueryCommandOutput>;
   public run(options: QueryRunOptions & { return: 'input' }): QueryInput;
   public run(options?: QueryRunOptions): Promise<QueryRunOutput<M, E> | QueryCommandOutput> | QueryInput {
+    this.setAssociatedIndexName();
     this.buildQueryInput(options?.extraInput);
 
     if (options?.return === 'input') {
@@ -66,7 +78,8 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
 
   public partitionKey<Q extends Query<M, E>, K extends EntityKey<E> & TablePartitionKeys<M, E>>(this: Q, key: K) {
     this.maybePushKeyLogicalOperator();
-    this.setAssociatedIndexName(key);
+    const attributes = Dynamode.storage.getEntityAttributes(this.entity.name);
+    this.partitionKeyMetadata = attributes[key as string];
 
     return {
       eq: (value: EntityValue<E, K>): Q => this.eq(this.keyOperators, key, value),
@@ -75,7 +88,8 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
 
   public sortKey<Q extends Query<M, E>, K extends EntityKey<E> & TableSortKeys<M, E>>(this: Q, key: K) {
     this.maybePushKeyLogicalOperator();
-    this.setAssociatedIndexName(key);
+    const attributes = Dynamode.storage.getEntityAttributes(this.entity.name);
+    this.sortKeyMetadata = attributes[key as string];
 
     return {
       eq: (value: EntityValue<E, K>): Q => this.eq(this.keyOperators, key, value),
@@ -101,12 +115,71 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
     }
   }
 
-  private setAssociatedIndexName<K extends EntityKey<E> & TablePartitionKeys<M, E>>(key: K) {
-    const attributes = Dynamode.storage.getEntityAttributes(this.entity.name);
-    const { indexName } = attributes[key as string];
+  private setAssociatedIndexName() {
+    if (this.input.IndexName) {
+      return;
+    }
+
+    const { partitionKeyMetadata, sortKeyMetadata } = this;
+    if (!partitionKeyMetadata) {
+      throw new ValidationError('You need to use ".partitionKey()" method before calling ".run()"');
+    }
+
+    // Primary key
+    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.role !== 'index') {
+      return;
+    }
+
+    // GSI with sort key
+    if (partitionKeyMetadata.role === 'index' && sortKeyMetadata?.role === 'index') {
+      const pkIndexes: IndexMetadata[] = partitionKeyMetadata.indexes;
+      const skIndexes: IndexMetadata[] = sortKeyMetadata.indexes;
+
+      const commonIndexes = pkIndexes.filter((pkIndex) => skIndexes.some((skIndex) => skIndex.name === pkIndex.name));
+      if (commonIndexes.length === 0) {
+        throw new ValidationError(
+          `No common indexes found for "${partitionKeyMetadata.propertyName}" and "${sortKeyMetadata.propertyName}"`,
+        );
+      }
+
+      if (commonIndexes.length > 1) {
+        throw new ValidationError(
+          `Multiple common indexes found for "${partitionKeyMetadata.propertyName}" and "${sortKeyMetadata.propertyName}"`,
+        );
+      }
+
+      this.input.IndexName = commonIndexes[0].name;
+      return;
+    }
+
+    // GSI without sort key
+    if (partitionKeyMetadata.role === 'index' && !sortKeyMetadata) {
+      const possibleIndexes = partitionKeyMetadata.indexes;
+
+      if (possibleIndexes.length > 1) {
+        throw new ValidationError(
+          `Multiple indexes found for "${partitionKeyMetadata.propertyName}", please use ".indexName(${possibleIndexes
+            .map((index) => index.name)
+            .join(' | ')})" method to specify the index name`,
+        );
+      }
+
+      this.input.IndexName = possibleIndexes[0].name;
+      return;
+    }
+
+    // LSI with sort key
+    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.role === 'index') {
+      const possibleIndexes = sortKeyMetadata.indexes;
+
+      if (possibleIndexes.length > 1) {
+        throw new ValidationError(
+          `Multiple indexes found for "${sortKeyMetadata.propertyName}", an LSI can only have one index`,
+        );
+      }
 
-    if (indexName) {
-      this.input.IndexName = indexName;
+      this.input.IndexName = possibleIndexes[0].name;
+      return;
     }
   }
 

lib/retriever/index.ts

@@ -5,7 +5,7 @@ import Entity from '@lib/entity';
 import { buildGetProjectionExpression } from '@lib/entity/helpers/buildExpressions';
 import { convertRetrieverLastKeyToAttributeValues } from '@lib/entity/helpers/converters';
 import { EntityKey } from '@lib/entity/types';
-import { Metadata, TableRetrieverLastKey } from '@lib/table/types';
+import { Metadata, TableIndexNames, TableRetrieverLastKey } from '@lib/table/types';
 import { AttributeNames, AttributeValues } from '@lib/utils';
 
 export default class RetrieverBase<M extends Metadata<E>, E extends typeof Entity> extends Condition<E> {
@@ -20,6 +20,11 @@ export default class RetrieverBase<M extends Metadata<E>, E extends typeof Entit
     };
   }
 
+  public indexName(name: TableIndexNames<M, E>) {
+    this.input.IndexName = String(name);
+    return this;
+  }
+
   public limit(count: number) {
     this.input.Limit = count;
     return this;