From 507cd8c46d2964b6a6f899c8d42803e4812eafb1 Mon Sep 17 00:00:00 2001
From: Timur Moziev <code@timurrin.ru>
Date: Tue, 19 Mar 2024 01:07:22 +0300
Subject: [PATCH] add foreign key support

---
 migrate.js      | 15 +++++++++++++++
 migrate.test.js | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/migrate.js b/migrate.js
index 118fba6..1058c3f 100644
--- a/migrate.js
+++ b/migrate.js
@@ -47,6 +47,8 @@ function getColumnQueryPart(columnName, column) {
   columnQuery.push(`"${columnName}"`);
   if (column.type === "ID") {
     columnQuery.push("INTEGER PRIMARY KEY AUTOINCREMENT");
+  } else if (column.type === "FOREIGN") {
+    columnQuery.push("INTEGER");
   } else {
     columnQuery.push(column.type.toUpperCase());
     // if (column.primaryKey) {
@@ -75,6 +77,7 @@ function getTableCreationSqlQuery(name, columns) {
 
   const primaryKeys = [];
   const uniques = [];
+  const foreigns = [];
 
   for (const [columnName, column] of Object.entries(columns)) {
     if (column.type !== "ID" && column.primaryKey) {
@@ -85,6 +88,10 @@ function getTableCreationSqlQuery(name, columns) {
       uniques.push(`"${columnName}"`);
     }
 
+    if (column.type === "FOREIGN" && column.table) {
+      foreigns.push({ column: columnName, table: column.table });
+    }
+
     columnsQuery.push(getColumnQueryPart(columnName, column));
   }
 
@@ -96,6 +103,14 @@ function getTableCreationSqlQuery(name, columns) {
     columnsQuery.push(`UNIQUE(${uniques.join(", ")})`);
   }
 
+  if (foreigns.length > 0) {
+    for (const foreign of foreigns) {
+      columnsQuery.push(
+        `FOREIGN KEY ("${foreign.column}") REFERENCES "${foreign.table}"("id")`
+      );
+    }
+  }
+
   return `CREATE TABLE "${name}" (${columnsQuery.join(", ")});`;
 }
 
diff --git a/migrate.test.js b/migrate.test.js
index 6ad79d9..f0c6a87 100644
--- a/migrate.test.js
+++ b/migrate.test.js
@@ -73,6 +73,38 @@ describe("Creating new table", function () {
   });
 });
 
+describe("Creating two tables and link foreign key", function () {
+  it("should create a table 'species'", function () {
+    createBaseMigration();
+  });
+  it("should create a table 'people'", function () {
+    migrations.createTable("people", {
+      id: { type: "ID" },
+      name: { type: "TEXT", notNull: true, default: "Unnamed person" },
+      talisman: { type: "FOREIGN", table: "species" },
+    });
+  });
+  it("should return correct SQL query", function () {
+    const queries = migrations.getMigrationsSqlQueries({});
+    assert.strictEqual(queries[0].query, "BEGIN TRANSACTION;");
+    assert.strictEqual(
+      queries[1].query,
+      `INSERT INTO "migrations" ("revision", "dw_version", "date_migrated") VALUES (?, ?, ?);`
+    );
+    assert.strictEqual(queries[1].args.length, 3);
+    assert.strictEqual(
+      queries[2].query,
+      `CREATE TABLE "species" ("id" INTEGER PRIMARY KEY AUTOINCREMENT, "name" TEXT NOT NULL DEFAULT 'Unnamed species', "origin" TEXT, "population" INTEGER);`
+    );
+    assert.strictEqual(
+      queries[3].query,
+      `CREATE TABLE "people" ("id" INTEGER PRIMARY KEY AUTOINCREMENT, "name" TEXT NOT NULL DEFAULT 'Unnamed person', "talisman" INTEGER, FOREIGN KEY ("talisman") REFERENCES "species"("id"));`
+    );
+    assert.strictEqual(queries[4].query, "COMMIT TRANSACTION;");
+    assert.strictEqual(queries[5].query, "VACUUM;");
+  });
+});
+
 describe("Creating new column", function () {
   it("should create a table 'species'", function () {
     createBaseMigration();