From bfbf03887310bac3b79a8fdc8366079cab823135 Mon Sep 17 00:00:00 2001
From: Ulf Wiger <ulf@feuerlabs.com>
Date: Fri, 24 Feb 2012 09:11:24 +0100
Subject: [PATCH] list_tables/1 now returns table names as they were first
 given.

The create_table/3 function will cast all table names to binaries, allowing
them to be given as atoms, binaries or strings. In order to recall what the
original type was, the CREATE TABLE statement includes a constant CHECK
constraint, e.g. CHECK('bin'='bin'), which is then parsed by the list_tables
function. The constraint is put last in the statement, and should not
interfere with any other constraints. Hopefully, the sqlite compiler is good
enough that this constant expression doesn't add too much overhead. :)
---
 src/sqlite3.erl     | 18 ++++++++++++++++--
 src/sqlite3_lib.erl | 22 +++++++++++++++++-----
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/sqlite3.erl b/src/sqlite3.erl
index 038aa1d..0d8c01c 100644
--- a/src/sqlite3.erl
+++ b/src/sqlite3.erl
@@ -863,10 +863,10 @@ handle_call(close, _From, State) ->
     Reply = ok,
     {stop, normal, Reply, State};
 handle_call(list_tables, _From, State) ->
-    SQL = "select name from sqlite_master where type='table';",
+    SQL = "select name, sql from sqlite_master where type='table';",
     Data = do_sql_exec(SQL, State),
     TableList = proplists:get_value(rows, Data),
-    TableNames = [erlang:list_to_atom(erlang:binary_to_list(Name)) || {Name} <- TableList],
+    TableNames = [cast_table_name(Name, SQLx) || {Name,SQLx} <- TableList],
     {reply, TableNames, State};
 handle_call({table_info, Tbl}, _From, State) when is_atom(Tbl) ->
     % make sure we only get table info.
@@ -1242,6 +1242,20 @@ build_primary_key_constraint(Tail, []) ->
 build_primary_key_constraint(Tail, Acc) ->
     {{primary_key, lists:reverse(Acc)}, Tail}.
 
+cast_table_name(Bin, SQL) ->
+    case re:run(SQL,<<"CHECK\\((.*)=(.*)\\)\\)">>,[{capture,all_but_first,binary}]) of
+	{match, [<<"'bin'">>, <<"'bin'">>]} ->
+	    Bin;
+	{match, [<<"'lst'">>, <<"'lst'">>]} ->
+	    binary_to_atom(Bin, latin1);
+	{match, [<<"'am'">>, <<"'am'">>]} ->
+	    binary_to_atom(Bin, latin1);
+	_ ->
+	    %% backwards compatible
+	    binary_to_atom(Bin, latin1)
+    end.
+
+
 %% conflict_clause(["ON", "CONFLICT", ResolutionString | Tail]) ->
 %%     Resolution = case ResolutionString of
 %%                      "ROLLBACK" -> rollback;
diff --git a/src/sqlite3_lib.erl b/src/sqlite3_lib.erl
index 9695fa2..b0bb2b7 100644
--- a/src/sqlite3_lib.erl
+++ b/src/sqlite3_lib.erl
@@ -205,8 +205,10 @@ read_cols_sql(Columns) ->
 %%--------------------------------------------------------------------
 -spec create_table_sql(table(), table_info()) -> iolist().
 create_table_sql(Tbl, Columns) ->
-    ["CREATE TABLE ", to_iolist(Tbl), " (",
-     map_intersperse(fun column_sql_for_create_table/1, Columns, ", "), ");"].
+    {Type, TName} = encode_tbl(Tbl),
+    ["CREATE TABLE ", TName, " (",
+     map_intersperse(fun column_sql_for_create_table/1, Columns, ", "),
+     " CHECK('", Type, "'='", Type, "'));"].
 
 %%--------------------------------------------------------------------
 %% @spec create_table_sql(Tbl :: table(), ColumnData, TableConstraints) -> iolist()
@@ -221,10 +223,11 @@ create_table_sql(Tbl, Columns) ->
 %%--------------------------------------------------------------------
 -spec create_table_sql(table(), table_info(), table_constraints()) -> iolist().
 create_table_sql(Tbl, Columns, TblConstraints) ->
-    ["CREATE TABLE ", to_iolist(Tbl), " (",
+    {Type, TName} = encode_tbl(Tbl),
+    ["CREATE TABLE ", TName, " (",
      map_intersperse(fun column_sql_for_create_table/1, Columns, ", "), ", ",
      table_constraint_sql(TblConstraints), 
-     ");"].
+     " CHECK('", Type, "'='", Type, "'));"].
 
 %%--------------------------------------------------------------------
 %% @spec update_sql(Tbl, Key, Value, Data) -> iolist()
@@ -240,7 +243,8 @@ create_table_sql(Tbl, Columns, TblConstraints) ->
 %%--------------------------------------------------------------------
 -spec update_sql(table(), atom(), sql_value(), [{atom(), sql_value()}]) -> iolist().
 update_sql(Tbl, Key, Value, Data) ->
-    ["UPDATE ", to_iolist(Tbl), " SET ", update_set_sql(Data), 
+    {_, TName} = encode_tbl(Tbl),
+    ["UPDATE ", TName, " SET ", update_set_sql(Data), 
      " WHERE ", atom_to_list(Key), " = ", value_to_sql(Value), ";"].
 
 %%--------------------------------------------------------------------
@@ -420,6 +424,14 @@ indexed_column_sql({ColumnName, asc}) -> [atom_to_list(ColumnName), " ASC"];
 indexed_column_sql({ColumnName, desc}) -> [atom_to_list(ColumnName), " DESC"];
 indexed_column_sql(ColumnName) -> atom_to_list(ColumnName).
 
+
+encode_tbl(A) when is_atom(A) ->
+    {"am", atom_to_list(A)};
+encode_tbl(B) when is_binary(B) ->
+    {"bin", B};
+encode_tbl(L) when is_list(L) ->
+    {"lst", list_to_binary(L)}.
+
 to_iolist(A) when is_atom(A) ->
     atom_to_list(A);
 to_iolist(L) when is_list(L) ->