Fix: Cannot add multiple attribute decorators for primary key

examples/playground.ts

@@ -0,0 +1,51 @@
+import attribute from '../dist/decorators';
+import Entity from '../dist/entity';
+import TableManager from '../dist/table';
+
+export class TableUser extends Entity {
+  @attribute.gsi.partitionKey.string({ indexName: 'index-1' })
+  @attribute.partitionKey.string()
+  domain: string;
+
+  @attribute.sortKey.string()
+  @attribute.gsi.sortKey.string({ indexName: 'index-1' })
+  email: string;
+
+  constructor(data: { domain: string; email: string }) {
+    super(data);
+    this.domain = data.domain;
+    this.email = data.email;
+  }
+}
+
+export const UserdataTableManager = new TableManager(TableUser, {
+  tableName: 'USERS_TABLE_NAME',
+  partitionKey: 'domain',
+  sortKey: 'email',
+  indexes: {
+    'index-1': {
+      partitionKey: 'domain',
+      sortKey: 'email',
+    },
+  },
+});
+
+const EntityManager = UserdataTableManager.entityManager();
+
+async function test() {
+  const entity = await EntityManager.put(
+    new TableUser({
+      domain: 'test',
+      email: 'not_empty',
+    }),
+  );
+  console.log(entity);
+}
+
+async function createTable() {
+  const table = await UserdataTableManager.createTable();
+  console.log(table);
+}
+
+// createTable();
+test();

lib/decorators/helpers/decorateAttribute.ts

@@ -19,7 +19,7 @@ export function decorateAttribute(
         throw new Error(`Index name is required for ${role} attribute`);
       }
 
