From 5ffe51e3a03d4bc49e1716bf87421ce6105b65fb Mon Sep 17 00:00:00 2001
From: Emil Lerch <emil@lerch.org>
Date: Thu, 1 Feb 2024 10:52:44 -0800
Subject: [PATCH] add key schema stuff

---
 src/createtable.zig | 74 +++++++++++++++++++++++++++++++++++++++++++--
 src/ddb.zig         |  7 +++++
 2 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/src/createtable.zig b/src/createtable.zig
index 17e5c20..7176358 100644
--- a/src/createtable.zig
+++ b/src/createtable.zig
@@ -29,6 +29,27 @@ pub fn handler(request: *AuthenticatedRequest, writer: anytype) ![]const u8 {
     var parsed = try std.json.parseFromSlice(std.json.Value, allocator, request.event_data, .{});
     defer parsed.deinit();
     const request_params = try parseRequest(request, parsed, writer);
+    // Parsing does most validation for us, but we also need to make sure that
+    // the attributes specified in the key schema actually exist
+    var found_keys: u2 = if (request_params.table_info.range_key_attribute_name == null) 0b01 else 0b00;
+    for (request_params.table_info.attribute_definitions) |def| {
+        if (std.mem.eql(u8, def.name, request_params.table_info.hash_key_attribute_name)) {
+            found_keys |= 0b10;
+            continue;
+        }
+        if (request_params.table_info.range_key_attribute_name) |n| {
+            if (std.mem.eql(u8, def.name, n))
+                found_keys |= 0b01;
+        }
+    }
+    if (found_keys != 0b11)
+        try returnException(
+            request,
+            .bad_request,
+            error.ValidationException,
+            writer,
+            "Attribute names in KeySchema must also exist in AttributeDefinitions",
+        );
     defer {
         for (request_params.table_info.attribute_definitions) |d| {
             allocator.free(d.*.name);
@@ -133,6 +154,8 @@ fn parseRequest(
         .table_info = .{
             .attribute_definitions = undefined,
             .table_key = undefined,
+            .hash_key_attribute_name = undefined,
+            .range_key_attribute_name = null,
         },
     };
     // This is a new table, so we will generate a random key for table data
@@ -233,13 +256,60 @@ fn parseRequest(
                     writer,
                     "KeySchema must be an array",
                 );
-            if (val.array.items.len == 0)
+            if (val.array.items.len < 1 or val.array.items.len > 2)
                 try returnException(
                     request,
                     .bad_request,
                     error.ValidationException,
                     writer,
-                    "KeySchema array cannot be empty",
+                    "KeySchema array must consist of only 1 or 2 elements",
+                );
+            var found_hash = false;
+            for (val.array.items) |item| {
+                //     {
+                //       "AttributeName": "Artist",
+                //       "KeyType": "HASH"
+                //     },
+                if (item != .object)
+                    try returnException(
+                        request,
+                        .bad_request,
+                        error.ValidationException,
+                        writer,
+                        "KeySchemaElement must be an object",
+                    );
+                const attribute_name = item.object.get("AttributeName");
+                const key_type = item.object.get("KeyType");
+                if (attribute_name == null or key_type == null or attribute_name.? != .string or key_type.? != .string)
+                    try returnException(
+                        request,
+                        .bad_request,
+                        error.ValidationException,
+                        writer,
+                        "KeySchemaElement must contain AttributeName and KeyType strings",
+                    );
+                if (!std.mem.eql(u8, key_type.?.string, "HASH") and !std.mem.eql(u8, key_type.?.string, "RANGE"))
+                    try returnException(
+                        request,
+                        .bad_request,
+                        error.ValidationException,
+                        writer,
+                        "Invalid KeyType. Valid values are HASH | RANGE",
+                    );
+                const is_hash = std.mem.eql(u8, key_type.?.string, "HASH");
+                found_hash = found_hash or is_hash;
+                if (is_hash)
+                    request_params.table_info.hash_key_attribute_name = attribute_name.?.string
+                else
+                    request_params.table_info.range_key_attribute_name = attribute_name.?.string;
+            }
+            if (!found_hash)
+                try returnException(
+                    request,
+                    .bad_request,
+                    error.ValidationException,
+                    writer,
+                    "KeySchema missing hash key",
                 );
             continue;
         }
diff --git a/src/ddb.zig b/src/ddb.zig
index 4fe3148..cbdb9ff 100644
--- a/src/ddb.zig
+++ b/src/ddb.zig
@@ -53,6 +53,11 @@ pub const TableInfo = struct {
     // gsi_description_list: []const u8, // Not sure how this is used
     // sqlite_index: []const u8, // Not sure how this is used
     table_key: [encryption.encoded_key_length]u8,
+
+    // DDB Local is using sqlite_index here, which seems very much overkill
+    // as everything can be determined by just the name...
+    hash_key_attribute_name: []const u8,
+    range_key_attribute_name: ?[]const u8,
 };
 
 pub const TableArray = struct {
@@ -287,6 +292,8 @@ fn testCreateTable(allocator: std.mem.Allocator, account_id: []const u8) !sqlite
     var table_info: TableInfo = .{
         .table_key = undefined,
         .attribute_definitions = definitions[0..],
+        .hash_key_attribute_name = "Artist",
+        .range_key_attribute_name = null,
     };
     encryption.randomEncodedKey(&table_info.table_key);
     try createDdbTable(