Skip to content

Instantly share code, notes, and snippets.

@RobinL

RobinL/add_pk_to_sqlite_table.py

Created November 17, 2017 09:14
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save RobinL/c27fef037cbd2b24c909d5275b01adfb to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save RobinL/c27fef037cbd2b24c909d5275b01adfb to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. RobinL created this gist Nov 17, 2017.
    46 changes: 46 additions & 0 deletions add_pk_to_sqlite_table.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,46 @@
    import re

    def get_create_table_string(tablename, connection):
    sql = """
    select * from sqlite_master where name = "{}" and type = "table"
    """.format(tablename)
    result = connection.execute(sql)

    create_table_string = result.fetchmany()[0][4]
    return create_table_string

    def add_pk_to_create_table_string(create_table_string, colname):
    regex = "(\n.+{}[^,]+)(,)".format(colname)
    return re.sub(regex, "\\1 PRIMARY KEY,", create_table_string)

    def add_pk_to_sqlite_table(tablename, index_column, connection):
    cts = get_create_table_string(tablename, connection)
    cts = add_pk_to_create_table_string(cts, index_column)
    template = """
    BEGIN TRANSACTION;
    ALTER TABLE {tablename} RENAME TO {tablename}_old_;
    {cts};
    INSERT INTO {tablename} SELECT * FROM {tablename}_old_;
    DROP TABLE {tablename}_old_;
    COMMIT TRANSACTION;
    """

    create_and_drop_sql = template.format(tablename = tablename, cts = cts)
    connection.executescript(create_and_drop_sql)

    # Example:

    # import pandas as pd
    # import sqlite3

    # df = pd.DataFrame({"a": [1,2,3], "b": [2,3,4]})
    # con = sqlite3.connect("deleteme.db")
    # df.to_sql("df", con, if_exists="replace")

    # add_pk_to_sqlite_table("df", "index", con)
    # r = con.execute("select sql from sqlite_master where name = 'df' and type = 'table'")
    # print(r.fetchone()[0])