4

Allow upserting a pandas dataframe to a postgres table (equivalent to df.to_sql(...

 1 year ago
source link: https://gist.github.com/pedrovgp/b46773a1240165bf2b1448b3f70bed32
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

jBbL81 commented Jan 7, 2022

edited

how do i solve this error:

UndefinedColumn                           Traceback (most recent call last)
~\AppData\Local\Programs\Python\Python38\lib\site-packages\sqlalchemy\engine\base.py in _execute_context(self, dialect, constructor, statement, parameters, execution_options, *args, **kw)
   1801                 if not evt_handled:
-> 1802                     self.dialect.do_execute(
   1803                         cursor, statement, parameters, context

~\AppData\Local\Programs\Python\Python38\lib\site-packages\sqlalchemy\engine\default.py in do_execute(self, cursor, statement, parameters, context)
    731     def do_execute(self, cursor, statement, parameters, context=None):
--> 732         cursor.execute(statement, parameters)
    733 

UndefinedColumn: column "None" named in key does not exist

Author

pedrovgp commented Jan 7, 2022

It does not seem related to this snippet. Can you send the head(5) of your dataframe? I am guessing there is something wrong with the column names. Maybe duplicated column names? Or empty ones?

jBbL81 commented Jan 7, 2022

edited

here's the full code:

scraped_data_list = []

n = 1
while n <= 2:
    get_page_number = driver.find_element(By.CSS_SELECTOR, 'span.pageNumberElement').text
    print(f'Working on page {get_page_number}...')

    soup = BeautifulSoup(driver.page_source, 'lxml')
    listings = soup.select('li.productListItem')

    for listing in listings[:1]:
        title = listing.select_one('.bc-heading').get_text().strip()
        link = listing.select_one('.bc-link').get('href')
        link = link.split('?')[0]
        link = 'https://www.audible.com' + link
        date = listing.select_one('li.releaseDateLabel').get_text()
        date = date.split(':')[1].strip() # removes 'Release date:' and any whitespace
        date = datetime.strptime(date, "%m-%d-%y").strftime("%Y-%m-%d")
        
        try:
            rating = listing.select('li.ratingsLabel span')[6].get_text()
            rating = rating.split(' ')[0] # removes 'ratings'
            rating = rating.replace(',', '')
            rating = int(rating)
        except:
            rating = 0
        
        scraped_data_dict = {
            'title': title,
            'rating': rating,
            'date': date,
            'link': link
        }

        scraped_data_list.append(scraped_data_dict)

    driver.find_element(By.CSS_SELECTOR, 'span.nextButton a').click()
    time.sleep(random.randint(5, 10))    
    n += 1

df = pd.DataFrame(scraped_data_list)
df

i only have 4 rows in my dataframe (my scraper is still in the testing phase so i didn't want to scrape thousands of rows) and here's the full error (sorry that it's in image format):

Author

pedrovgp commented Jan 11, 2022

You did not set an index for your table, the index name is None by default. Just set an index or index name and it should work.

Chunkford commented Jan 12, 2022

edited

Thanks for creating this.
I needed to use a different schema and set the dtypes, if it's of any use?

  ....
  from typing import Dict
  
  
  def upsert_df(df: pd.DataFrame, table_name: str, engine: sqlalchemy.engine.Engine, schema: str='public', dtypes: Dict=None):
      """Implements the equivalent of pd.DataFrame.to_sql(..., if_exists='update')
      (which does not exist). Creates or updates the db records based on the
      dataframe records.
      Conflicts to determine update are based on the dataframes index.
      This will set primary keys on the table equal to the index names
      1. Create a temp table from the dataframe
      2. Insert/update from temp table into table_name
      Returns: True if successful
      """
  
      # If the table does not exist, we should just use to_sql to create it
      if not engine.execute(
          f"""SELECT EXISTS (
              SELECT FROM information_schema.tables 
              WHERE  table_schema = '{schema}'
              AND    table_name   = '{table_name}');
              """
      ).first()[0]:
          df.to_sql(table_name, engine, schema=schema, dtype=dtypes)
          return True
  
      # If it already exists...
      temp_table_name = f"temp_{uuid.uuid4().hex[:6]}"
      df.to_sql(temp_table_name, engine, schema=schema, dtype=dtypes, index=True)
  
      index = list(df.index.names)
      index_sql_txt = ", ".join([f'"{i}"' for i in index])
      columns = list(df.columns)
      headers = index + columns
      headers_sql_txt = ", ".join(
          [f'"{i}"' for i in headers]
      )  # index1, index2, ..., column 1, col2, ...
  
      # col1 = exluded.col1, col2=excluded.col2
      update_column_stmt = ", ".join([f'"{col}" = EXCLUDED."{col}"' for col in columns])
  
      # For the ON CONFLICT clause, postgres requires that the columns have unique constraint
      query_pk = f"""
      ALTER TABLE "{schema}"."{table_name}" ADD CONSTRAINT {table_name}_unique_constraint_for_upsert UNIQUE ({index_sql_txt});
      """
      try:
          engine.execute(query_pk)
      except Exception as e:
          # relation "unique_constraint_for_upsert" already exists
          if not 'unique_constraint_for_upsert" already exists' in e.args[0]:
              raise e
  
      # Compose and execute upsert query
      query_upsert = f"""
      INSERT INTO "{schema}"."{table_name}" ({headers_sql_txt}) 
      SELECT {headers_sql_txt} FROM "{schema}"."{temp_table_name}"
      ON CONFLICT ({index_sql_txt}) DO UPDATE 
      SET {update_column_stmt};
      """
      engine.execute(query_upsert)
      engine.execute(f'DROP TABLE "{schema}"."{temp_table_name}"')
  
      return True
      
...

bachkukkik commented Mar 9, 2022

edited

As I use engine from sqlmodel which has SQLAlchemy 2.0 underneath. Direct engine.execute() is no longer supported. It needs to be inside connection object so I adjusted your code as below

Also I'm not sure about This InternalError: (psycopg2.errors.InFailedSqlTransaction) current transaction is aborted, commands ignored until end of transaction block when connection object tries to execute "INSERT INTO ..." but adding conn.execute(sqlmodel.text("commit")) helps

ps.1 type(engine) from sqlmodel is sqlalchemy.future.engine.Engine
ps.2 The strings inside conn.execute() needs to be a class of sqlalchemy.sql.elements.TextClause so i cast it by sqlmodel.text(...)

Here is my code:

import pandas as pd
import sqlalchemy
import uuid
import os
import sqlmodel

## https://gist.github.com/pedrovgp/b46773a1240165bf2b1448b3f70bed32
def sqlmodel_upsert_df(engine: sqlalchemy.future.engine.Engine, df: pd.core.frame.DataFrame, db_name: str, table_name: str, if_exists='fail'):
    """
    Implements the equivalent of pd.DataFrame.to_sql(..., if_exists='update')
    (which does not exist). Creates or updates the db records based on the
    dataframe records.
    Conflicts to determine update are based on the dataframes index.
    This will set primary keys on the table equal to the index names
    1. Create a temp table from the dataframe
    Returns: True if successful
    2. Insert/update from temp table into table_name

    Example
    --------
    >>> from modules import database_utils
    >>> import pandas as pd
    >>> import sqlmodel
    >>> df = pd.DataFrame(...)
    >>> engine = sqlmodel.create_engine('postgresql://postgres:password@localhost:5432/postgres')
    >>> database_utils.sqlmodel_upsert_df(df.set_index('some_key'), 'some_table_name', engine)
    """
    ## Check input validity
    assert (isinstance(db_name, str)) and (len(table_name) > 0)
    assert (isinstance(table_name, str)) and (len(table_name) > 0)
    assert (isinstance(df, pd.core.frame.DataFrame)) and (df.shape[0] > 0)

    ## Try to create `db_name` if not exist
    sqlmodel_create_database(engine, db_name)

    with engine.connect() as conn:
        try:
            assert engine.url.database == db_name
        except AssertionError:
            engine = sqlmodel_create_postgres_engine(
                    engine.url.host,
                    engine.url.port,
                    engine.url.username,
                    engine.url.password,
                    db_name
                )

    if if_exists == 'replace':
        with engine.connect() as conn:
            df.to_sql(table_name, engine, if_exists=if_exists)
            return True

    # If the table does not exist, we should just use to_sql to create it
    with engine.connect() as conn:
        if not conn.execute(
            sqlmodel.text(
                f"""SELECT EXISTS (
                    SELECT FROM information_schema.tables 
                    WHERE  table_schema = 'public'
                    AND    table_name   = '{table_name}');
                    """
            )
        ).first()[0]:
            df.to_sql(table_name, engine)
            return True

    # If it already exists...
    temp_table_name = f"temp_{uuid.uuid4().hex[:6]}"
    df.to_sql(temp_table_name, engine, index=True)

    index = list(df.index.names)
    index_sql_txt = ", ".join([f'"{i}"' for i in index])
    columns = list(df.columns)
    headers = index + columns
    headers_sql_txt = ", ".join(
        [f'"{i}"' for i in headers]
    )  # index1, index2, ..., column 1, col2, ...

    # col1 = exluded.col1, col2=excluded.col2
    update_column_stmt = ", ".join([f'"{col}" = EXCLUDED."{col}"' for col in columns])

    # For the ON CONFLICT clause, postgres requires that the columns have unique constraint
    query_pk = f"""
    ALTER TABLE "{table_name}" ADD CONSTRAINT {table_name}_unique_constraint_for_upsert UNIQUE ({index_sql_txt});
    """
    with engine.connect() as conn:
        try:
            conn.execute(sqlmodel.text(query_pk))
        except Exception as e:
            # relation "unique_constraint_for_upsert" already exists
            if not 'unique_constraint_for_upsert" already exists' in e.args[0]:
                raise e

        # Compose and execute upsert query
        query_upsert = f"""
        INSERT INTO "{table_name}" ({headers_sql_txt}) 
        SELECT {headers_sql_txt} FROM "{temp_table_name}"
        ON CONFLICT ({index_sql_txt}) DO UPDATE 
        SET {update_column_stmt};
        """
        conn.execute(sqlmodel.text("commit"))
        conn.execute(sqlmodel.text(query_upsert))
        conn.execute(sqlmodel.text(f'DROP TABLE "{temp_table_name}"'))
        conn.commit()

        return True

df_test = pd.DataFrame(...)
engine = sqlmodel.create_engine(os.getenv("DB_STR"))
upsert_df(df_test.set_index('some_key'), 'table_name', engine)
## True

Hi here this is so useful and impressive @pedrovgp would you have a similar function for SQLite engine instead of PostgreSQL I believe the syntax would vary a bit but keen on trying something similar on SQLite

Author

pedrovgp commented Oct 24, 2022

Hi Pascal, I never tried anything like it in SQLite. You can try providing a sqlite engine (instead of a postgres one) to the function to probe it.


About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK