Python Pandas to_sql, как создать таблицу с первичным ключом?

Я хотел бы создать таблицу MySQL с помощью функции to_sql от Pandas, которая имеет первичный ключ (обычно полезно иметь первичный ключ в таблице mysql) следующим образом:

group_export.to_sql(con = db, name = config.table_group_export, if_exists = 'replace', flavor = 'mysql', index = False)

но при этом создается таблица без какого-либо первичного ключа (или даже без индекса).

В документации упоминается параметр index_label, который в сочетании с параметром index можно использовать для создания индекса, но не упоминается ни один параметр для первичных ключей.

Документация

Ответ 1

Отказ от ответственности: этот ответ более экспериментальный, чем практический, но, возможно, стоит упомянуть.

Я обнаружил, что класс pandas.io.sql.SQLTable назвал аргумент key, и если вы присвоите ему имя поля, тогда это поле станет первичным ключом:

К сожалению, вы не можете просто передать этот аргумент из функции DataFrame.to_sql(). Чтобы использовать его, вы должны:

создать pandas.io.SQLDatabase экземпляр

engine = sa.create_engine('postgresql:///somedb')
pandas_sql = pd.io.sql.pandasSQL_builder(engine, schema=None, flavor=None)

определить функцию, аналогичную pandas.io.SQLDatabase.to_sql(), но с дополнительным аргументом *kwargs, который передается в объект pandas.io.SQLTable, созданный внутри него (я только что скопировал оригинальный метод to_sql() и добавил *kwargs):

def to_sql_k(self, frame, name, if_exists='fail', index=True,
           index_label=None, schema=None, chunksize=None, dtype=None, **kwargs):
    if dtype is not None:
        from sqlalchemy.types import to_instance, TypeEngine
        for col, my_type in dtype.items():
            if not isinstance(to_instance(my_type), TypeEngine):
                raise ValueError('The type of %s is not a SQLAlchemy '
                                 'type ' % col)

    table = pd.io.sql.SQLTable(name, self, frame=frame, index=index,
                     if_exists=if_exists, index_label=index_label,
                     schema=schema, dtype=dtype, **kwargs)
    table.create()
    table.insert(chunksize)

вызов этой функции с помощью экземпляра SQLDatabase и блока данных, который вы хотите сохранить
```
to_sql_k(pandas_sql, df2save, 'tmp',
        index=True, index_label='id', keys='id', if_exists='replace')
```

И мы получаем что-то вроде

CREATE TABLE public.tmp
(
  id bigint NOT NULL DEFAULT nextval('tmp_id_seq'::regclass),
...
)

в базе данных.

PS Вы можете, конечно, использовать функции обезьяны-патча DataFrame, io.SQLDatabase и io.to_sql(), чтобы использовать это обходное решение с удобством.

Ответ 2

Просто добавьте первичный ключ после загрузки таблицы с помощью pandas.

group_export.to_sql(con=engine, name=example_table, if_exists='replace', 
                    flavor='mysql', index=False)

with engine.connect() as con:
    con.execute('ALTER TABLE `example_table` ADD PRIMARY KEY (`ID_column`);')

Ответ 3

automap_base from sqlalchemy.ext.automap (tableNamesDict - это dict с таблицами Pandas):

metadata = MetaData()
metadata.reflect(db.engine, only=tableNamesDict.values())
Base = automap_base(metadata=metadata)
Base.prepare()

Что бы сработало отлично, за исключением одной проблемы, automap требует, чтобы таблицы имели первичный ключ. Хорошо, не проблема, я уверен, что Pandas to_sql имеет способ указать первичный ключ... nope. Здесь он становится немного взломанным:

for df in dfs.keys():
    cols = dfs[df].columns
    cols = [str(col) for col in cols if 'id' in col.lower()]
    schema = pd.io.sql.get_schema(dfs[df],df, con=db.engine, keys=cols)
    db.engine.execute('DROP TABLE ' + df + ';')
    db.engine.execute(schema)
    dfs[df].to_sql(df,con=db.engine, index=False, if_exists='append')

Я перебираю dict из DataFrames, получаю список столбцов для первичного ключа (т.е. те, которые содержат id), используйте get_schema для создания пустых таблиц, затем добавьте DataFrame к таблице.

Теперь, когда у вас есть модели, вы можете явно указывать и использовать их (т.е. User = Base.classes.user) с помощью session.query или создавать список всех классов с чем-то вроде этого:

alchemyClassDict = {}
for t in Base.classes.keys():
    alchemyClassDict[t] = Base.classes[t]

И запрос с помощью:

res = db.session.query(alchemyClassDict['user']).first()