diff --git a/requirements.txt b/requirements.txt index d00bfac..43b710f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ flask mock +psycopg2 diff --git a/smash/__init__.py b/smash/__init__.py index d14dda5..2ef607b 100644 --- a/smash/__init__.py +++ b/smash/__init__.py @@ -1,6 +1,6 @@ import os from flask import Flask -from . import config, database, log, models +from . import config, database_postgresql, log, models log.configure_logging() @@ -36,7 +36,7 @@ if 'SECRETKEY' in conf.config: else: exit("Secret key not set.") -db = database.Database(conf.config["DBNAME"]) +db = database_postgresql.PostgreSQL(conf.config["DBNAME"]) models.init_models(db) diff --git a/smash/database_postgresql.py b/smash/database_postgresql.py new file mode 100644 index 0000000..8679afb --- /dev/null +++ b/smash/database_postgresql.py @@ -0,0 +1,66 @@ +import os +import psycopg2 +import urlparse + +class DatabasePostgreSQL(object): + """Interface to the heroku PostgreSQL database plugin. + """ + def __init__(self): + urlparse.uses_netloc.append("postgres") + self.url = urlparse.urlparse(os.environ["DATABASE_URL"]) + + + @property + def conn(self): + return psycopg2.connect( + database=self.url.path[1:], + user=self.url.username, + password=self.url.password, + host=self.url.hostname, + port=self.url.port + ) + + + @property + def cursor(self): + return self.conn.cursor() + + + def select(self, table, fields, condition=None): + cur = self.cursor + + if condition is not None: + cur.execute("SELECT {} FROM {} WHERE {}".format(fields, table, + condition)) + else: + cur.execute("SELECT {} FROM {}".format(fields, table)) + + return cur.fetchall() + + + def create_table(self, table, cols): + conn = self.conn + cur = conn.cursor() + + query = "CREATE TABLE {}({})".format(table, cols) + cur.execute(query) + conn.commit() + + + def insert(self, table, columns, values, params): + conn = self.conn + cur = conn.cursor() + + query = "INSERT INTO {}({}) VALUES ({})".format(table, columns, values) + cur.execute(query, params) + conn.commit() + return cur + + + def delete(self, table, condition): + conn = self.conn + cur = conn.cursor() + + query = "DELETE FROM {} WHERE {}".format(table, condition) + cur.execute(query) + conn.commit() diff --git a/smash/models.py b/smash/models.py index ae260d6..d8f46d3 100644 --- a/smash/models.py +++ b/smash/models.py @@ -1,5 +1,3 @@ -import sqlite3 - class Model(object): #Name of the table to be created tablename = "abstractmodel" diff --git a/smash/views.py b/smash/views.py index c7c2335..86035a0 100644 --- a/smash/views.py +++ b/smash/views.py @@ -1,6 +1,6 @@ import datetime import logging -import sqlite3 +import psycopg2 from flask import render_template, Markup, request, abort, session from smash import app, conf, db @@ -155,7 +155,7 @@ def add_new(): (tag,) ) tid = cur.lastrowid - except sqlite3.IntegrityError: + except psycopg2.IntegrityError: logger.warning("Tag {} already exists".format(tag)) if tid != -1: @@ -166,7 +166,7 @@ def add_new(): "?, ?", (tid, qid) ) - except sqlite3.Error: + except psycopg2.Error: logger.warning("Database error while inserting into tagsToQuotes") return render_template( "message.html",