From 9611e5edf0adab7c4a1ed75efd8c21ea431e2e00 Mon Sep 17 00:00:00 2001
From: David Krause <david.krause.de@gmail.com>
Date: Thu, 13 Sep 2018 19:32:33 -0700
Subject: [PATCH] METAMODEL-1197: Fixed Kafka module's Column-to-Table
 references

Closes #188
---
 CHANGES.md                                             |  1 +
 .../org/apache/metamodel/schema/MutableColumn.java     |  6 ++++++
 .../org/apache/metamodel/kafka/KafkaDataContext.java   | 10 +++++-----
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index f9b79ef9d..d37939c46 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -3,6 +3,7 @@
  * [METAMODEL-1194] - Improved rewriting of CREATE TABLE for Apache Hive when user specifies primary key(s).
  * [METAMODEL-1192] - Add support for Oracle "TIMESTAMP WITH (LOCAL_)TIME_ZONE" column type.
  * [METAMODEL-1193] - Updated release process to produce SHA-512 signature files.
+ * [METAMODEL-1197] - Fixed references in Kafka module's Columns to Tables.
 
 ### Apache MetaModel 5.1.0
 
diff --git a/core/src/main/java/org/apache/metamodel/schema/MutableColumn.java b/core/src/main/java/org/apache/metamodel/schema/MutableColumn.java
index 050de154f..b1b6346ed 100644
--- a/core/src/main/java/org/apache/metamodel/schema/MutableColumn.java
+++ b/core/src/main/java/org/apache/metamodel/schema/MutableColumn.java
@@ -56,6 +56,12 @@ public MutableColumn(String name, ColumnType type) {
         setType(type);
     }
 
+    public MutableColumn(String name, ColumnType type, Table table) {
+        this(name);
+        setType(type);
+        setTable(table);
+    }
+
     public MutableColumn(String name, ColumnType type, Table table, int columnNumber, Boolean nullable) {
         this(name, type);
         setColumnNumber(columnNumber);
diff --git a/kafka/src/main/java/org/apache/metamodel/kafka/KafkaDataContext.java b/kafka/src/main/java/org/apache/metamodel/kafka/KafkaDataContext.java
index a4dea2846..7c9790b16 100644
--- a/kafka/src/main/java/org/apache/metamodel/kafka/KafkaDataContext.java
+++ b/kafka/src/main/java/org/apache/metamodel/kafka/KafkaDataContext.java
@@ -94,11 +94,11 @@ protected Schema getMainSchema() throws MetaModelException {
 
         for (String topic : topics) {
             final MutableTable table = new MutableTable(topic, schema);
-            table.addColumn(new MutableColumn(COLUMN_PARTITION, ColumnType.INTEGER));
-            table.addColumn(new MutableColumn(COLUMN_OFFSET, ColumnType.BIGINT));
-            table.addColumn(new MutableColumn(COLUMN_TIMESTAMP, ColumnType.TIMESTAMP));
-            table.addColumn(new MutableColumn(COLUMN_KEY, ColumnTypeImpl.convertColumnType(keyClass)));
-            table.addColumn(new MutableColumn(COLUMN_VALUE, ColumnTypeImpl.convertColumnType(valueClass)));
+            table.addColumn(new MutableColumn(COLUMN_PARTITION, ColumnType.INTEGER, table).setPrimaryKey(true));
+            table.addColumn(new MutableColumn(COLUMN_OFFSET, ColumnType.BIGINT, table).setPrimaryKey(true));
+            table.addColumn(new MutableColumn(COLUMN_TIMESTAMP, ColumnType.TIMESTAMP, table));
+            table.addColumn(new MutableColumn(COLUMN_KEY, ColumnTypeImpl.convertColumnType(keyClass), table));
+            table.addColumn(new MutableColumn(COLUMN_VALUE, ColumnTypeImpl.convertColumnType(valueClass), table));
             schema.addTable(table);
         }
         return schema;