import unittest
import psycopg2
import pandas as pd
import re

class TestPostgresQuery(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.connection = psycopg2.connect(
            host="localhost",
            user="agent",
            password="qwerty",
            dbname="testdb",
            sslmode="disable"
        )
        cls.cursor = cls.connection.cursor()

    def test_optimized_query_present(self):
        """Check that the optimized query is present and not empty."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        # There must be at least two EXPLAIN queries: original and optimized
        condition = len(queries) >= 2 and len(queries[1]) > 0
        _dynamic_test(
            self,
            condition,
            "Optimized query is present in the solution.",
            "Optimized query is missing or incomplete. Ensure you include a complete EXPLAIN and query block for the optimized solution."
        )

    def test_all_drivers_present(self):
        """Check that all drivers are present in the optimized query result."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        self.assertTrue(len(queries) >= 2, "Optimized EXPLAIN query is missing.")
        user_query = queries[1]
        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        self.cursor.execute("SELECT driver_id, first_name, last_name FROM drivers ORDER BY driver_id")
        expected = self.cursor.fetchall()
        expected_driver_ids = set(row[0] for row in expected)
        result_driver_ids = set(result_df['driver_id'])
        condition = expected_driver_ids == result_driver_ids
        missing = expected_driver_ids - result_driver_ids
        extra = result_driver_ids - expected_driver_ids
        failure_msg = ''
        if missing:
            failure_msg += f"Missing drivers: {missing}. "
        if extra:
            failure_msg += f"Unexpected drivers: {extra}."
        _dynamic_test(
            self,
            condition,
            "All drivers are present in the optimized query result.",
            failure_msg
        )

    def test_total_rides_correct(self):
        """Check that total_rides column is correct for each driver."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        user_query = queries[1]
        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        self.cursor.execute("SELECT driver_id, COUNT(*) FROM rides GROUP BY driver_id")
        ride_counts = dict(self.cursor.fetchall())
        for _, row in result_df.iterrows():
            driver_id = row['driver_id']
            expected = ride_counts.get(driver_id, 0)
            actual = row['total_rides']
            condition = actual == expected
            _dynamic_test(
                self,
                condition,
                f"total_rides is correct for driver_id {driver_id}.",
                f"total_rides incorrect for driver_id {driver_id}: expected {expected}, got {actual}."
            )

    def test_avg_rating_correct(self):
        """Check that avg_rating is correct for each driver, 0 if no ratings."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        user_query = queries[1]
        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # Build expected averages
        self.cursor.execute("""
            SELECT r.driver_id, AVG(rt.rating)::float
            FROM rides r
            JOIN ratings rt ON r.ride_id = rt.ride_id
            GROUP BY r.driver_id
        """)
        ratings_avg = dict(self.cursor.fetchall())
        self.cursor.execute("SELECT driver_id FROM drivers")
        all_driver_ids = [row[0] for row in self.cursor.fetchall()]
        for _, row in result_df.iterrows():
            driver_id = row['driver_id']
            expected = ratings_avg.get(driver_id, 0) or 0
            actual = float(row['avg_rating'])
            # Allow rounding to 2 decimals
            condition = round(actual, 2) == round(expected, 2)
            _dynamic_test(
                self,
                condition,
                f"avg_rating is correct for driver_id {driver_id}.",
                f"avg_rating incorrect for driver_id {driver_id}: expected {expected}, got {actual}."
            )

    def test_drivers_with_no_rides_and_ratings(self):
        """Drivers with no rides or ratings should have 0 for total_rides and avg_rating."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        user_query = queries[1]
        self.cursor.execute(user_query)
        rows = self.cursor.fetchall()
        columns = [desc[0] for desc in self.cursor.description]
        result_df = pd.DataFrame(rows, columns=columns)
        # Find drivers with no rides
        self.cursor.execute("SELECT driver_id FROM drivers WHERE driver_id NOT IN (SELECT driver_id FROM rides)")
        no_ride_drivers = [row[0] for row in self.cursor.fetchall()]
        for driver_id in no_ride_drivers:
            driver_row = result_df[result_df['driver_id'] == driver_id]
            if not driver_row.empty:
                total_rides = int(driver_row.iloc[0]['total_rides'])
                avg_rating = float(driver_row.iloc[0]['avg_rating'])
                condition = total_rides == 0 and avg_rating == 0
                _dynamic_test(
                    self,
                    condition,
                    f"Driver {driver_id} with no rides/ratings has 0 for both columns.",
                    f"Driver {driver_id} with no rides/ratings does not have 0 for both columns."
                )

    def test_explain_output_present(self):
        """Check that EXPLAIN output for the optimized query is present."""
        with open("query.sql", "r") as file:
            content = file.read()
        queries = [q.strip() for q in content.split('EXPLAIN') if q.strip()]
        self.assertTrue(len(queries) >= 2, "Optimized EXPLAIN query is missing.")
        explain_query = 'EXPLAIN ' + queries[1]
        try:
            self.cursor.execute(explain_query)
            explain_rows = self.cursor.fetchall()
            condition = len(explain_rows) > 0
        except Exception as e:
            condition = False
        _dynamic_test(
            self,
            condition,
            "EXPLAIN output for the optimized query is present.",
            "EXPLAIN output for the optimized query is missing or not executable."
        )

def _dynamic_test(test_case, condition, success_message, failure_message):
    if condition:
        test_case._testMethodName = success_message
        test_case.assertTrue(True, success_message)
    else:
        test_case._testMethodName = failure_message
        test_case.fail(failure_message)

def normalize_text(text):
    text = text.lower()
    text = re.sub(r"\s{2,}", " ", text)
    text = re.sub(r"\s*([,:?])\s*", r"\1 ", text)
    return text.strip()

if __name__ == "__main__":
    unittest.main()


test_solution.py

Dive deep into the world of SQL query optimization by mastering the EXPLAIN statement and understanding how query planners work. This course blends engaging theory with hands-on, real-world tasks to help you analyze, interpret, and optimize SQL queries using EXPLAIN and related tools.

Learn the fundamentals of the EXPLAIN statement, how query planners work, and how to interpret basic query plans.

Delve into more complex query plans, including multi-table joins, subqueries, and the impact of indexes.

Apply your knowledge to complex, real-world scenarios involving large datasets, reporting queries, and performance troubleshooting.

Challenge: Optimizing a Ratings Analysis Query

解答