[Do not merge] Speed up LB tests

admin/sql/truncate_tables.sql

@@ -0,0 +1,19 @@
+BEGIN;
+
+TRUNCATE "user"                         CASCADE;
+TRUNCATE data_dump                      CASCADE;
+TRUNCATE spotify_auth                   CASCADE;
+TRUNCATE external_service_oauth         CASCADE;
+TRUNCATE listens_importer               CASCADE;
+TRUNCATE recording_feedback             CASCADE;
+TRUNCATE missing_musicbrainz_data       CASCADE;
+TRUNCATE user_relationship              CASCADE;
+TRUNCATE recommendation_feedback        CASCADE;
+TRUNCATE user_timeline_event            CASCADE;
+TRUNCATE hide_user_timeline_event       CASCADE;
+TRUNCATE reported_users                 CASCADE;
+TRUNCATE pinned_recording               CASCADE;
+TRUNCATE release_color                  CASCADE;
+TRUNCATE user_setting                   CASCADE;
+
+COMMIT;

listenbrainz/db/__init__.py

@@ -1,3 +1,4 @@
+from typing import Optional
 
 import sqlalchemy
 from sqlalchemy import create_engine
@@ -13,7 +14,7 @@
 # public dump
 SCHEMA_VERSION_CORE = 8
 
-engine = None
+engine: Optional[sqlalchemy.engine.Engine] = None
 
 DUMP_DEFAULT_THREAD_COUNT = 4
 

listenbrainz/db/dump.py

@@ -639,7 +639,7 @@ def copy_table(cursor, location, columns, table_name):
         cursor.copy_expert(query, f)
 
 
-def add_dump_entry(timestamp):
+def add_dump_entry(connection, timestamp):
     """ Adds an entry to the data_dump table with specified time.
 
         Args:
@@ -648,44 +648,39 @@ def add_dump_entry(timestamp):
         Returns:
             id (int): the id of the new entry added
     """
-
-    with db.engine.connect() as connection:
-        result = connection.execute(sqlalchemy.text("""
-                INSERT INTO data_dump (created)
-                     VALUES (TO_TIMESTAMP(:ts))
-                  RETURNING id
-            """), {
-            'ts': timestamp,
-        })
-        return result.fetchone()['id']
+    result = connection.execute(sqlalchemy.text("""
+            INSERT INTO data_dump (created)
+                 VALUES (TO_TIMESTAMP(:ts))
+              RETURNING id
+        """), {
+        'ts': timestamp,
+    })
+    return result.fetchone()['id']
 
 
-def get_dump_entries():
+def get_dump_entries(connection):
     """ Returns a list of all dump entries in the data_dump table
     """
