import repository from arizona
[raven.git] / apps / gacks / gacksinvoice.py
1 import datetime
2 import fcntl
3 import os
4 import time
5
6 from gacksexcep import *
7 from gackshandle import *
8 import MySQLdb
9 import MySQLdb.cursors
10
11 STATE_PENDING = 0
12 STATE_APPLY = 1
13 STATE_APPLIED = 2
14 STATE_INVOICED = 3
15
16 KIND_CYCLE_CHARGE = 0    # besteffort
17 KIND_SLOT_CHARGE = 1     # silver
18 KIND_INVOICE=10  # invoice applied to balance
19
20 def _format_dt(dt):
21     return "%04d-%02d-%02d %02d:%02d:%02d.%06d" % (dt.year, dt.month, dt.day, dt.hour, dt.minute, dt.second, dt.microsecond)
22
23 def _format_timestamp(x):
24     x = float(x)
25     dt = datetime.datetime.fromtimestamp(x)
26     return _format_dt(dt)
27
28 class GacksCharge:
29     def __init__(self, id, object_id, kind_id, date, amount=0.0, state=STATE_PENDING, parent_id=-1):
30         self.id = int(id)
31         self.object_id = int(object_id)
32         self.kind_id = int(kind_id)
33         self.date = float(date)
34         self.amount = float(amount)
35         self.state = int(state)
36         self.last_state = int(state)    # the state before someone changed it
37         self.parent_id = int(parent_id)
38
39 class GacksInvoice:
40     def __init__(self, account_id=None, manager=None):
41         self.manager = manager
42         self.account_id = int(account_id)
43         self.charges = []
44
45     def add_charge(self, id, object_id, kind_id, date, amount, state, parent_id=-1):
46         charge = GacksCharge(id, object_id, kind_id, date, float(amount), int(state), int(parent_id))
47         self.charges.append(charge)
48
49     def get_total(self):
50         total = 0.0
51         for charge in self.charges:
52            total = total + charge.amount
53         return total
54
55     def get_account_name(self):
56         return self.manager.get_account_name(self.account_id)
57
58     def dump(self):
59         print "account:", self.get_account_name()
60         for charge in self.charges:
61            object_name = self.manager.get_object_name(charge.object_id)
62            print "  ", object_name, charge.amount
63
64     def set_state(self, state):
65         for charge in self.charges:
66             charge.state=state
67
68     def apply_to_balance(self):
69         self.manager.apply_invoice_to_balance(self)
70
71 class GacksInvoiceManager:
72     def __init__(self, tableName=None, accounts=None):
73         if tableName is None:
74              tableName = "gacksinvoice"
75
76         self.tablename = tableName
77         self.objTablename = self.tablename + "_obj"
78         self.accounts = accounts
79         self.dbname = "gacks"
80         self.address = "localhost"
81         self.user = "gacks"
82         self.passwordfile = "/root/mysql-gacks-pass.txt"
83         self.rootpasswordfile = "/root/mysql-root-pass.txt"
84         self.init_database()
85         self.open()
86
87     def get_password(self):
88         return open(self.passwordfile,"r").readline().strip()
89
90     def get_root_password(self):
91         return open(self.rootpasswordfile,"r").readline().strip()
92
93     def reset(self):
94         self.execute("DROP TABLE IF EXISTS `" + self.tablename + "`;")
95         self.execute("DROP TABLE IF EXISTS `" + self.objTablename + "`;")
96         self.create()
97
98     def open(self):
99         self.connect()
100         self.create()
101
102     def connect(self, db=None, user=None, passwd=None, address=None):
103         if not db:\r
104             db = self.dbname\r
105         if not user:\r
106             user = self.user\r
107         if not passwd:\r
108             passwd = self.get_password()\r
109         if not address:\r
110             address = self.address\r
111         self.conn = MySQLdb.connect(host=address, user=user, passwd=passwd, db=db, cursorclass=MySQLdb.cursors.DictCursor)\r
112         self.cursor = self.conn.cursor()\r
113
114     def init_database(self):
115         self.connect(db="information_schema")\r
116         self.execute("SELECT * FROM SCHEMATA WHERE SCHEMA_NAME = %s", (self.dbname))
117         row = self.cursor.fetchone();
118         if row == None:
119             self.connect(db="information_schema", user="root", passwd=self.get_root_password())
120             self.execute("CREATE DATABASE IF NOT EXISTS `" + self.dbname + "`;")
121             self.execute("GRANT ALL PRIVILEGES ON `" + self.dbname +"` . * TO %s@'%%';", (self.user))
122             self.execute("GRANT ALL PRIVILEGES ON `" + self.dbname +"` . * TO %s@'localhost';", (self.user))
123
124     def create(self):
125         querystr = "CREATE TABLE IF NOT EXISTS `" + self.tablename + "`" + """\r
126                        (`id` INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
127                         `account_id`    INT,
128                         `object_id`     INT,
129                         `parent_id`     INT,
130                         `kind_id`       INT,
131                         `date`          TIMESTAMP,
132                         `amount`        FLOAT,
133                         `state`         INT,
134                         INDEX(`account_id`),
135                         INDEX(`account_id`, `kind_id`),
136                         UNIQUE(`account_id`, `kind_id`, `object_id`, `date`));"""
137         self.execute(querystr);
138
139         querystr = "CREATE TABLE IF NOT EXISTS `" + self.objTablename + "`" + """
140                        (`id` INT NOT NULL AUTO_INCREMENT PRIMARY KEY,
141                         `name`          VARCHAR(255),
142                         UNIQUE(`name`),
143                         INDEX(`name`));"""
144         self.execute(querystr);
145
146     def execute(self, query, args=None):
147         if args is not None:
148              query = query % self.conn.literal(args)
149         #print query
150         self.cursor.execute(query)
151
152     def execute_dict(self, query, args=None):
153         self.execute(query, args)
154         return self.cursor.fetchall()
155
156     def get_object_id(self, name):
157         if type(name)==str:
158             records = self.execute_dict("SELECT id FROM `" + self.objTablename + "` WHERE name=%s", (name))
159             if len(records)==0:
160                 self.execute("INSERT INTO `" + self.objTablename+ "` (`name`) VALUES(%s);", (name))
161                 records = self.execute_dict("SELECT id FROM `" + self.objTablename + "` WHERE name=%s", (name))
162             return records[0]["id"]
163         else:
164             object_id = int(name)
165
166         return object_id
167
168     def get_account_id(self, account):
169          if type(account)==str:
170              account_id = self.accounts.get_id(account, "user")
171              if account_id == None:
172                  raise GacksAccountNotFound(account)
173          else: # int
174              account_id = account
175
176          return account_id
177
178     def get_object_name(self, id):
179         records = self.execute_dict("SELECT name FROM `" + self.objTablename + "` WHERE id=%s", (str(id)))
180         if len(records)==0:
181             return None
182         else:
183             return records[0]["name"]
184
185     def get_account_name(self, id):
186         return self.accounts.get_name(id, "user")
187
188     def get_invoice(self, account, state=STATE_PENDING):
189         account_id = self.get_account_id(account)
190         records = self.execute_dict("SELECT *,UNIX_TIMESTAMP(`date`) FROM `" + self.tablename + "` WHERE account_id=%s AND state=%s;", (str(account_id),str(state)))
191         invoice = GacksInvoice(account_id, manager=self)
192         for record in records:
193             invoice.add_charge(id=record["id"], kind_id=record["kind_id"], object_id=record["object_id"], date=record["UNIX_TIMESTAMP(`date`)"], amount=record["amount"], state=record["state"], parent_id=record["parent_id"])
194         return invoice
195
196     def add_charge(self, account, object, kind_id, date, amount, state=STATE_PENDING, parent_id=-1):
197          account_id = self.get_account_id(account)
198          object_id = self.get_object_id(object)
199
200          try:
201              self.execute("INSERT INTO `" + self.tablename + """`
202                              (account_id, object_id, kind_id, date, amount, state, parent_id)
203                              VALUES(%s, %s, %s, %s, %s, %s, %s);""",
204                                (str(account_id),
205                                 str(object_id),
206                                 str(kind_id),
207                                 _format_timestamp(date),
208                                 str(amount),
209                                 str(state),
210                                 str(parent_id)))
211          except MySQLdb.IntegrityError:
212              # it's a duplicate
213              print "duplicate:", account, kind_id, object, date
214              pass
215
216     def apply_invoice_to_balance(self, invoice):
217         # first apply everything to the database, in a nice atomic way
218         tNow = time.time()
219         self.execute("BEGIN;")
220         try:
221             # compute the adjusted balance
222             adjustBalance = 0.0
223             for charge in invoice.charges:
224                 if (charge.last_state == STATE_PENDING) and (charge.state == STATE_APPLY):
225                     adjustBalance = adjustBalance + charge.amount
226
227             self.add_charge(invoice.account_id, "invoice_applied", KIND_INVOICE, tNow, -adjustBalance, STATE_INVOICED);
228
229             parent_id = self.execute_dict("SELECT LAST_INSERT_ID();")[0]["LAST_INSERT_ID()"]
230
231             # mark all of the charges as applied, and set the parent_id
232             for charge in invoice.charges:
233                 if (charge.last_state == STATE_PENDING) and (charge.state == STATE_APPLY):
234                     self.execute("UPDATE `" + self.tablename + "` SET state=%s, parent_id=%s WHERE id=%s;",
235                                  (str(STATE_APPLIED), str(parent_id), str(charge.id)))
236
237             # rely on some cooperation with the accountmanager module to apply the invoice to the
238             self.execute("UPDATE `" + self.accounts.tablename + "` SET balance=balance-%s WHERE id=%s AND kind=%s;",
239                          (str(adjustBalance),
240                           str(invoice.account_id),
241                           "user"))
242
243             self.execute("COMMIT;")
244         except:
245             # something went wrong!
246             self.execute("ROLLBACK;")
247             raise
248
249         # now update the invoice class. The caller probably doesn't care, but just
250         # in case he does...
251         for charge in invoice.charges:
252             if (charge.last_state == STATE_PENDING) and (charge.state == STATE_APPLY):
253                 charge.state = STATE_APPLIED;
254
255     def get_accounts_with_invoices(self, state=STATE_PENDING):
256         records = self.execute_dict("SELECT DISTINCT(account_id) FROM `" + self.tablename + "` WHERE `state`=%s;", (str(state)))
257         return [rec["account_id"] for rec in records]
258
259     def apply_invoices(self):
260         accounts_ids = self.get_accounts_with_invoices()
261         for account_id in account_ids:
262             inv = self.get_invoice(account_id)
263             inv.set_state(STATE_APPLY)
264             inv.apply_to_balance()
265
266
267