add options to create custom primary id and type with shorthand format

dataset/__init__.py

@@ -7,9 +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
+from sqlalchemy import Integer, Text
 
-__all__ = ['Database', 'Table', 'freeze', 'connect', 'Integer', 'String']
+__all__ = ['Database', 'Table', 'freeze', 'connect']
 
 
 def connect(url=None, schema=None, reflectMetadata=True):

dataset/persistence/database.py

@@ -7,7 +7,7 @@ from migrate.versioning.util import construct_engine
 from sqlalchemy.pool import NullPool
 from sqlalchemy.schema import MetaData, Column, Index
 from sqlalchemy.schema import Table as SQLATable
-from sqlalchemy import Integer, String
+from sqlalchemy import Integer, Text
 
 from dataset.persistence.table import Table
 from dataset.persistence.util import ResultIter
@@ -97,13 +97,13 @@ class Database(object):
         return list(set(self.metadata.tables.keys() +
                         self._tables.keys()))
 
-    def create_table(self, table_name, primary_id='id', primary_type=Integer):
+    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)) or an Integer as specified with primary_type flag. 
+        Text 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,22 +112,25 @@ 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)
+            # custom id and type
+            table2 = db.create_table('population2', 'age')
+            table3 = db.create_table('population3', primary_id='race', primary_type='Text')
         """
         self._acquire()
         try:
             log.debug("Creating table: %s on %r" % (table_name, self.engine))
             table = SQLATable(table_name, self.metadata)
-            if primary_type is Integer:
+            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:
+            elif primary_type is 'Text':
-                col = Column(primary_id, String(255), primary_key=True)
+                col = Column(primary_id, Text, primary_key=True)
             else:
-                raise DatasetException("The primary_type has to be either int or str.")
+                raise DatasetException(
+                    "The primary_type has to be either 'Integer' or 'Text'.")
 
             table.append_column(col)
             table.create(self.engine)
@@ -157,10 +160,12 @@ class Database(object):
         finally:
             self._release()
 
-    def get_table(self, table_name):
+    def get_table(self, table_name, primary_id='id', primary_type='Integer'):
         """
         Smart wrapper around *load_table* and *create_table*. Either loads a table
         or creates it if it doesn't exist yet.
+        For short-hand to create a table with custom id and type using [], where
+        table_name, primary_id, and primary_type are specified as a tuple
 
         Returns a :py:class:`Table <dataset.Table>` instance.
         ::
@@ -168,6 +173,10 @@ class Database(object):
             table = db.get_table('population')
             # you can also use the short-hand syntax:
             table = db['population']
+
+            # custom id and type
+            table2 = db['population2', 'age'] # default type is 'Integer'
+            table3 = db['population3', 'race', 'Text']
         """
         if table_name in self._tables:
             return Table(self, self._tables[table_name])
@@ -176,12 +185,15 @@ class Database(object):
             if self.engine.has_table(table_name, schema=self.schema):
                 return self.load_table(table_name)
             else:
-                return self.create_table(table_name)
+                return self.create_table(table_name, primary_id, primary_type)
         finally:
             self._release()
 
     def __getitem__(self, table_name):
-        return self.get_table(table_name)
+        if type(table_name) is tuple:
+            return self.get_table(*table_name[:3])
+        else:
+            return self.get_table(table_name)
 
     def query(self, query, **kw):
         """

test/test_persistence.py

@@ -2,7 +2,7 @@ import os
 import unittest
 from datetime import datetime
 
-from dataset import connect, Integer, String
+from dataset import connect
 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, primary_type=String)
+        table = self.db.create_table("foo2", pid, 'Text')
         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, primary_type=Integer)
+        table = self.db.create_table("foo3", primary_id = pid)
         assert table.table.exists()
         assert len(table.table.columns) == 1, table.table.columns
         assert pid in table.table.c, table.table.c
@@ -58,6 +58,31 @@ class DatabaseTestCase(unittest.TestCase):
         with self.assertRaises(IntegrityError):
             table.insert({'int_id': 123})
 
+    def test_create_table_shorthand1(self):
+        pid = "int_id"
+        table = self.db['foo4', pid]
+        assert table.table.exists
+        assert len(table.table.columns) == 1, table.table.columns
+        assert pid in table.table.c, table.table.c
+
+        table.insert({'int_id': 123})
+        table.insert({'int_id': 124})
+        assert table.find_one(int_id = 123)[0] == 123
+        assert table.find_one(int_id = 124)[0] == 124
+        with self.assertRaises(IntegrityError):
+            table.insert({'int_id': 123})
+
+    def test_create_table_shorthand2(self):
+        pid = "string_id"
+        table = self.db['foo5', pid, 'Text']
+        assert table.table.exists
+        assert len(table.table.columns) == 1, table.table.columns
+        assert pid in table.table.c, table.table.c
+
+        table.insert({
+            'string_id': 'foobar'})
+        assert table.find_one(string_id = 'foobar')[0] == 'foobar'
+
     def test_load_table(self):
         tbl = self.db.load_table('weather')
         assert tbl.table == self.tbl.table