-
-    with db.engine.connect() as connection:
-        result = connection.execute(sqlalchemy.text("""
-                SELECT id, created
-                  FROM data_dump
-              ORDER BY created DESC
-            """))
-
-        return [dict(row) for row in result]
-
-
-def get_dump_entry(dump_id):
-    with db.engine.connect() as connection:
-        result = connection.execute(sqlalchemy.text("""
+    result = connection.execute(sqlalchemy.text("""
             SELECT id, created
               FROM data_dump
-             WHERE id = :dump_id
-        """), {
-            'dump_id': dump_id,
-        })
-        if result.rowcount > 0:
-            return dict(result.fetchone())
-        return None
+          ORDER BY created DESC
+        """))
+
+    return [dict(row) for row in result]
+
+
+def get_dump_entry(connection, dump_id):
+    result = connection.execute(sqlalchemy.text("""
+        SELECT id, created
+          FROM data_dump
+         WHERE id = :dump_id
+    """), {
+        'dump_id': dump_id,
+    })
+    if result.rowcount > 0:
+        return dict(result.fetchone())
+    return None
 
 
 def import_postgres_dump(private_dump_archive_path=None,

listenbrainz/db/dump_manager.py

@@ -36,7 +36,7 @@
 from listenbrainz.db.year_in_music import insert_playlists
 from listenbrainz.listenstore.dump_listenstore import DumpListenStore
 from listenbrainz.utils import create_path
-from listenbrainz.webserver import create_app
+from listenbrainz.webserver import create_app, db_conn, ts_conn
 from listenbrainz.db.dump import check_ftp_dump_ages
 
 
@@ -92,9 +92,9 @@ def create_full(location, threads, dump_id, do_listen_dump: bool, do_spark_dump:
         ls = DumpListenStore(app)
         if dump_id is None:
             end_time = datetime.now()
-            dump_id = db_dump.add_dump_entry(int(end_time.strftime('%s')))
+            dump_id = db_dump.add_dump_entry(db_conn, int(end_time.strftime('%s')))
         else:
-            dump_entry = db_dump.get_dump_entry(dump_id)
+            dump_entry = db_dump.get_dump_entry(db_conn, dump_id)
             if dump_entry is None:
                 current_app.logger.error("No dump with ID %d found", dump_id)
                 sys.exit(-1)
@@ -158,16 +158,16 @@ def create_incremental(location, threads, dump_id):
         ls = DumpListenStore(app)
         if dump_id is None:
             end_time = datetime.now()
-            dump_id = db_dump.add_dump_entry(int(end_time.strftime('%s')))
+            dump_id = db_dump.add_dump_entry(db_conn, int(end_time.strftime('%s')))
         else:
-            dump_entry = db_dump.get_dump_entry(dump_id)
+            dump_entry = db_dump.get_dump_entry(db_conn, dump_id)
             if dump_entry is None:
                 current_app.logger.error(
                     "No dump with ID %d found, exiting!", dump_id)
                 sys.exit(-1)
             end_time = dump_entry['created']
 
-        prev_dump_entry = db_dump.get_dump_entry(dump_id - 1)
+        prev_dump_entry = db_dump.get_dump_entry(db_conn, dump_id - 1)
         if prev_dump_entry is None:  # incremental dumps must have a previous dump in the series
             current_app.logger.error(
                 "Invalid dump ID %d, could not find previous dump", dump_id)
@@ -292,7 +292,7 @@ def import_dump(private_archive, private_timescale_archive,
                                      threads)
         if listen_archive:
             from listenbrainz.webserver.timescale_connection import _ts as ls
-            ls.import_listens_dump(listen_archive, threads)
+            ls.import_listens_dump(ts_conn, listen_archive, threads)
 
     sys.exit(0)
 

listenbrainz/db/external_service_oauth.py

@@ -5,8 +5,8 @@
 import sqlalchemy
 
 
-def save_token(user_id: int, service: ExternalServiceType, access_token: str, refresh_token: str,
-               token_expires_ts: int, record_listens: bool, scopes: List[str]):
+def save_token(connection, user_id: int, service: ExternalServiceType, access_token: str,
+               refresh_token: str, token_expires_ts: int, record_listens: bool, scopes: List[str]):
     """ Add a row to the external_service_oauth table for specified user with corresponding tokens and information.
 
     Args:
@@ -26,80 +26,80 @@ def save_token(user_id: int, service: ExternalServiceType, access_token: str, re
     # be explicitly set to the default value (which would have been used if the row was
     # inserted instead).
     token_expires = utils.unix_timestamp_to_datetime(token_expires_ts)
-    with db.engine.connect() as connection:
-        result = connection.execute(sqlalchemy.text("""
-            INSERT INTO external_service_oauth
-            (user_id, service, access_token, refresh_token, token_expires, scopes)
+    result = connection.execute(sqlalchemy.text("""
+        INSERT INTO external_service_oauth
+        (user_id, service, access_token, refresh_token, token_expires, scopes)
+        VALUES
+        (:user_id, :service, :access_token, :refresh_token, :token_expires, :scopes)
+        ON CONFLICT (user_id, service)
+        DO UPDATE SET
+            user_id = EXCLUDED.user_id,
+            service = EXCLUDED.service,
+            access_token = EXCLUDED.access_token,
+            refresh_token = EXCLUDED.refresh_token,
+            token_expires = EXCLUDED.token_expires,
+            scopes = EXCLUDED.scopes,
+        -- using clock_timestamp because value returned by now() is always fixed during a transaction
+        -- this creates issues during tests where we want to test that the last_updated field was updated
+             last_updated = clock_timestamp()
+        RETURNING id
+        """), {
+            "user_id": user_id,
+            "service": service.value,
+            "access_token": access_token,
+            "refresh_token": refresh_token,
+            "token_expires": token_expires,
+            "scopes": scopes,
+        })
+
+    if record_listens:
+        external_service_oauth_id = result.fetchone()['id']
+        connection.execute(sqlalchemy.text("""
+            INSERT INTO listens_importer
+            (external_service_oauth_id, user_id, service)
             VALUES
-            (:user_id, :service, :access_token, :refresh_token, :token_expires, :scopes)
-            ON CONFLICT (user_id, service)
-            DO UPDATE SET
+            (:external_service_oauth_id, :user_id, :service)
+            ON CONFLICT (user_id, service) DO UPDATE SET
+                external_service_oauth_id = EXCLUDED.external_service_oauth_id,
                 user_id = EXCLUDED.user_id,
                 service = EXCLUDED.service,
-                access_token = EXCLUDED.access_token,
-                refresh_token = EXCLUDED.refresh_token,
-                token_expires = EXCLUDED.token_expires,
-                scopes = EXCLUDED.scopes,
-                last_updated = NOW()
-            RETURNING id
+                last_updated = NULL,
+                latest_listened_at = NULL,
+                error_message = NULL
             """), {
-                "user_id": user_id,
-                "service": service.value,
-                "access_token": access_token,
-                "refresh_token": refresh_token,
-                "token_expires": token_expires,
-                "scopes": scopes,
-            })
-
-        if record_listens:
-            external_service_oauth_id = result.fetchone()['id']
-            connection.execute(sqlalchemy.text("""
-                INSERT INTO listens_importer
-                (external_service_oauth_id, user_id, service)
-                VALUES
-                (:external_service_oauth_id, :user_id, :service)
-                ON CONFLICT (user_id, service) DO UPDATE SET
-                    external_service_oauth_id = EXCLUDED.external_service_oauth_id,
-                    user_id = EXCLUDED.user_id,
-                    service = EXCLUDED.service,
-                    last_updated = NULL,
-                    latest_listened_at = NULL,
-                    error_message = NULL
-                """), {
-                "external_service_oauth_id": external_service_oauth_id,
-                "user_id": user_id,
-                "service": service.value
-            })
+            "external_service_oauth_id": external_service_oauth_id,
+            "user_id": user_id,
+            "service": service.value
+        })
 
 
-def delete_token(user_id: int, service: ExternalServiceType, remove_import_log: bool):
+def delete_token(connection, user_id: int, service: ExternalServiceType, remove_import_log: bool):
     """ Delete a user from the external service table.
 
     Args:
         user_id: the ListenBrainz row ID of the user
         service: the service for which the token should be deleted
         remove_import_log: whether the (user, service) combination should be removed from the listens_importer table also
     """
