From 95cc5dd411f0ca690dd8c65e2a602664d547e22c Mon Sep 17 00:00:00 2001
From: Yi Xie <dnatag@gmail.com>
Date: Fri, 13 Sep 2013 09:42:35 -0400
Subject: [PATCH] simplify creat_table interface

---
 dataset/__init__.py             |  3 ++-
 dataset/persistence/database.py | 24 ++++++++++++++----------
 test/test_persistence.py        |  6 +++---
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/dataset/__init__.py b/dataset/__init__.py
index 4441aee..f6cf604 100644
--- a/dataset/__init__.py
+++ b/dataset/__init__.py
@@ -7,8 +7,9 @@ warnings.filterwarnings(
 from dataset.persistence.database import Database
 from dataset.persistence.table import Table
 from dataset.freeze.app import freeze
+from sqlalchemy import Integer, String
 
-__all__ = ['Database', 'Table', 'freeze', 'connect']
+__all__ = ['Database', 'Table', 'freeze', 'connect', 'Integer', 'String']
 
 
 def connect(url=None, schema=None, reflectMetadata=True):
diff --git a/dataset/persistence/database.py b/dataset/persistence/database.py
index 20a1519..1ae27b2 100644
--- a/dataset/persistence/database.py
+++ b/dataset/persistence/database.py
@@ -11,7 +11,7 @@ from sqlalchemy import Integer, String
 
 from dataset.persistence.table import Table
 from dataset.persistence.util import ResultIter
-
+from dataset.util import DatasetException
 
 log = logging.getLogger(__name__)
 
@@ -97,14 +97,14 @@ class Database(object):
         return list(set(self.metadata.tables.keys() +
                         self._tables.keys()))
 
-    def create_table(self, table_name, primary_id=None, is_integer=False):
+    def create_table(self, table_name, primary_id='id', primary_type=Integer):
         """
         Creates a new table. The new table will automatically have an `id` column 
         unless specified via optional parameter primary_id, which will be used 
         as the primary key of the table. Automatic id is set to be an 
         auto-incrementing integer, while the type of custom primary_id can be a 
-        string (i.e. VARCHAR(255)) by default or an integer via integer flag. 
-        The caller will be responsible for the uniqueness of manually specified primary_id.
+        String (i.e. VARCHAR(255)) or an Integer as specified with primary_type flag. 
+        The caller will be responsible for the uniqueness of manual primary_id.
 
         This custom id feature is only available via direct create_table call. 
 
@@ -112,18 +112,22 @@ class Database(object):
         ::
 
             table = db.create_table('population')
+            table2 = db.create_table('population2', primary_id='age', primary_type=Integer)
+            table3 = db.create_table('population3', primary_id='race', primary_type=String)
         """
         self._acquire()
         try:
             log.debug("Creating table: %s on %r" % (table_name, self.engine))
             table = SQLATable(table_name, self.metadata)
-            if not primary_id:
-                col = Column('id', Integer, primary_key=True)
+            if primary_type is Integer:
+                auto_flag = False
+                if primary_id is 'id':
+                    auto_flag = True
+                col = Column(primary_id, Integer, primary_key=True, autoincrement=auto_flag)
+            elif primary_type is String:
+                col = Column(primary_id, String(255), primary_key=True)
             else:
-                if is_integer:
-                    col = Column(primary_id, Integer, primary_key=True, autoincrement=False)
-                else:
-                    col = Column(primary_id, String(255), primary_key=True)
+                raise DatasetException("The primary_type has to be either int or str.")
 
             table.append_column(col)
             table.create(self.engine)
diff --git a/test/test_persistence.py b/test/test_persistence.py
index 80614b8..063032c 100644
--- a/test/test_persistence.py
+++ b/test/test_persistence.py
@@ -2,7 +2,7 @@ import os
 import unittest
 from datetime import datetime
 
-from dataset import connect
+from dataset import connect, Integer, String
 from dataset.util import DatasetException
 from sample_data import TEST_DATA
 from sqlalchemy.exc import IntegrityError
@@ -35,7 +35,7 @@ class DatabaseTestCase(unittest.TestCase):
 
     def test_create_table_custom_id1(self):
         pid = "string_id"
-        table = self.db.create_table("foo2", primary_id = pid)
+        table = self.db.create_table("foo2", primary_id = pid, primary_type=String)
         assert table.table.exists()
         assert len(table.table.columns) == 1, table.table.columns
         assert pid in table.table.c, table.table.c
@@ -46,7 +46,7 @@ class DatabaseTestCase(unittest.TestCase):
 
     def test_create_table_custom_id2(self):
         pid = "int_id"
-        table = self.db.create_table("foo3", primary_id = pid, is_integer=True)
+        table = self.db.create_table("foo3", primary_id = pid, primary_type=Integer)
         assert table.table.exists()
         assert len(table.table.columns) == 1, table.table.columns
         assert pid in table.table.c, table.table.c