diff --git a/recommender/scripts/database.py b/recommender/scripts/database.py index 0ff492d..408927f 100644 --- a/recommender/scripts/database.py +++ b/recommender/scripts/database.py @@ -14,7 +14,10 @@ def get_db_connection(): db_password = os.getenv('MOVIE_MATCH_DB_PASSWORD') db_host = os.getenv('MOVIE_MATCH_DB_HOST') - return create_engine(f'postgresql+psycopg2://{db_username}:{db_password}@{db_host}/movie_match', pool_recycle=3600).connect() + return create_engine( + f'postgresql+psycopg2://{db_username}:{db_password}@{db_host}/movie_match', + pool_recycle=3600 + ).connect() def get_voted_media(conn: Connection, user_id: str) -> DataFrame: diff --git a/recommender/scripts/predict_media_votes.py b/recommender/scripts/predict_media_votes.py index 7b96433..8646b95 100644 --- a/recommender/scripts/predict_media_votes.py +++ b/recommender/scripts/predict_media_votes.py @@ -1,3 +1,5 @@ +from sqlalchemy import text + from database import get_db_connection, get_all_user_ids from predict import predict_votes @@ -9,17 +11,25 @@ p_media = predict_votes(conn, user_id) for idx, row in p_media.iterrows(): + print(f'inserting vote prediction {row["vote_type"]} for {row["media_id"]}.') + conn.execute( - 'insert into media_user_vote_prediction (' - 'media_id, user_id, predicted_vote, created_at, updated_at' - ') values (' - '%(media_id)s, %(user_id)s, %(predicted_vote)s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP' - ')' - ' on conflict (media_id, user_id) do update set' - ' predicted_vote = %(predicted_vote)s, updated_at = CURRENT_TIMESTAMP', + text( + 'insert into media_user_vote_prediction (' + 'media_id, user_id, predicted_vote, created_at, updated_at' + ') values (' + ':media_id, :user_id, :predicted_vote, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP' + ')' + ' on conflict (media_id, user_id) do update set' + ' predicted_vote = :predicted_vote, updated_at = CURRENT_TIMESTAMP' + ), { 'media_id': row['media_id'], 'user_id': user_id, 'predicted_vote': row['vote_type'], } ) + + conn.commit() + +conn.close()