-    with db.engine.connect() as connection:
+    connection.execute(sqlalchemy.text("""
+        DELETE FROM external_service_oauth
+              WHERE user_id = :user_id AND service = :service
+    """), {
+        "user_id": user_id,
+        "service": service.value
+    })
+    if remove_import_log:
         connection.execute(sqlalchemy.text("""
-            DELETE FROM external_service_oauth
-                  WHERE user_id = :user_id AND service = :service
+            DELETE FROM listens_importer
+                WHERE user_id = :user_id AND service = :service
         """), {
             "user_id": user_id,
             "service": service.value
         })
-        if remove_import_log:
-            connection.execute(sqlalchemy.text("""
-                DELETE FROM listens_importer
-                    WHERE user_id = :user_id AND service = :service
-            """), {
-                "user_id": user_id,
-                "service": service.value
-            })
 
 
-def update_token(user_id: int, service: ExternalServiceType, access_token: str,
-                 refresh_token: str, expires_at: int):
+def update_token(connection, user_id: int, service: ExternalServiceType,
+                 access_token: str, refresh_token: str, expires_at: int):
     """ Update the token for user with specified LB user ID and external service.
 
     Args:
@@ -110,49 +110,49 @@ def update_token(user_id: int, service: ExternalServiceType, access_token: str,
         expires_at: the unix timestamp at which the access token expires
     """
     token_expires = utils.unix_timestamp_to_datetime(expires_at)
-    with db.engine.connect() as connection:
-        connection.execute(sqlalchemy.text("""
-            UPDATE external_service_oauth
-               SET access_token = :access_token
-                 , refresh_token = :refresh_token
-                 , token_expires = :token_expires
-                 , last_updated = now()
-             WHERE user_id = :user_id AND service = :service
-        """), {
-            "access_token": access_token,
-            "refresh_token": refresh_token,
-            "token_expires": token_expires,
-            "user_id": user_id,
-            "service": service.value
-        })
+    connection.execute(sqlalchemy.text("""
+        UPDATE external_service_oauth
+           SET access_token = :access_token
+             , refresh_token = :refresh_token
+             , token_expires = :token_expires
+        -- using clock_timestamp because value returned by now() is always fixed during a transaction
+        -- this creates issues during tests where we want to test that the last_updated field was updated
+             , last_updated = clock_timestamp()
+         WHERE user_id = :user_id AND service = :service
+    """), {
+        "access_token": access_token,
+        "refresh_token": refresh_token,
+        "token_expires": token_expires,
+        "user_id": user_id,
+        "service": service.value
+    })
 
 
-def get_token(user_id: int, service: ExternalServiceType) -> Union[dict, None]:
+def get_token(connection, user_id: int, service: ExternalServiceType) -> Union[dict, None]:
     """ Get details for user with specified user ID and service.
 
     Args:
         user_id: the ListenBrainz row ID of the user
         service: the service for which the token should be fetched
     """
-    with db.engine.connect() as connection:
-        result = connection.execute(sqlalchemy.text("""
-            SELECT user_id
-                 , "user".musicbrainz_id
-                 , "user".musicbrainz_row_id
-                 , service
-                 , access_token
-                 , refresh_token
-                 , last_updated
-                 , token_expires
-                 , scopes
-              FROM external_service_oauth
-              JOIN "user"
-                ON "user".id = external_service_oauth.user_id
-             WHERE user_id = :user_id AND service = :service
-            """), {
-                'user_id': user_id,
-                'service': service.value
-            })
-        if result.rowcount > 0:
-            return dict(result.fetchone())
-        return None
+    result = connection.execute(sqlalchemy.text("""
+        SELECT user_id
+             , "user".musicbrainz_id
+             , "user".musicbrainz_row_id
+             , service
+             , access_token
+             , refresh_token
+             , last_updated
+             , token_expires
+             , scopes
+          FROM external_service_oauth
+          JOIN "user"
+            ON "user".id = external_service_oauth.user_id
+         WHERE user_id = :user_id AND service = :service
+        """), {
+            'user_id': user_id,
+            'service': service.value
+        })
+    if result.rowcount > 0:
+        return dict(result.fetchone())
+    return None