{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def get_failed_jobs(start_date=None, end_date=None):\n",
    "    url = 'https://verify.tidb.ai/api/v1/failed-jobs'\n",
    "    headers = {'accept': 'application/json'}\n",
    "    params = {}\n",
    "\n",
    "    # Add optional start_date and end_date parameters if provided\n",
    "    if start_date:\n",
    "        params['start_date'] = start_date\n",
    "    if end_date:\n",
    "        params['end_date'] = end_date\n",
    "\n",
    "    try:\n",
    "        response = requests.get(url, headers=headers, params=params)\n",
    "        response.raise_for_status()  # Raises an HTTPError for bad responses\n",
    "        return response.json()\n",
    "    except requests.exceptions.RequestException as e:\n",
    "        print(f\"An error occurred: {e}\")\n",
    "        return None\n",
    "    \n",
    "def add_job_to_dataset(job_id):\n",
    "    url = f'https://verify.tidb.ai/api/v1/add-job-to-dataset/{job_id}'\n",
    "    headers = {'accept': 'application/json'}\n",
    "    data = ''\n",
    "\n",
    "    try:\n",
    "        response = requests.post(url, headers=headers, data=data)\n",
    "        if response.status_code == 200:\n",
    "            print(\"Job successfully added to the dataset.\")\n",
    "            return response.json()\n",
    "        else:\n",
    "            print(f\"Failed to add job to dataset. Status code: {response.status_code}\")\n",
    "            print(f\"Response: {response.text}\")\n",
    "            return None\n",
    "    except requests.exceptions.RequestException as e:\n",
    "        print(f\"An error occurred: {e}\")\n",
    "        return None\n",
    "    \n",
    "def get_regression_dataset(ids:list[int]=None, limit=22, offset=0):\n",
    "    url = 'https://verify.tidb.ai/api/v1/regression-dataset'\n",
    "    headers = {'accept': 'application/json'}\n",
    "    params = {}\n",
    "    if ids:\n",
    "        params['ids'] = ','.join(map(str, ids))\n",
    "    else:\n",
    "        params = {\n",
    "            'limit': limit,\n",
    "            'offset': offset\n",
    "        }\n",
    "\n",
    "    try:\n",
    "        response = requests.get(url, headers=headers, params=params)\n",
    "        response.raise_for_status()\n",
    "        return response.json()\n",
    "    except requests.exceptions.RequestException as e:\n",
    "        print(f\"An error occurred: {e}\")\n",
    "        return None\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from datetime import datetime, timedelta\n",
    "\n",
    "current_date = datetime.now().strftime('%Y-%m-%d')\n",
    "one_day_ago = (datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d')\n",
    "\n",
    "failed_jobs = get_failed_jobs(start_date=one_day_ago, end_date=current_date)\n",
    "print(f\"failed job: {len(failed_jobs['jobs'])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for job in failed_jobs['jobs']:\n",
    "    print(f\"Job ID: {job['id']}, status: {job['status']}\")\n",
    "    print(f\"Job message: {job['message']}\")\n",
    "    print(f\"Time: {job['started_at']}\")\n",
    "    print(\"Job Runs:\\n\")\n",
    "    for run in job['query_runs']:\n",
    "        detail = json.dumps(run, indent=2)\n",
    "        print(detail, \"\\n\")\n",
    "    print(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_job_to_dataset('ae12a30a-9350-463e-b1ba-08adef04a256')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "33"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = get_regression_dataset(limit=50)\n",
    "len(dataset['samples'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sample ID: 5, status: AVAILABLE\n",
      "steps: [{'explanation': 'Create a table with SHARD_ROW_ID_BITS to distribute row IDs across multiple regions.', 'sql': 'CREATE TABLE t_shard (c int) SHARD_ROW_ID_BITS = 4;'}, {'explanation': 'Create a table with AUTO_RANDOM primary key to avoid monotonic increasing IDs.', 'sql': 'CREATE TABLE t_auto_random (id BIGINT PRIMARY KEY AUTO_RANDOM, c int);'}, {'explanation': 'Pre-split a table into multiple regions to distribute the load.', 'sql': 'SPLIT TABLE t_shard BETWEEN (0) AND (10000000) REGIONS 16;'}, {'explanation': 'Create a table with a hash function to distribute writes for a timestamp column.', 'sql': 'CREATE TABLE events (id BIGINT PRIMARY KEY AUTO_RANDOM, event_time TIMESTAMP, event_hash INT GENERATED ALWAYS AS (CONV(SUBSTR(MD5(event_time), 1, 8), 16, 10)) STORED, INDEX idx_event_hash (event_hash));'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Create a table with SHARD_ROW_ID_BITS to distribute row IDs across multiple regions.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The table 't_shard' was created with SHARD_ROW_ID_BITS = 4, which distributes row IDs across multiple regions.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.316 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE t_shard (c int) SHARD_ROW_ID_BITS = 4;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-2: {\n",
      "  \"explanation\": \"Create a table with AUTO_RANDOM primary key to avoid monotonic increasing IDs.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The table 't_auto_random' was created with an AUTO_RANDOM primary key, which avoids monotonic increasing IDs.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.298 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE t_auto_random (id BIGINT PRIMARY KEY AUTO_RANDOM, c int);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-3: {\n",
      "  \"explanation\": \"Pre-split a table into multiple regions to distribute the load.\",\n",
      "  \"llm_verification\": \"The SQL query results do not satisfy the criteria outlined in the expected SQL result explanation. The query failed with an OperationalError indicating that 'SPLIT REGION' is not supported on TiDB Serverless.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"(pymysql.err.OperationalError) (20001, \\\"'SPLIT REGION' is not supported on TiDB Serverless, see https://docs.pingcap.com/tidbcloud/limited-sql-features#statements for more details.\\\")\\n[SQL: SPLIT TABLE t_shard BETWEEN (0) AND (10000000) REGIONS 16;]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"SPLIT TABLE t_shard BETWEEN (0) AND (10000000) REGIONS 16;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": \"(pymysql.err.OperationalError) (20001, \\\"'SPLIT REGION' is not supported on TiDB Serverless, see https://docs.pingcap.com/tidbcloud/limited-sql-features#statements for more details.\\\")\\n[SQL: SPLIT TABLE t_shard BETWEEN (0) AND (10000000) REGIONS 16;]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\",\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 6, status: AVAILABLE\n",
      "steps: [{'explanation': 'Enable general log.', 'sql': 'SET GLOBAL tidb_general_log = 1;'}, {'explanation': 'Create a sample table.', 'sql': 'CREATE TABLE test (id INT PRIMARY KEY);'}, {'explanation': 'Insert some sample data.', 'sql': 'INSERT INTO test VALUES (1), (2), (3);'}, {'explanation': 'Select data from the sample table.', 'sql': 'SELECT * FROM test;'}, {'explanation': 'Check the status of the general log.', 'sql': \"SHOW VARIABLES LIKE 'tidb_general_log';\"}, {'explanation': 'Disable general log.', 'sql': 'SET GLOBAL tidb_general_log = 0;'}, {'explanation': 'Drop the sample table.', 'sql': 'DROP TABLE test;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Enable general log.\",\n",
      "  \"llm_verification\": \"The SQL query did not satisfy the expected criteria because it resulted in an 'Access denied' error. The general log was not enabled as expected.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"(pymysql.err.OperationalError) (1227, 'Access denied; you need (at least one of) the RESTRICTED_VARIABLES_ADMIN privilege(s) for this operation, see https://docs.pingcap.com/tidbcloud/limited-sql-features#system-variables for more details.')\\n[SQL: SET GLOBAL tidb_general_log = 1;]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"SET GLOBAL tidb_general_log = 1;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": \"(pymysql.err.OperationalError) (1227, 'Access denied; you need (at least one of) the RESTRICTED_VARIABLES_ADMIN privilege(s) for this operation, see https://docs.pingcap.com/tidbcloud/limited-sql-features#system-variables for more details.')\\n[SQL: SET GLOBAL tidb_general_log = 1;]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\",\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 7, status: AVAILABLE\n",
      "steps: [{'explanation': \"Create a sample table named 'employees'.\", 'sql': 'CREATE TABLE employees (id INT PRIMARY KEY, name VARCHAR(50), email VARCHAR(100), hire_date DATE);'}, {'explanation': \"Add a new column named 'department'.\", 'sql': 'ALTER TABLE employees ADD COLUMN department VARCHAR(50);'}, {'explanation': \"Modify the data type of the 'name' column to VARCHAR(100).\", 'sql': 'ALTER TABLE employees MODIFY COLUMN name VARCHAR(100);'}, {'explanation': \"Rename the 'hire_date' column to 'start_date'.\", 'sql': 'ALTER TABLE employees CHANGE COLUMN hire_date start_date DATE;'}, {'explanation': \"Add a unique constraint on the 'email' column.\", 'sql': 'ALTER TABLE employees ADD UNIQUE KEY (email);'}, {'explanation': \"Modify the 'department' column to be NOT NULL with a default value of 'Unassigned'.\", 'sql': \"ALTER TABLE employees MODIFY COLUMN department VARCHAR(50) NOT NULL DEFAULT 'Unassigned';\"}, {'explanation': \"Drop the 'start_date' column.\", 'sql': 'ALTER TABLE employees DROP COLUMN start_date;'}, {'explanation': 'Create a new table with the desired structure including an AUTO_INCREMENT column.', 'sql': \"CREATE TABLE new_employees (id INT PRIMARY KEY, name VARCHAR(100), email VARCHAR(100), department VARCHAR(50) NOT NULL DEFAULT 'Unassigned', employee_number INT AUTO_INCREMENT UNIQUE, KEY (employee_number));\"}, {'explanation': 'Copy data from the old table to the new table.', 'sql': 'INSERT INTO new_employees (id, name, email, department) SELECT id, name, email, department FROM employees;'}, {'explanation': 'Rename tables to swap the old and new tables.', 'sql': 'RENAME TABLE employees TO old_employees, new_employees TO employees;'}, {'explanation': 'Drop the old table if no longer needed.', 'sql': 'DROP TABLE old_employees;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Create a sample table named 'employees'.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The table 'employees' was successfully created.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.276 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE employees (id INT PRIMARY KEY, name VARCHAR(50), email VARCHAR(100), hire_date DATE);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-2: {\n",
      "  \"explanation\": \"Add a new column named 'department'.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The column 'department' was successfully added to the 'employees' table.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.503 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees ADD COLUMN department VARCHAR(50);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-3: {\n",
      "  \"explanation\": \"Modify the data type of the 'name' column to VARCHAR(100).\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The data type of the 'name' column was successfully modified to VARCHAR(100).\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.268 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees MODIFY COLUMN name VARCHAR(100);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-4: {\n",
      "  \"explanation\": \"Rename the 'hire_date' column to 'start_date'.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The 'hire_date' column was successfully renamed to 'start_date'.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.283 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees CHANGE COLUMN hire_date start_date DATE;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-5: {\n",
      "  \"explanation\": \"Add a unique constraint on the 'email' column.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The unique constraint on the 'email' column was successfully added.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.586 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees ADD UNIQUE KEY (email);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-6: {\n",
      "  \"explanation\": \"Modify the 'department' column to be NOT NULL with a default value of 'Unassigned'.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The 'department' column was successfully modified to be NOT NULL with a default value of 'Unassigned'.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.345 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees MODIFY COLUMN department VARCHAR(50) NOT NULL DEFAULT 'Unassigned';\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-7: {\n",
      "  \"explanation\": \"Drop the 'start_date' column.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The 'start_date' column was successfully dropped from the 'employees' table.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.499 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE employees DROP COLUMN start_date;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-8: {\n",
      "  \"explanation\": \"Create a new table with the desired structure including an AUTO_INCREMENT column.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The table 'new_employees' was successfully created with the desired structure including an AUTO_INCREMENT column.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.273 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE new_employees (id INT PRIMARY KEY, name VARCHAR(100), email VARCHAR(100), department VARCHAR(50) NOT NULL DEFAULT 'Unassigned', employee_number INT AUTO_INCREMENT UNIQUE, KEY (employee_number));\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-9: {\n",
      "  \"explanation\": \"Copy data from the old table to the new table.\",\n",
      "  \"llm_verification\": \"The SQL query results do not satisfy the criteria outlined in the expected SQL result explanation. The expected result was to copy data from the old table to the new table, but the query result indicates that 0 rows were affected, meaning no data was copied.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.165 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"INSERT INTO new_employees (id, name, email, department) SELECT id, name, email, department FROM employees;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True}\n",
      "Sample ID: 9, status: AVAILABLE\n",
      "steps: [{'explanation': 'Create a table with an index.', 'sql': 'CREATE TABLE example_table (id INT PRIMARY KEY, name VARCHAR(255), age INT, INDEX idx_name (name));'}, {'explanation': 'Add an index with lower priority.', 'sql': 'ALTER TABLE example_table ADD INDEX idx_age (age);'}, {'explanation': 'Query DDL job information.', 'sql': 'ADMIN SHOW DDL JOBS;'}, {'explanation': 'Query original DDL statements for a specific job ID (replace 12345 with the actual job ID).', 'sql': 'ADMIN SHOW DDL JOB QUERIES 12345;'}, {'explanation': 'Recover index data for a specific index.', 'sql': 'ADMIN RECOVER INDEX example_table idx_name;'}, {'explanation': 'Drop the example table after the demo.', 'sql': 'DROP TABLE example_table;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Create a table with an index.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The table 'example_table' was created with an index on the 'name' column.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.274 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE example_table (id INT PRIMARY KEY, name VARCHAR(255), age INT, INDEX idx_name (name));\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-2: {\n",
      "  \"explanation\": \"Add an index with lower priority.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The index 'idx_age' was successfully added to the 'example_table' with the expected result.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.618 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ALTER TABLE example_table ADD INDEX idx_age (age);\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-3: {\n",
      "  \"explanation\": \"Query DDL job information.\",\n",
      "  \"llm_verification\": \"The SQL query results satisfy the criteria outlined in the expected SQL result explanation. The 'ADMIN SHOW DDL JOBS' query successfully returned the DDL job information, including the creation of the 'example_table' and the addition of the 'idx_age' index.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      111,\n",
      "      \"db_3a69e6e1\",\n",
      "      \"example_table\",\n",
      "      \"add index\",\n",
      "      \"public\",\n",
      "      107,\n",
      "      109,\n",
      "      0,\n",
      "      \"2024-09-10T04:06:37\",\n",
      "      \"2024-09-10T04:06:37\",\n",
      "      \"2024-09-10T04:06:37\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      110,\n",
      "      \"db_3a69e6e1\",\n",
      "      \"example_table\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      107,\n",
      "      109,\n",
      "      0,\n",
      "      \"2024-09-10T04:06:36\",\n",
      "      \"2024-09-10T04:06:36\",\n",
      "      \"2024-09-10T04:06:36\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      108,\n",
      "      \"db_3a69e6e1\",\n",
      "      \"\",\n",
      "      \"create schema\",\n",
      "      \"public\",\n",
      "      107,\n",
      "      0,\n",
      "      0,\n",
      "      \"2024-09-10T04:06:31\",\n",
      "      \"2024-09-10T04:06:31\",\n",
      "      \"2024-09-10T04:06:31\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      106,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"\",\n",
      "      \"drop schema\",\n",
      "      \"none\",\n",
      "      96,\n",
      "      0,\n",
      "      0,\n",
      "      \"2024-08-28T14:00:57\",\n",
      "      \"2024-08-28T14:00:57\",\n",
      "      \"2024-08-28T14:00:57\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      105,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"validation_jobs\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      96,\n",
      "      104,\n",
      "      0,\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      103,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"user_sessions\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      96,\n",
      "      102,\n",
      "      0,\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      101,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"db_instances\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      96,\n",
      "      100,\n",
      "      0,\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"2024-08-28T12:33:24\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      99,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"alembic_version\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      96,\n",
      "      98,\n",
      "      0,\n",
      "      \"2024-08-28T11:42:00\",\n",
      "      \"2024-08-28T11:42:00\",\n",
      "      \"2024-08-28T11:42:01\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      97,\n",
      "      \"tcs_sql_playground\",\n",
      "      \"\",\n",
      "      \"create schema\",\n",
      "      \"public\",\n",
      "      96,\n",
      "      0,\n",
      "      0,\n",
      "      \"2024-08-28T11:28:45\",\n",
      "      \"2024-08-28T11:28:45\",\n",
      "      \"2024-08-28T11:28:46\",\n",
      "      \"synced\"\n",
      "    ],\n",
      "    [\n",
      "      95,\n",
      "      \"mysql\",\n",
      "      \"auto_analyze_tasks_history\",\n",
      "      \"create table\",\n",
      "      \"public\",\n",
      "      1,\n",
      "      94,\n",
      "      0,\n",
      "      \"2024-08-28T07:36:41\",\n",
      "      \"2024-08-28T07:36:41\",\n",
      "      \"2024-08-28T07:36:41\",\n",
      "      \"synced\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ADMIN SHOW DDL JOBS;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-4: {\n",
      "  \"explanation\": \"Query original DDL statements for a specific job ID (replace 12345 with the actual job ID).\",\n",
      "  \"llm_verification\": \"The SQL query results do not satisfy the criteria outlined in the expected SQL result explanation. The expected result was to query original DDL statements for a specific job ID, but the result was an empty array.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [],\n",
      "  \"sql\": \"ADMIN SHOW DDL JOB QUERIES 12345;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 10, status: AVAILABLE\n",
      "steps: [{'explanation': 'Query to view the current GC leader in the TiDB cluster.', 'sql': \"SELECT * FROM mysql.tidb WHERE VARIABLE_NAME IN ('tidb_gc_leader_desc', 'tidb_gc_leader_uuid');\"}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Query to view the current GC leader in the TiDB cluster.\",\n",
      "  \"llm_verification\": \"The SQL query did not execute successfully due to a permission error. The user does not have the necessary permissions to access the 'mysql.tidb' table, so the expected result of viewing the current GC leader in the TiDB cluster could not be verified.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"(pymysql.err.OperationalError) (1142, \\\"SELECT command denied to user 'HCnFsDAR71bkfJd.h4xI6gnm'@'%' for table 'tidb', see https://docs.pingcap.com/tidbcloud/limited-sql-features#system-tables for more details.\\\")\\n[SQL: SELECT * FROM mysql.tidb WHERE VARIABLE_NAME IN ('tidb_gc_leader_desc', 'tidb_gc_leader_uuid');]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"SELECT * FROM mysql.tidb WHERE VARIABLE_NAME IN ('tidb_gc_leader_desc', 'tidb_gc_leader_uuid');\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": \"(pymysql.err.OperationalError) (1142, \\\"SELECT command denied to user 'HCnFsDAR71bkfJd.h4xI6gnm'@'%' for table 'tidb', see https://docs.pingcap.com/tidbcloud/limited-sql-features#system-tables for more details.\\\")\\n[SQL: SELECT * FROM mysql.tidb WHERE VARIABLE_NAME IN ('tidb_gc_leader_desc', 'tidb_gc_leader_uuid');]\\n(Background on this error at: https://sqlalche.me/e/20/e3q8)\",\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 14, status: AVAILABLE\n",
      "steps: [{'explanation': 'Check if TiFlash is available.', 'sql': 'SELECT * FROM information_schema.tiflash_replica;'}, {'explanation': 'Create a TiFlash replica for a table.', 'sql': 'ALTER TABLE your_table SET TIFLASH REPLICA 1;'}, {'explanation': 'Verify the replica status.', 'sql': \"SELECT * FROM information_schema.tiflash_replica WHERE TABLE_SCHEMA = 'your_database' AND TABLE_NAME = 'your_table';\"}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Check if TiFlash is available.\",\n",
      "  \"llm_verification\": \"The SQL query result is empty, which indicates that TiFlash is not available. The expected result should have shown some entries if TiFlash was available.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [],\n",
      "  \"sql\": \"SELECT * FROM information_schema.tiflash_replica;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 60005, status: AVAILABLE\n",
      "steps: [{'explanation': 'Connect to your TiDB instance using the MySQL client.', 'sql': 'mysql --host 127.0.0.1 --port 4000 -u root'}, {'explanation': 'Run the following SQL command to check the TiDB version.', 'sql': 'SELECT tidb_version();'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Connect to your TiDB instance using the MySQL client.\",\n",
      "  \"llm_verification\": \"The SQL query resulted in a syntax error and did not successfully connect to the TiDB instance using the MySQL client. The expected result was a successful connection.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"(pymysql.err.ProgrammingError) (1064, 'You have an error in your SQL syntax; check the manual that corresponds to your TiDB version for the right syntax to use line 1 column 5 near \\\"mysql --host 127.0.0.1 --port 4000 -u root\\\" ')\\n[SQL: mysql --host 127.0.0.1 --port 4000 -u root]\\n(Background on this error at: https://sqlalche.me/e/20/f405)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"mysql --host 127.0.0.1 --port 4000 -u root\",\n",
      "  \"sql_error_code\": 1064,\n",
      "  \"sql_error_message\": \"You have an error in your SQL syntax; check the manual that corresponds to your TiDB version for the right syntax to use line 1 column 5 near \\\"mysql --host 127.0.0.1 --port 4000 -u root\\\" \",\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 60008, status: AVAILABLE\n",
      "steps: [{'explanation': 'Check the status of ongoing DDL operations.', 'sql': 'ADMIN SHOW DDL;'}, {'explanation': \"Check the progress of a long-running 'add index' operation.\", 'sql': 'ADMIN SHOW DDL JOB QUERIES LIMIT 1;'}, {'explanation': 'Check the recent DDL job queries history.', 'sql': 'ADMIN SHOW DDL JOB QUERIES RECENT 5;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Check the status of ongoing DDL operations.\",\n",
      "  \"llm_verification\": \"The SQL query results show the status of ongoing DDL operations, which matches the expected criteria.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      56,\n",
      "      \"2c3193e7-21b8-4545-93f5-e03d35ae7877\",\n",
      "      \"localhost:4000\",\n",
      "      \"\",\n",
      "      \"2c3193e7-21b8-4545-93f5-e03d35ae7877\",\n",
      "      \"\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ADMIN SHOW DDL;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-2: {\n",
      "  \"explanation\": \"Check the progress of a long-running 'add index' operation.\",\n",
      "  \"llm_verification\": \"The SQL query results do not match the expected criteria. The expected result was to check the progress of a long-running 'add index' operation, but the actual result shows a 'CREATE DATABASE' operation.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"108\",\n",
      "      \"CREATE DATABASE IF NOT EXISTS db_6aa45ea0;\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ADMIN SHOW DDL JOB QUERIES LIMIT 1;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 90002, status: AVAILABLE\n",
      "steps: [{'explanation': 'Check for ongoing DDL operations.', 'sql': 'ADMIN SHOW DDL;'}, {'explanation': 'Check for ongoing backup tasks.', 'sql': 'SHOW BACKUPS;'}, {'explanation': 'Check for ongoing restore tasks.', 'sql': 'SHOW RESTORES;'}, {'explanation': 'Cancel a DDL job if necessary (replace <job_id> with the actual job ID).', 'sql': 'ADMIN CANCEL DDL JOBS <job_id>;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Check for ongoing DDL operations.\",\n",
      "  \"llm_verification\": \"The SQL query results indicate that there is an ongoing DDL operation, which does not satisfy the criteria outlined in the expected SQL result explanation.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      56,\n",
      "      \"2b75c5be-cfca-463d-a1f8-66414f341ab7\",\n",
      "      \"localhost:4000\",\n",
      "      \"\",\n",
      "      \"2b75c5be-cfca-463d-a1f8-66414f341ab7\",\n",
      "      \"\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"ADMIN SHOW DDL;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 90007, status: AVAILABLE\n",
      "steps: [{'explanation': 'Create a table that might benefit from scattering.', 'sql': 'CREATE TABLE large_table (id INT PRIMARY KEY, data VARCHAR(255));'}, {'explanation': 'Insert a large amount of data into the table.', 'sql': \"INSERT INTO large_table (id, data) SELECT n, CONCAT('Data-', LPAD(n, 10, '0')) FROM (SELECT (@row_number:=@row_number + 1) AS n FROM (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t1, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t2, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t3, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t4, (SELECT @row_number:=0) t0) numbers LIMIT 1000000;\"}, {'explanation': 'While the scattering is in progress, you can continue to query the table.', 'sql': 'SELECT COUNT(*) FROM large_table;'}, {'explanation': 'Check the status of the Regions for this table.', 'sql': \"SELECT * FROM information_schema.tikv_region_status WHERE table_name = 'large_table';\"}, {'explanation': 'Clear the table after the demo.', 'sql': 'DROP TABLE large_table;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Create a table that might benefit from scattering.\",\n",
      "  \"llm_verification\": \"The SQL query successfully created the table 'large_table' with the specified columns. The result matches the expected outcome of creating a table that might benefit from scattering.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 0 row affected (0.263 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"CREATE TABLE large_table (id INT PRIMARY KEY, data VARCHAR(255));\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": true,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "Step-2: {\n",
      "  \"explanation\": \"Insert a large amount of data into the table.\",\n",
      "  \"llm_verification\": \"The SQL query was expected to insert a large amount of data into the table, specifically 1,000,000 rows. However, the result indicates that only 625 rows were affected. This does not meet the expected criteria of inserting a large amount of data.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"Query OK, 625 row affected (0.172 sec)\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"INSERT INTO large_table (id, data) SELECT n, CONCAT('Data-', LPAD(n, 10, '0')) FROM (SELECT (@row_number:=@row_number + 1) AS n FROM (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t1, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t2, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t3, (SELECT 0 UNION ALL SELECT 1 UNION ALL SELECT 2 UNION ALL SELECT 3 UNION ALL SELECT 4) t4, (SELECT @row_number:=0) t0) numbers LIMIT 1000000;\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True, 'augment_sql_extraction': True}\n",
      "Sample ID: 90011, status: AVAILABLE\n",
      "steps: [{'explanation': 'Check the current max_connections setting.', 'sql': \"SHOW VARIABLES LIKE 'max_connections';\"}, {'explanation': 'Modify the max_connections setting temporarily (this change will be reset when the server restarts).', 'sql': 'SET GLOBAL max_connections = 1000;'}]\n",
      "SQL Execution History:\n",
      "\n",
      "Step-1: {\n",
      "  \"explanation\": \"Check the current max_connections setting.\",\n",
      "  \"llm_verification\": \"The SQL query result does not satisfy the expected criteria. The 'max_connections' setting should be a positive integer, but the result shows '0'.\",\n",
      "  \"message\": null,\n",
      "  \"results\": [\n",
      "    [\n",
      "      \"max_connections\",\n",
      "      \"0\"\n",
      "    ]\n",
      "  ],\n",
      "  \"sql\": \"SHOW VARIABLES LIKE 'max_connections';\",\n",
      "  \"sql_error_code\": null,\n",
      "  \"sql_error_message\": null,\n",
      "  \"success\": false,\n",
      "  \"warnings\": []\n",
      "} \n",
      "\n",
      "\n",
      "\n",
      "meta: {'augment_llm_verification': True}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "def clean_and_parse_json(meta):\n",
    "    if isinstance(meta, dict):\n",
    "        return meta\n",
    "    elif isinstance(meta, str):\n",
    "        cleaned_meta = meta.replace('\\\\n', '').strip()\n",
    "        try:\n",
    "            return json.loads(cleaned_meta)\n",
    "        except json.JSONDecodeError as e:\n",
    "            print(f\"JSONDecodeError: {e} in metadata: {cleaned_meta}\")\n",
    "            return None\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def print_sample(sample):\n",
    "    if sample['status'] == \"ARCHIVED\":\n",
    "        return\n",
    "    meta = clean_and_parse_json(sample['meta'])\n",
    "    if meta.get('augment_llm_verification', False) is False:\n",
    "        return\n",
    "    \n",
    "    print(f\"Sample ID: {sample['id']}, status: {sample['status']}\")\n",
    "    print(f\"steps: {sample['steps']}\")\n",
    "    print(f\"SQL Execution History:\\n\")\n",
    "    n = 0\n",
    "    for run in sample['query_runs']:\n",
    "        n += 1\n",
    "        detail = json.dumps(run, indent=2)\n",
    "        print(f\"Step-{n}:\", detail, \"\\n\")\n",
    "    print(\"\\n\")\n",
    "    \n",
    "    if sample['gold_answer']:\n",
    "        print(f\"gold answer:\")\n",
    "        print(f\"    no_reasonable_example: {sample['gold_answer']['sql_extraction']['no_reasonable_example']}\")\n",
    "        print(f\"    steps: {json.dumps(json.loads(sample['gold_answer']['sql_extraction']['steps']), indent=2)}\")\n",
    "    print(f\"meta: {meta}\")\n",
    "\n",
    "for sample in dataset['samples']:\n",
    "    print_sample(sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>sample_id</th>\n",
       "      <th>QA_content</th>\n",
       "      <th>no_reasonable_example</th>\n",
       "      <th>steps</th>\n",
       "      <th>metadata</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>User question: TiDB 的 view 是否能更新？\\n\\nAnswer:\\n...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"First, create a sa...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>User question: How to check table statistic up...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Check the health o...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>User question: tidb 查看表的 column flag 值\\n\\n\\nAn...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a sample ta...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5</td>\n",
       "      <td>User question: 高并发下索引写热点问题如何处理 中文\\n\\nAnswer:\\n...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a table wit...</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>6</td>\n",
       "      <td>User question: general log 怎么打开\\n\\n\\nAnswer:\\n...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>7</td>\n",
       "      <td>User question: Below are the results of my ver...</td>\n",
       "      <td>False</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_llm_verification': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>9</td>\n",
       "      <td>User question: 总结一下 tidb 6.0 到目前的所有 ddl 相关的改进\\...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a table nam...</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>10</td>\n",
       "      <td>User question: 如何查看当前 gc leader\\n\\n\\nAnswer:\\n...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>11</td>\n",
       "      <td>User question: 请编写一个存储过程，可以遍历多个表，并将查询到的数据插入到另外...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create the depende...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>13</td>\n",
       "      <td>User question: how to create clustered index ?...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a table wit...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>14</td>\n",
       "      <td>User question: how to configure storage class ...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Check if TiFlash i...</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>60002</td>\n",
       "      <td>User question: rename需要什么权限\\n\\nAnswer:\\n要执行重命名...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>60005</td>\n",
       "      <td>User question: how to check tidb version\\n\\n\\n...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Connect to your Ti...</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>60006</td>\n",
       "      <td>User question: INSERT INTO db.emp (orgId, empI...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create the sample ...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>60008</td>\n",
       "      <td>User question: how to set alert in tidb give a...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>60011</td>\n",
       "      <td>User question: How to view names of indexes of...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a sample ta...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>60013</td>\n",
       "      <td>User question: how ot upsert\\n\\nAnswer:\\nTo pe...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create the 'users'...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>90001</td>\n",
       "      <td>User question: How to grant SYSTEM_VARIABLES_A...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>90002</td>\n",
       "      <td>User question: how to upgrade tidb\\n\\n\\nAnswer...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Check for ongoing ...</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>90003</td>\n",
       "      <td>User question: any TLS usage but not mTLS\\n\\n\\...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>90004</td>\n",
       "      <td>User question: load data 报错 commit forced to q...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>90005</td>\n",
       "      <td>User question: 如何查 GC safe point\\n\\n\\nAnswer:\\...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>90006</td>\n",
       "      <td>User question: what is placement rule?\\n\\n\\nAn...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>90007</td>\n",
       "      <td>User question: curl http://{TiDBIP}:10080/tabl...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_llm_verification': True, 'augment_sq...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>90008</td>\n",
       "      <td>User question: do i need create index mysql fo...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a sample ta...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>90009</td>\n",
       "      <td>User question: How to grant SYSTEM_VARIABLES_A...</td>\n",
       "      <td>True</td>\n",
       "      <td>[]</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>90010</td>\n",
       "      <td>User question: 怎么查看表的属性\\n\\nAnswer:\\n在 TiDB 中，有...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Create a sample ta...</td>\n",
       "      <td>{'augment_sql_extraction': True}</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>90011</td>\n",
       "      <td>User question: Calculate max connection\\n\\nAns...</td>\n",
       "      <td>False</td>\n",
       "      <td>[\\n  {\\n    \"explanation\": \"Check the current ...</td>\n",
       "      <td>{'augment_llm_verification': True}</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    sample_id                                         QA_content  \\\n",
       "0           1  User question: TiDB 的 view 是否能更新？\\n\\nAnswer:\\n...   \n",
       "1           2  User question: How to check table statistic up...   \n",
       "2           3  User question: tidb 查看表的 column flag 值\\n\\n\\nAn...   \n",
       "3           5  User question: 高并发下索引写热点问题如何处理 中文\\n\\nAnswer:\\n...   \n",
       "4           6  User question: general log 怎么打开\\n\\n\\nAnswer:\\n...   \n",
       "5           7  User question: Below are the results of my ver...   \n",
       "6           9  User question: 总结一下 tidb 6.0 到目前的所有 ddl 相关的改进\\...   \n",
       "7          10  User question: 如何查看当前 gc leader\\n\\n\\nAnswer:\\n...   \n",
       "8          11  User question: 请编写一个存储过程，可以遍历多个表，并将查询到的数据插入到另外...   \n",
       "9          13  User question: how to create clustered index ?...   \n",
       "10         14  User question: how to configure storage class ...   \n",
       "11      60002  User question: rename需要什么权限\\n\\nAnswer:\\n要执行重命名...   \n",
       "12      60005  User question: how to check tidb version\\n\\n\\n...   \n",
       "13      60006  User question: INSERT INTO db.emp (orgId, empI...   \n",
       "14      60008  User question: how to set alert in tidb give a...   \n",
       "15      60011  User question: How to view names of indexes of...   \n",
       "16      60013  User question: how ot upsert\\n\\nAnswer:\\nTo pe...   \n",
       "17      90001  User question: How to grant SYSTEM_VARIABLES_A...   \n",
       "18      90002  User question: how to upgrade tidb\\n\\n\\nAnswer...   \n",
       "19      90003  User question: any TLS usage but not mTLS\\n\\n\\...   \n",
       "20      90004  User question: load data 报错 commit forced to q...   \n",
       "21      90005  User question: 如何查 GC safe point\\n\\n\\nAnswer:\\...   \n",
       "22      90006  User question: what is placement rule?\\n\\n\\nAn...   \n",
       "23      90007  User question: curl http://{TiDBIP}:10080/tabl...   \n",
       "24      90008  User question: do i need create index mysql fo...   \n",
       "25      90009  User question: How to grant SYSTEM_VARIABLES_A...   \n",
       "26      90010  User question: 怎么查看表的属性\\n\\nAnswer:\\n在 TiDB 中，有...   \n",
       "27      90011  User question: Calculate max connection\\n\\nAns...   \n",
       "\n",
       "    no_reasonable_example                                              steps  \\\n",
       "0                   False  [\\n  {\\n    \"explanation\": \"First, create a sa...   \n",
       "1                   False  [\\n  {\\n    \"explanation\": \"Check the health o...   \n",
       "2                   False  [\\n  {\\n    \"explanation\": \"Create a sample ta...   \n",
       "3                   False  [\\n  {\\n    \"explanation\": \"Create a table wit...   \n",
       "4                    True                                                 []   \n",
       "5                   False                                                 []   \n",
       "6                   False  [\\n  {\\n    \"explanation\": \"Create a table nam...   \n",
       "7                    True                                                 []   \n",
       "8                   False  [\\n  {\\n    \"explanation\": \"Create the depende...   \n",
       "9                   False  [\\n  {\\n    \"explanation\": \"Create a table wit...   \n",
       "10                  False  [\\n  {\\n    \"explanation\": \"Check if TiFlash i...   \n",
       "11                   True                                                 []   \n",
       "12                  False  [\\n  {\\n    \"explanation\": \"Connect to your Ti...   \n",
       "13                  False  [\\n  {\\n    \"explanation\": \"Create the sample ...   \n",
       "14                   True                                                 []   \n",
       "15                  False  [\\n  {\\n    \"explanation\": \"Create a sample ta...   \n",
       "16                  False  [\\n  {\\n    \"explanation\": \"Create the 'users'...   \n",
       "17                   True                                                 []   \n",
       "18                  False  [\\n  {\\n    \"explanation\": \"Check for ongoing ...   \n",
       "19                   True                                                 []   \n",
       "20                   True                                                 []   \n",
       "21                   True                                                 []   \n",
       "22                   True                                                 []   \n",
       "23                   True                                                 []   \n",
       "24                  False  [\\n  {\\n    \"explanation\": \"Create a sample ta...   \n",
       "25                   True                                                 []   \n",
       "26                  False  [\\n  {\\n    \"explanation\": \"Create a sample ta...   \n",
       "27                  False  [\\n  {\\n    \"explanation\": \"Check the current ...   \n",
       "\n",
       "                                             metadata  \n",
       "0                    {'augment_sql_extraction': True}  \n",
       "1                    {'augment_sql_extraction': True}  \n",
       "2                    {'augment_sql_extraction': True}  \n",
       "3   {'augment_llm_verification': True, 'augment_sq...  \n",
       "4   {'augment_llm_verification': True, 'augment_sq...  \n",
       "5                  {'augment_llm_verification': True}  \n",
       "6   {'augment_llm_verification': True, 'augment_sq...  \n",
       "7   {'augment_llm_verification': True, 'augment_sq...  \n",
       "8                    {'augment_sql_extraction': True}  \n",
       "9                    {'augment_sql_extraction': True}  \n",
       "10  {'augment_llm_verification': True, 'augment_sq...  \n",
       "11                   {'augment_sql_extraction': True}  \n",
       "12  {'augment_llm_verification': True, 'augment_sq...  \n",
       "13                   {'augment_sql_extraction': True}  \n",
       "14  {'augment_llm_verification': True, 'augment_sq...  \n",
       "15                   {'augment_sql_extraction': True}  \n",
       "16                   {'augment_sql_extraction': True}  \n",
       "17                   {'augment_sql_extraction': True}  \n",
       "18  {'augment_llm_verification': True, 'augment_sq...  \n",
       "19                   {'augment_sql_extraction': True}  \n",
       "20                   {'augment_sql_extraction': True}  \n",
       "21                   {'augment_sql_extraction': True}  \n",
       "22                   {'augment_sql_extraction': True}  \n",
       "23  {'augment_llm_verification': True, 'augment_sq...  \n",
       "24                   {'augment_sql_extraction': True}  \n",
       "25                   {'augment_sql_extraction': True}  \n",
       "26                   {'augment_sql_extraction': True}  \n",
       "27                 {'augment_llm_verification': True}  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import json\n",
    "\n",
    "import json\n",
    "\n",
    "def clean_and_parse_json(meta):\n",
    "    if isinstance(meta, dict):\n",
    "        return meta\n",
    "    elif isinstance(meta, str):\n",
    "        cleaned_meta = meta.replace('\\\\n', '').strip()\n",
    "        try:\n",
    "            return json.loads(cleaned_meta)\n",
    "        except json.JSONDecodeError as e:\n",
    "            print(f\"JSONDecodeError: {e} in metadata: {cleaned_meta}\")\n",
    "            return None\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "candidates = []\n",
    "for sample in dataset['samples']:\n",
    "    if sample['status'] == \"ARCHIVED\":\n",
    "        continue\n",
    "    candidates.append({\n",
    "        \"sample_id\": sample['id'],\n",
    "        \"QA_content\": sample['request']['qa_content'],\n",
    "        \"no_reasonable_example\": sample['gold_answer']['sql_extraction']['no_reasonable_example'] if sample['gold_answer'] and 'sql_extraction' in sample['gold_answer'] else False,\n",
    "        \"steps\": sample['gold_answer']['sql_extraction']['steps'] if sample['gold_answer'] and 'sql_extraction' in sample['gold_answer'] else [],\n",
    "        \"metadata\": clean_and_parse_json(sample['meta'])\n",
    "    })\n",
    "\n",
    "candidates_df = pd.DataFrame(candidates, columns=['sample_id', 'QA_content','no_reasonable_example', 'steps', 'metadata'])\n",
    "candidates_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "candidates_df.to_pickle(\"./dspy_utils/sample_candidates_pendding.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from app.repositories.regression_dataset import RegressionDatasetRepo\n",
    "from app.db import Scoped_Session\n",
    "import pandas as pd\n",
    "\n",
    "session = Scoped_Session()\n",
    "\n",
    "available_candidates_df = pd.read_pickle(\"./dspy_utils/sample_candidates_pendding.pkl\")\n",
    "available_candidates_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def label_sample(row):\n",
    "    with Scoped_Session() as session:\n",
    "        RegressionDatasetRepo.label_sample(session, row['sample_id'], {\n",
    "            \"sql_extraction\": {\n",
    "                \"no_reasonable_example\": row['no_reasonable_example'],\n",
    "                \"steps\": row['steps']\n",
    "            }\n",
    "        }, row['metadata'])\n",
    "\n",
    "for index, row in available_candidates_df.iterrows():\n",
    "    print(f\"labeling sample {row['sample_id']}\")\n",
    "    label_sample(row)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "verification",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
