prazy1208/text2sql
0
1"""2Generate synthetic data for healthcare_schema, retail_schema, and finance_schema.3Uses Faker, pandas, and SQLAlchemy. Loads DB credentials from .env.4Run: python generate_data.py5"""6 7import logging8import os9from datetime import timedelta10 11import pandas as pd12from dotenv import load_dotenv13from faker import Faker14from sqlalchemy import create_engine, text15 16# ---------------------------------------------------------------------------17# Logging18# ---------------------------------------------------------------------------19logging.basicConfig(20 level=logging.INFO,21 format="%(asctime)s [%(levelname)s] %(message)s",22 datefmt="%Y-%m-%d %H:%M:%S",23)24logger = logging.getLogger(__name__)25 26# Load .env before reading env vars27load_dotenv()28 29# ---------------------------------------------------------------------------30# Engine31# ---------------------------------------------------------------------------32 33 34def get_engine():35 """Build SQLAlchemy engine from .env (DATABASE_URL or DB_* variables)."""36 database_url = os.getenv("DATABASE_URL")37 if database_url:38 return create_engine(database_url)39 host = os.getenv("DB_HOST", "localhost")40 port = os.getenv("DB_PORT", "5432")41 user = os.getenv("DB_USER", "postgres")42 password = os.getenv("DB_PASSWORD", "")43 dbname = os.getenv("DB_NAME", "postgres")44 url = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"45 return create_engine(url)46 47 48def _sync_sequence(engine, schema: str, table: str, column: str):49 """Set the table's SERIAL sequence to the current max value."""50 stmt = text(51 f"SELECT setval(pg_get_serial_sequence('{schema}.{table}', '{column}'), "52 f"(SELECT COALESCE(MAX({column}), 1) FROM {schema}.{table}))"53 )54 with engine.connect() as conn:55 conn.execute(stmt)56 conn.commit()57 58 59# ---------------------------------------------------------------------------60# Healthcare61# ---------------------------------------------------------------------------62 63INSURANCE_TYPES = ["Medicare", "Medicaid", "Private", "Employer", "Uninsured"]64GENDERS = ["Male", "Female", "Other"]65DEPARTMENTS = ["Emergency", "Surgery", "Cardiology", "Neurology", "Pediatrics", "Oncology", "General"]66VISIT_TYPES = ["Inpatient", "Outpatient", "Emergency", "Follow-up"]67SEVERITY_LEVELS = ["Low", "Medium", "High", "Critical"]68DIAGNOSIS_CODES = ["ICD-10-A01", "ICD-10-B02", "ICD-10-C03", "ICD-10-D04", "ICD-10-E05", "ICD-10-F06", "ICD-10-G07", "ICD-10-H08"]69 70 71def generate_healthcare_data(engine):72 """Generate patients (1000), visits (2000), diagnoses (3000) in healthcare_schema."""73 fake = Faker()74 fake.seed_instance(42)75 logger.info("Generating healthcare_schema data...")76 77 # Patients: 100078 n_patients = 100079 patients = []80 for i in range(1, n_patients + 1):81 reg_date = fake.date_between(start_date="-5y", end_date="today")82 dob = fake.date_between(start_date="-90y", end_date="-18y")83 patients.append({84 "patient_id": i,85 "first_name": fake.first_name(),86 "last_name": fake.last_name(),87 "date_of_birth": dob,88 "gender": fake.random_element(GENDERS),89 "city": fake.city(),90 "state": fake.state_abbr(),91 "insurance_type": fake.random_element(INSURANCE_TYPES),92 "registration_date": reg_date,93 })94 df_patients = pd.DataFrame(patients)95 df_patients.to_sql("patients", engine, schema="healthcare_schema", if_exists="append", index=False)96 logger.info(" Inserted %d patients.", len(df_patients))97 98 # Visits: 2000 (linked to patients)99 n_visits = 2000100 visits = []101 for i in range(1, n_visits + 1):102 patient_id = fake.random_int(min=1, max=n_patients)103 admission = fake.date_between(start_date="-2y", end_date="today")104 discharge = admission + timedelta(days=fake.random_int(min=0, max=14))105 cost = round(float(fake.random_number(digits=4, fix_len=False)) + fake.random.random() * 100, 2)106 visits.append({107 "visit_id": i,108 "patient_id": patient_id,109 "admission_date": admission,110 "discharge_date": discharge,111 "department": fake.random_element(DEPARTMENTS),112 "visit_type": fake.random_element(VISIT_TYPES),113 "total_cost": cost,114 })115 df_visits = pd.DataFrame(visits)116 df_visits.to_sql("visits", engine, schema="healthcare_schema", if_exists="append", index=False)117 logger.info(" Inserted %d visits.", len(df_visits))118 119 # Diagnoses: 3000 (linked to visits)120 n_diagnoses = 3000121 diagnoses = []122 for i in range(1, n_diagnoses + 1):123 visit_id = fake.random_int(min=1, max=n_visits)124 code = fake.random_element(DIAGNOSIS_CODES) + "-" + str(fake.random_int(10, 99))125 diagnoses.append({126 "diagnosis_id": i,127 "visit_id": visit_id,128 "diagnosis_code": code,129 "diagnosis_description": fake.sentence(nb_words=6),130 "severity_level": fake.random_element(SEVERITY_LEVELS),131 })132 df_diagnoses = pd.DataFrame(diagnoses)133 df_diagnoses.to_sql("diagnoses", engine, schema="healthcare_schema", if_exists="append", index=False)134 logger.info(" Inserted %d diagnoses.", len(df_diagnoses))135 for tbl, col in [("patients", "patient_id"), ("visits", "visit_id"), ("diagnoses", "diagnosis_id")]:136 _sync_sequence(engine, "healthcare_schema", tbl, col)137 logger.info("Healthcare data generation complete.")138 139 140# ---------------------------------------------------------------------------141# Retail142# ---------------------------------------------------------------------------143 144CATEGORIES = ["Electronics", "Clothing", "Home", "Sports", "Books", "Toys", "Health", "Grocery"]145BRANDS = ["Acme", "Beta", "Gamma", "Delta", "Epsilon", "Zeta", "Omega"]146 147 148def generate_retail_data(engine):149 """Generate customers (1000), products (500), orders (3000) in retail_schema."""150 fake = Faker()151 fake.seed_instance(43)152 logger.info("Generating retail_schema data...")153 154 # Customers: 1000155 n_customers = 1000156 customers = []157 for i in range(1, n_customers + 1):158 first, last = fake.first_name(), fake.last_name()159 customers.append({160 "customer_id": i,161 "first_name": first,162 "last_name": last,163 "email": fake.ascii_safe_email(),164 "city": fake.city(),165 "state": fake.state_abbr(),166 "signup_date": fake.date_between(start_date="-3y", end_date="today"),167 })168 df_customers = pd.DataFrame(customers)169 df_customers.to_sql("customers", engine, schema="retail_schema", if_exists="append", index=False)170 logger.info(" Inserted %d customers.", len(df_customers))171 172 # Products: 500173 n_products = 500174 products = []175 for i in range(1, n_products + 1):176 price = round(fake.random.uniform(5.0, 500.0), 2)177 products.append({178 "product_id": i,179 "product_name": fake.catch_phrase().replace(".", ""),180 "category": fake.random_element(CATEGORIES),181 "brand": fake.random_element(BRANDS),182 "price": price,183 "launch_date": fake.date_between(start_date="-4y", end_date="today"),184 })185 df_products = pd.DataFrame(products)186 df_products.to_sql("products", engine, schema="retail_schema", if_exists="append", index=False)187 logger.info(" Inserted %d products.", len(df_products))188 189 # Orders: 3000 (linked to customers and products)190 n_orders = 3000191 orders = []192 for i in range(1, n_orders + 1):193 customer_id = fake.random_int(min=1, max=n_customers)194 product_id = fake.random_int(min=1, max=n_products)195 quantity = fake.random_int(min=1, max=10)196 # Price looked up would require a join; use a plausible total197 unit_price = fake.random.uniform(5.0, 500.0)198 total_amount = round(quantity * unit_price, 2)199 orders.append({200 "order_id": i,201 "customer_id": customer_id,202 "product_id": product_id,203 "order_date": fake.date_between(start_date="-2y", end_date="today"),204 "quantity": quantity,205 "total_amount": total_amount,206 })207 df_orders = pd.DataFrame(orders)208 df_orders.to_sql("orders", engine, schema="retail_schema", if_exists="append", index=False)209 logger.info(" Inserted %d orders.", len(df_orders))210 for tbl, col in [("customers", "customer_id"), ("products", "product_id"), ("orders", "order_id")]:211 _sync_sequence(engine, "retail_schema", tbl, col)212 logger.info("Retail data generation complete.")213 214 215# ---------------------------------------------------------------------------216# Finance217# ---------------------------------------------------------------------------218 219ACCOUNT_TYPES = ["savings", "checking"]220TRANSACTION_TYPES = ["debit", "credit"]221LOAN_TYPES = ["home", "auto", "personal"]222 223 224def generate_finance_data(engine):225 """Generate accounts (1000), transactions (5000), loans (800) in finance_schema."""226 fake = Faker()227 fake.seed_instance(44)228 logger.info("Generating finance_schema data...")229 230 # Accounts: 1000231 n_accounts = 1000232 accounts = []233 for i in range(1, n_accounts + 1):234 balance = round(fake.random.uniform(100.0, 100000.0), 2)235 accounts.append({236 "account_id": i,237 "customer_name": fake.name(),238 "account_type": fake.random_element(ACCOUNT_TYPES),239 "branch_city": fake.city(),240 "opening_date": fake.date_between(start_date="-10y", end_date="today"),241 "current_balance": balance,242 })243 df_accounts = pd.DataFrame(accounts)244 df_accounts.to_sql("accounts", engine, schema="finance_schema", if_exists="append", index=False)245 logger.info(" Inserted %d accounts.", len(df_accounts))246 247 # Transactions: 5000 (linked to accounts)248 n_transactions = 5000249 transactions = []250 for i in range(1, n_transactions + 1):251 account_id = fake.random_int(min=1, max=n_accounts)252 ttype = fake.random_element(TRANSACTION_TYPES)253 amount = round(fake.random.uniform(10.0, 5000.0), 2)254 transactions.append({255 "transaction_id": i,256 "account_id": account_id,257 "transaction_date": fake.date_between(start_date="-1y", end_date="today"),258 "transaction_type": ttype,259 "amount": amount,260 "description": fake.sentence(nb_words=4),261 })262 df_transactions = pd.DataFrame(transactions)263 df_transactions.to_sql("transactions", engine, schema="finance_schema", if_exists="append", index=False)264 logger.info(" Inserted %d transactions.", len(df_transactions))265 266 # Loans: 800 (linked to accounts)267 n_loans = 800268 loans = []269 for i in range(1, n_loans + 1):270 account_id = fake.random_int(min=1, max=n_accounts)271 start_date = fake.date_between(start_date="-5y", end_date="today")272 end_date = start_date + timedelta(days=fake.random_int(365, 3600))273 loan_amount = round(fake.random.uniform(5000.0, 500000.0), 2)274 rate = round(fake.random.uniform(3.0, 15.0), 2)275 loans.append({276 "loan_id": i,277 "account_id": account_id,278 "loan_type": fake.random_element(LOAN_TYPES),279 "loan_amount": loan_amount,280 "interest_rate": rate,281 "loan_start_date": start_date,282 "loan_end_date": end_date,283 })284 df_loans = pd.DataFrame(loans)285 df_loans.to_sql("loans", engine, schema="finance_schema", if_exists="append", index=False)286 logger.info(" Inserted %d loans.", len(df_loans))287 for tbl, col in [("accounts", "account_id"), ("transactions", "transaction_id"), ("loans", "loan_id")]:288 _sync_sequence(engine, "finance_schema", tbl, col)289 logger.info("Finance data generation complete.")290 291 292# ---------------------------------------------------------------------------293# Main294# ---------------------------------------------------------------------------295 296if __name__ == "__main__":297 logger.info("Starting data generation.")298 engine = get_engine()299 generate_healthcare_data(engine)300 generate_retail_data(engine)301 generate_finance_data(engine)302 logger.info("All data generation finished.")303