New -pk/--primary-key options, closes #22

Author: simonw

README.md

@@ -133,6 +133,7 @@ They will be populated with IDs that reference the new derived tables.
                                    datetimes
       -df, --datetime-format TEXT  One or more custom date format strings to try
                                    when parsing dates/datetimes
+      -pk, --primary-key TEXT      One or more columns to use as the primary key
       -f, --fts TEXT               One or more columns to use to populate a full-
                                    text index
       -i, --index TEXT             Add index on this column (or a compound index

csvs_to_sqlite/cli.py

@@ -53,6 +53,9 @@
 @click.option('--datetime-format', '-df', multiple=True, help=(
     "One or more custom date format strings to try when parsing dates/datetimes"
 ))
+@click.option('--primary-key', '-pk', multiple=True, help=(
+    "One or more columns to use as the primary key"
+))
 @click.option('--fts', '-f', multiple=True, help=(
     "One or more columns to use to populate a full-text index"
 ))
@@ -76,6 +79,7 @@ def cli(
     date,
     datetime,
     datetime_format,
+    primary_key,
     fts,
     index,
     shape,
@@ -155,6 +159,7 @@ def cli(
         else:
             to_sql_with_foreign_keys(
                 conn, df, df.table_name, foreign_keys, sql_type_overrides,
+                primary_keys=primary_key,
                 index_fks=not no_index_fks
             )
             created_tables[df.table_name] = df

csvs_to_sqlite/utils.py

@@ -194,7 +194,7 @@ def drop_table(conn, table):
     conn.execute('DROP TABLE [{}]'.format(table))
 
 
-def get_create_table_sql(table_name, df, index=True, sql_type_overrides=None, **extra_args):
+def get_create_table_sql(table_name, df, index=True, sql_type_overrides=None, primary_keys=None):
     # Create a temporary table with just the first row
     # We do this in memory because we just want to get the
     # CREATE TABLE statement
@@ -226,7 +226,7 @@ def get_create_table_sql(table_name, df, index=True, sql_type_overrides=None, **
                 # Everything was NaN or an integer-float - switch type:
                 sql_type_overrides[column] = 'INTEGER'
 
-    df[:1].to_sql(table_name, conn, index=index, dtype=sql_type_overrides, **extra_args)
+    df[:1].to_sql(table_name, conn, index=index, dtype=sql_type_overrides)
     sql = conn.execute(
         'select sql from sqlite_master where name = ?', [table_name]
     ).fetchone()[0]
@@ -235,11 +235,17 @@ def get_create_table_sql(table_name, df, index=True, sql_type_overrides=None, **
             'PRAGMA table_info([{}])'.format(table_name)
         )
     ]
+    if primary_keys:
+        # Rewrite SQL to add PRIMARY KEY (col1, col2) at end
+        assert sql[-1] == ')'
+        sql = sql[:-1] + '  ,PRIMARY KEY ({cols})\n)'.format(
+            cols=', '.join('[{}]'.format(col) for col in primary_keys)
+        )
     return sql, columns
 
 
-def to_sql_with_foreign_keys(conn, df, name, foreign_keys, sql_type_overrides=None, index_fks=False):
-    create_sql, columns = get_create_table_sql(name, df, index=False, sql_type_overrides=sql_type_overrides)
+def to_sql_with_foreign_keys(conn, df, name, foreign_keys, sql_type_overrides=None, primary_keys=None, index_fks=False):
+    create_sql, columns = get_create_table_sql(name, df, index=False, primary_keys=primary_keys, sql_type_overrides=sql_type_overrides)
     foreign_key_bits = []
     index_bits = []
     for column, (table, value_column) in foreign_keys.items():

tests/test_csvs_to_sqlite.py

@@ -23,6 +23,10 @@
 CSV_DATES_CUSTOM_FORMAT = '''headline,date
 Custom format,03/02/01'''
 
+CSV_CUSTOM_PRIMARY_KEYS = '''pk1,pk2,name
+one,one,11
+one,two,12
+two,one,21'''
 
 def test_flat():
     runner = CliRunner()
@@ -430,3 +434,22 @@ def test_extract_cols_no_fts():
             where type='table' and name like '%_fts'
             and sql like '%USING FTS%'
         ''').fetchall()
+
+
+def test_custom_primary_keys():
+    runner = CliRunner()
+    with runner.isolated_filesystem():
+        open('pks.csv', 'w').write(CSV_CUSTOM_PRIMARY_KEYS)
+        result = runner.invoke(
+            cli.cli, (
+                'pks.csv pks.db -pk pk1 --primary-key pk2'
+            ).split()
+        )
+        assert result.exit_code == 0
+        conn = sqlite3.connect('pks.db')
+        pks = [
+            r[1]
+            for r in conn.execute('PRAGMA table_info("pks")').fetchall()
+            if r[-1]
+        ]
+        assert ['pk1', 'pk2'] == pks