Added some DB logic