예제 #1
0
    def _create_sqlite_query(self, values: list, if_row_exists: str) -> str:
        def escape_col(col):
            # unbound column from its table
            # otherwise the column would compile as "table.col_name"
            # which we could not use in e.g. SQlite
            unbound_col = deepcopy(col)
            unbound_col.table = None
            return str(unbound_col.compile(dialect=self.connection.dialect))

        # prepare start of upsert (INSERT VALUES (...) ON CONFLICT)
        upsert = SQLCompiler(dialect=self.connection.dialect,
                             statement=self.table.insert().values(values))

        # append on conflict clause
        pk = [escape_col(c) for c in self.table.primary_key]
        non_pks = [
            escape_col(c) for c in self.table.columns
            if c not in list(self.table.primary_key)
        ]
        ondup = f'ON CONFLICT ({",".join(pk)})'
        # always use "DO NOTHING" if there are no primary keys
        if (not non_pks) or (if_row_exists == 'ignore'):
            ondup_action = 'DO NOTHING'
            upsert.string = ' '.join((upsert.string, ondup, ondup_action))
        elif if_row_exists == 'update':
            ondup_action = 'DO UPDATE SET'
            updates = ', '.join(f'{c}=EXCLUDED.{c}' for c in non_pks)
            upsert.string = ' '.join(
                (upsert.string, ondup, ondup_action, updates))
        return upsert
예제 #2
0
파일: upsert.py 프로젝트: mkrdip/pangres
def sqlite_upsert(engine, connection, table, values, if_row_exists):
    """
    Compiles and executes a SQlite ON CONFLICT...DO NOTHING or DO UPDATE
    statement.

    Parameters
    ----------
    engine : sqlalchemy.engine.base.Engine
    connection : sqlalchemy.engine.base.Connection
    table : sqlalchemy.sql.schema.Table
    values : list of dict
    if_row_exists : {'update', 'ignore'}
        * If 'update' issues a ON CONFLICT...DO UPDATE statement
        * If 'ignore' issues a ON CONFLICT...DO NOTHING statement

    Examples
    --------
    >>> import datetime
    >>> from sqlalchemy import create_engine
    >>> from pangres.examples import _TestsExampleTable
    >>> from pangres.helpers import PandasSpecialEngine
    >>> 
    >>> engine = create_engine('sqlite:///:memory:')
    >>> df = _TestsExampleTable.create_example_df(nb_rows=5)
    >>> df # doctest: +SKIP
    | profileid   | email             | timestamp                 |   size_in_meters | likes_pizza   | favorite_colors              |
    |:------------|:------------------|:--------------------------|-----------------:|:--------------|:-----------------------------|
    | abc0        | [email protected]  | 2007-10-11 23:15:06+00:00 |          1.93994 | False         | ['yellow', 'blue']           |
    | abc1        | [email protected]  | 2007-11-21 07:18:20+00:00 |          1.98637 | True          | ['blue', 'pink']             |
    | abc2        | [email protected] | 2002-09-30 17:55:09+00:00 |          1.55945 | True          | ['blue']                     |
    | abc3        | [email protected]      | 2007-06-13 22:08:36+00:00 |          2.2495  | True          | ['orange', 'blue']           |
    | abc4        | [email protected]     | 2004-11-22 04:54:09+00:00 |          2.2019  | False         | ['orange', 'yellow', 'blue'] |

    >>> pse = PandasSpecialEngine(engine=engine, df=df, table_name='test_upsert_sqlite')
    >>> 
    >>> insert_values = {'profileid':'abc5', 'email': '*****@*****.**',
    ...                  'timestamp': datetime.datetime(2019, 1, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
    ...                  'size_in_meters':1.9,
    ...                  'likes_pizza':True,
    ...                  'favorite_colors':['red', 'pink']}
    >>> 
    >>> sqlite_upsert(engine=engine, connection=engine.connect(), table=pse.table,
    ...               values=list(insert_values.values()), if_row_exists='update')  # doctest: +SKIP
    """
    def escape_col(col):
        # unbound column from its table
        # otherwise the column would compile as "table.col_name"
        # which we could not use in e.g. SQlite
        unbound_col = deepcopy(col)
        unbound_col.table = None
        return str(unbound_col.compile(dialect=engine.dialect))

    # prepare start of insert (INSERT VALUES (...) ON CONFLICT)
    pk = [escape_col(c) for c in table.primary_key]
    insert = SQLCompiler(dialect=engine.dialect,
                         statement=table.insert().values(values))

    # append on conflict clause
    pk = [escape_col(c) for c in table.primary_key]
    ondup = f'ON CONFLICT ({",".join(pk)})'
    if if_row_exists == 'ignore':
        ondup_action = 'DO NOTHING'
        insert.string = ' '.join((insert.string, ondup, ondup_action))
    elif if_row_exists == 'update':
        ondup_action = 'DO UPDATE SET'
        non_pks = [
            escape_col(c) for c in table.columns
            if c not in list(table.primary_key)
        ]
        updates = ', '.join(f'{c}=EXCLUDED.{c}' for c in non_pks)
        insert.string = ' '.join((insert.string, ondup, ondup_action, updates))
    connection.execute(insert)