-      return Dynamode.storage.registerAttribute(entityName, propertyName, {
+      return Dynamode.storage.registerIndex(entityName, propertyName, {
         propertyName,
         type,
         role: 'index',

lib/dynamode/storage/helpers/validator.ts

@@ -1,13 +1,10 @@
-/* eslint-disable @typescript-eslint/no-unused-vars */
-import type { AttributeRole, ValidateDecoratedAttribute, ValidateMetadataAttribute } from '@lib/dynamode/storage/types';
-import Entity from '@lib/entity';
-import { Metadata } from '@lib/table/types';
+import type { ValidateDecoratedAttribute, ValidateMetadataAttribute } from '@lib/dynamode/storage/types';
 import { DYNAMODE_ALLOWED_KEY_TYPES, ValidationError } from '@lib/utils';
 
 export function validateMetadataAttribute({
   attributes,
   name,
-  role,
+  validRoles,
   indexName,
   entityName,
 }: ValidateMetadataAttribute): void {
@@ -16,21 +13,17 @@ export function validateMetadataAttribute({
     throw new ValidationError(`Attribute "${name}" should be decorated in "${entityName}" Entity.`);
   }
 
-  if (attribute.role !== role) {
+  if (!validRoles.includes(attribute.role)) {
     throw new ValidationError(`Attribute "${name}" is decorated with a wrong role 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') {
+  if (indexName && !attribute.indexes) {
     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)) {
+  if (indexName && !attribute.indexes?.some((index) => index.name === indexName)) {
     throw new ValidationError(
       `Attribute "${name}" is not decorated with index "${indexName}" in "${entityName}" Entity.`,
     );
@@ -47,51 +40,37 @@ export function validateDecoratedAttribute({
   metadata,
   entityName,
 }: ValidateDecoratedAttribute): void {
-  const roleValidationMap: Record<AttributeRole, (v: ValidateDecoratedAttribute) => boolean> = {
-    partitionKey: ({ name, metadata }) => metadata.partitionKey !== name,
-    sortKey: ({ name, metadata }) => metadata.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 })) {
+  if (attribute.role === 'partitionKey' && metadata.partitionKey !== name) {
     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`,
+      `Attribute "${name}" is not defined as a partition key in "${entityName}" Entity's metadata.`,
     );
   }
-}
 
-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,
-    ...new Set(allIndexes),
-  ].filter((attribute) => !!attribute);
+  if (attribute.role === 'sortKey' && metadata.sortKey !== name) {
+    throw new ValidationError(`Attribute "${name}" is not defined as a sort key in "${entityName}" Entity's metadata.`);
+  }
 
-  if (metadataKeys.length !== new Set(metadataKeys).size) {
-    throw new ValidationError(`Duplicated metadata keys passed to "${entityName}" TableManager.`);
+  if (!attribute.indexes) {
+    return;
   }
+
+  attribute.indexes.forEach((index) => {
+    if (index.role === 'gsiPartitionKey' && metadata.indexes?.[index.name]?.partitionKey !== name) {
+      throw new ValidationError(
+        `Attribute "${name}" is not defined as a GSI partition key in "${entityName}" Entity's metadata for index named "${index.name}".`,
+      );
+    }
+
+    if (index.role === 'gsiSortKey' && metadata.indexes?.[index.name]?.sortKey !== name) {
+      throw new ValidationError(
+        `Attribute "${name}" is not defined as a GSI sort key in "${entityName}" Entity's metadata for index named "${index.name}".`,
+      );
+    }
+
+    if (index.role === 'lsiSortKey' && metadata.indexes?.[index.name]?.sortKey !== name) {
+      throw new ValidationError(
+        `Attribute "${name}" is not defined as a LSI sort key in "${entityName}" Entity's metadata for index named "${index.name}".`,
+      );
+    }
+  });
 }

lib/dynamode/storage/index.ts

@@ -1,9 +1,8 @@
-import {
-  validateDecoratedAttribute,
-  validateMetadataAttribute,
-  validateMetadataUniqueness,
-} from '@lib/dynamode/storage/helpers/validator';
+import { Except } from 'type-fest';
+
+import { validateDecoratedAttribute, validateMetadataAttribute } from '@lib/dynamode/storage/helpers/validator';
 import type {
+  AttributeIndexMetadata,
   AttributeMetadata,
   AttributesMetadata,
   EntitiesMetadata,
@@ -51,21 +50,43 @@ export default class DynamodeStorage {
     this.entities[entity.name] = entityMetadata;
   }
 
-  public registerAttribute(entityName: string, propertyName: string, value: AttributeMetadata): void {
+  public registerAttribute(
+    entityName: string,
+    propertyName: string,
+    value: Except<AttributeMetadata, 'indexes'>,
+  ): void {
     if (!this.entities[entityName]) {
       this.entities[entityName] = { attributes: {} } as EntityMetadata;
     }
 
     const attributeMetadata = this.entities[entityName].attributes[propertyName];
 
-    if (attributeMetadata && (attributeMetadata.role !== 'index' || value.role !== 'index')) {
+    // Throw error if attribute was already decorated and it's not an index
+    if (attributeMetadata && attributeMetadata.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);
+    // Otherwise, register attribute and copy over indexes
+    this.entities[entityName].attributes[propertyName] = {
+      ...value,
+      indexes: attributeMetadata?.indexes,
+    };
+  }
+
+  public registerIndex(entityName: string, propertyName: string, value: AttributeIndexMetadata): void {
+    if (!this.entities[entityName]) {
+      this.entities[entityName] = { attributes: {} } as EntityMetadata;
+    }
+
+    const attributeMetadata = this.entities[entityName].attributes[propertyName];
+
+    // Merge indexes if attribute was already decorated
+    if (attributeMetadata) {
+      attributeMetadata.indexes = [...(attributeMetadata.indexes ?? []), ...value.indexes];
+      return;
     }
 
+    // Register attribute with indexes
     this.entities[entityName].attributes[propertyName] = value;
   }
 
@@ -138,28 +159,25 @@ export default class DynamodeStorage {
     const metadata = this.getEntityMetadata(entityName);
     const attributes = this.getEntityAttributes(entityName);
 
-    // Validate metadata
-    validateMetadataUniqueness(entityName, metadata);
-
     // Validate decorated attributes
     Object.entries(attributes).forEach(([name, attribute]) =>
       validateDecoratedAttribute({ name, attribute, metadata, entityName }),
     );
 
     // Validate table partition key
-    validateMetadataAttribute({ name: metadata.partitionKey, attributes, role: 'partitionKey', entityName });
+    validateMetadataAttribute({ name: metadata.partitionKey, attributes, validRoles: ['partitionKey'], entityName });
 
     // Validate table sort key
     if (metadata.sortKey) {
-      validateMetadataAttribute({ name: metadata.sortKey, attributes, role: 'sortKey', entityName });
+      validateMetadataAttribute({ name: metadata.sortKey, attributes, validRoles: ['sortKey'], entityName });
     }
 
     // Validate table timestamps
     if (metadata.createdAt) {
-      validateMetadataAttribute({ name: metadata.createdAt, attributes, role: 'date', entityName });
+      validateMetadataAttribute({ name: metadata.createdAt, attributes, validRoles: ['date'], entityName });
     }
     if (metadata.updatedAt) {
-      validateMetadataAttribute({ name: metadata.updatedAt, attributes, role: 'date', entityName });
+      validateMetadataAttribute({ name: metadata.updatedAt, attributes, validRoles: ['date'], entityName });
     }
 
     // Validate table indexes
@@ -175,19 +193,31 @@ export default class DynamodeStorage {
         validateMetadataAttribute({
           name: index.partitionKey,
           attributes,
-          role: 'index',
           indexName,
           entityName,
+          validRoles: ['partitionKey', 'sortKey', 'index'],
         });
 
         if (index.sortKey) {
-          validateMetadataAttribute({ name: index.sortKey, attributes, role: 'index', indexName, entityName });
+          validateMetadataAttribute({
+            name: index.sortKey,
+            attributes,
+            indexName,
+            entityName,
+            validRoles: ['partitionKey', 'sortKey', 'index'],
+          });
         }
       }
 
       // Validate LSI
       if (index.sortKey && !index.partitionKey) {
-        validateMetadataAttribute({ name: index.sortKey, attributes, role: 'index', indexName, entityName });
+        validateMetadataAttribute({
+          name: index.sortKey,
+          attributes,
+          indexName,
+          entityName,
+          validRoles: ['sortKey', 'index'],
+        });
       }
     });
   }

lib/dynamode/storage/types.ts

@@ -24,20 +24,18 @@ type BaseAttributeMetadata = {
   suffix?: string;
 };
 
-export type NonIndexAttributeMetadata = BaseAttributeMetadata & {
-  role: Exclude<AttributeRole, 'index'>;
-  indexName?: never;
-};
-
 export type IndexMetadata = { name: string; role: AttributeIndexRole };
 
-export type IndexAttributeMetadata = BaseAttributeMetadata & {
+export type AttributeMetadata = BaseAttributeMetadata & {
+  role: AttributeRole;
+  indexes?: IndexMetadata[];
+};
+
+export type AttributeIndexMetadata = BaseAttributeMetadata & {
   role: 'index';
   indexes: IndexMetadata[];
 };
 
-export type AttributeMetadata = NonIndexAttributeMetadata | IndexAttributeMetadata;
-
 export type AttributesMetadata = {
   [attributeName: string]: AttributeMetadata;
 };
@@ -67,7 +65,7 @@ export type ValidateMetadataAttribute = {
   entityName: string;
   name: string;
   attributes: AttributesMetadata;
-  role: AttributeRole;
+  validRoles: AttributeRole[];
   indexName?: string;
 };
 

lib/query/index.ts

@@ -125,13 +125,18 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
       throw new ValidationError('You need to use ".partitionKey()" method before calling ".run()"');
     }
 
-    // Primary key
-    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.role !== 'index') {
+    // Primary key without sort key
+    if (partitionKeyMetadata.role === 'partitionKey' && !sortKeyMetadata) {
+      return;
+    }
+
+    // Primary key with sort key
+    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.role === 'sortKey') {
       return;
     }
 
     // GSI with sort key
-    if (partitionKeyMetadata.role === 'index' && sortKeyMetadata?.role === 'index') {
+    if (partitionKeyMetadata.indexes && sortKeyMetadata?.indexes) {
       const pkIndexes: IndexMetadata[] = partitionKeyMetadata.indexes;
       const skIndexes: IndexMetadata[] = sortKeyMetadata.indexes;
 
@@ -153,7 +158,7 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
     }
 
     // GSI without sort key
-    if (partitionKeyMetadata.role === 'index' && !sortKeyMetadata) {
+    if (partitionKeyMetadata.indexes && !sortKeyMetadata) {
       const possibleIndexes = partitionKeyMetadata.indexes;
 
       if (possibleIndexes.length > 1) {
@@ -169,7 +174,7 @@ export default class Query<M extends Metadata<E>, E extends typeof Entity> exten
     }
 
     // LSI with sort key
-    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.role === 'index') {
+    if (partitionKeyMetadata.role === 'partitionKey' && sortKeyMetadata?.indexes) {
       const possibleIndexes = sortKeyMetadata.indexes;
 
       if (possibleIndexes.length > 1) {

tests/e2e/condition.test.ts

@@ -2,7 +2,7 @@ import { afterAll, beforeAll, describe, expect, test, vi } from 'vitest';
 
 import { AttributeType } from '@lib/condition';
 
-import { mockDate, MockEntityManager, TEST_TABLE_NAME, TestTableManager } from '../fixtures';
+import { mockDate, MockEntityManager, TEST_TABLE_NAME, TestTableManager } from '../fixtures/TestTable';
 
 import { mockEntityFactory } from './mockEntityFactory';