LeetCode SQL Problem Statement to local database

leetcode_sql_problem_statements_to_db_tables.py

import re
from typing import Dict
from io import StringIO
from dateutil import parser

from sqlalchemy import create_engine
import pandas as pd

def is_date(string):
    if string.isnumeric():
        return False
    try: 
        parser.parse(string)
        return True
    except ValueError:
        return False

def get_date_cols(line):
    words = [word.strip() for word in line.split("|") if word.strip()]
    return [ix for ix, string in enumerate(words) if is_date(string.strip())]

def get_table_name(lines, name_ix):
    return lines[name_ix].split(" ")[0].lower()

def is_table_name(line):
    return line.split(" ")[0].isalpha()

def split_text(text):
    text= re.sub(' +', ' ', text)
    return [line for line in text.splitlines() if line]

def parse_table(split_text):
    divider_cols = [0, 1, 3, len(split_text) - 1]
    
    res = []
    for ix, line in enumerate(split_text):
        if ix in divider_cols:
            continue
        else:
            res.append(line[1:-1])
    return "\n".join(res).replace(" ","")

def text_to_tables(text):
    lines = split_text(text)
    table_ixs = [ix for ix, line in enumerate(lines) if line.split(" ")[0].isalpha()]

    tables = {}
    for ix, pos in enumerate(table_ixs):
        end = table_ixs[ix + 1] if ix + 1 < len(table_ixs) else None

        date_cols = get_date_cols(lines[pos+2])
        table_name = get_table_name(lines, pos)

        tables[table_name] = pd.read_csv(
            StringIO(parse_table(lines[pos:end])), 
            sep="|",
            parse_dates=True,
        )

    return tables

def write_tables(tables: Dict[str, pd.DataFrame]) -> None:
    engine = create_engine(DATABASE_URI)

    for table_name, df in tables.items():
        df.to_sql(table_name, engine, if_exists="replace")

    tables_written = "\n".join([f'"{table}"' for table in tables.keys()])
    print(f"Succesfully wrote\n{tables_written}\ntables to database!")


text = '''
Queries table:
+------------+-------------------+----------+--------+
| query_name | result            | position | rating |
+------------+-------------------+----------+--------+
| Dog        | Golden Retriever  | 1        | 5      |
| Dog        | German Shepherd   | 2        | 5      |
| Dog        | Mule              | 200      | 1      |
| Cat        | Shirazi           | 5        | 2      |
| Cat        | Siamese           | 3        | 3      |
| Cat        | Sphynx            | 7        | 4      |
+------------+-------------------+----------+--------+
'''

# ADD YOUR DATABSE CREDENTIALS
DATABASE_URI = "postgresql://USERNAME:PASSWORD@localhost:5432/leetcode"

write_tables(text_to_tables(text))

I think LeetCode has some of the best SQL practice questions I've come across but their platform is clearly not built to make it easy to work with SQL. I use this script so I can copy and paste their demo tables into my local database. This process is trivial and working with your own db client is far more enjoyable than the leetcode website (autocomplete, auto formatting, syntax highlighting).