Added database connection code