cleaning up code and adding connection management