dbtextdb.py 34 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199
  1. #!/usr/bin/python
  2. #
  3. # Copyright 2008 Google Inc. All Rights Reserved.
  4. """SQL-like access layer for dbtext.
  5. This module provides the glue for kamctl to interact with dbtext files
  6. using basic SQL syntax thus avoiding special case handling of dbtext.
  7. """
  8. __author__ = '[email protected] (Herman Sheremetyev)'
  9. import fcntl
  10. import os
  11. import shutil
  12. import sys
  13. import tempfile
  14. import time
  15. if 'DBTEXTDB_DEBUG' in os.environ:
  16. DEBUG = os.environ['DBTEXTDB_DEBUG']
  17. else:
  18. DEBUG = 0
  19. def Debug(msg):
  20. """Debug print method."""
  21. if DEBUG:
  22. print msg
  23. class DBText(object):
  24. """Provides connection to a dbtext database."""
  25. RESERVED_WORDS = ['SELECT', 'DELETE', 'UPDATE', 'INSERT', 'SET',
  26. 'VALUES', 'INTO', 'FROM', 'ORDER', 'BY', 'WHERE',
  27. 'COUNT', 'CONCAT', 'AND', 'AS']
  28. ALL_COMMANDS = ['SELECT', 'DELETE', 'UPDATE', 'INSERT']
  29. WHERE_COMMANDS = ['SELECT', 'DELETE', 'UPDATE']
  30. def __init__(self, location):
  31. self.location = location # location of dbtext tables
  32. self.tokens = [] # query broken up into tokens
  33. self.conditions = {} # args to the WHERE clause
  34. self.columns = [] # columns requested by SELECT
  35. self.table = '' # name of the table being queried
  36. self.header = {} # table header
  37. self.orig_data = [] # original table data used to diff after updates
  38. self.data = [] # table data as a list of dicts
  39. self.count = False # where or not using COUNT()
  40. self.aliases = {} # column aliases (SELECT AS)
  41. self.targets = {} # target columns-value pairs for INSERT/UPDATE
  42. self.args = '' # query arguments preceeding the ;
  43. self.command = '' # which command are we executing
  44. self.strings = [] # list of string literals parsed from the query
  45. self.parens = [] # list of parentheses parsed from the query
  46. self._str_placeholder = '__DBTEXTDB_PARSED_OUT_STRING__'
  47. self._paren_placeholder = '__DBTEXTDB_PARSED_OUT_PARENS__'
  48. if not os.path.isdir(location):
  49. raise ParseError(location + ' is not a directory')
  50. def _ParseOrderBy(self):
  51. """Parse out the column name to be used for ordering the dataset.
  52. Raises:
  53. ParseError: Invalid ORDER BY clause
  54. """
  55. self.order_by = ''
  56. if 'ORDER' in self.tokens:
  57. order_index = self.tokens.index('ORDER')
  58. if order_index != len(self.tokens) - 3:
  59. raise ParseError('ORDER must be followed with BY and column name')
  60. if self.tokens[order_index + 1] != 'BY':
  61. raise ParseError('ORDER must be followed with BY')
  62. self.order_by = self.tokens[order_index + 2]
  63. # strip off the order by stuff
  64. self.tokens.pop() # column name
  65. self.tokens.pop() # BY
  66. self.tokens.pop() # ORDER
  67. elif 'BY' in self.tokens:
  68. raise ParseError('BY must be preceeded by ORDER')
  69. Debug('Order by: ' + self.order_by)
  70. def _ParseConditions(self):
  71. """Parse out WHERE clause.
  72. Take everything after the WHERE keyword and convert it to a dict of
  73. name value pairs corresponding to the columns and their values that
  74. should be matched.
  75. Raises:
  76. ParseError: Invalid WHERE clause
  77. NotSupportedError: Unsupported syntax
  78. """
  79. self.conditions = {}
  80. Debug('self.tokens = %s' % self.tokens)
  81. if 'WHERE' not in self.tokens:
  82. return
  83. if self.command not in self.WHERE_COMMANDS:
  84. raise ParseError(self.command + ' cannot have a WHERE clause')
  85. if 'OR' in self.tokens:
  86. raise NotSupportedError('WHERE clause does not support OR operator')
  87. where_clause = self.tokens[self.tokens.index('WHERE') + 1:]
  88. self.conditions = self._ParsePairs(' '.join(where_clause), 'AND')
  89. for cond in self.conditions:
  90. self.conditions[cond] = self._EscapeChars(self.conditions[cond])
  91. Debug('Conditions are [%s]' % self.conditions)
  92. # pop off where clause
  93. a = self.tokens.pop()
  94. while a != 'WHERE':
  95. a = self.tokens.pop()
  96. Debug('self.tokens: %s' % self.tokens)
  97. def _ParseColumns(self):
  98. """Parse out the columns that need to be selected.
  99. Raises:
  100. ParseError: Invalid SELECT syntax
  101. """
  102. self.columns = []
  103. self.count = False
  104. self.aliases = {}
  105. col_end = 0
  106. # this is only valid for SELECT
  107. if self.command != 'SELECT':
  108. return
  109. if 'FROM' not in self.tokens:
  110. raise ParseError('SELECT must be followed by FROM')
  111. col_end = self.tokens.index('FROM')
  112. if not col_end: # col_end == 0
  113. raise ParseError('SELECT must be followed by column name[s]')
  114. cols_str = ' '.join(self.tokens[0:col_end])
  115. # check if there is a function modifier on the columns
  116. if self.tokens[0] == 'COUNT':
  117. self.count = True
  118. if col_end == 1:
  119. raise ParseError('COUNT must be followed by column name[s]')
  120. if not self.tokens[1].startswith(self._paren_placeholder):
  121. raise ParseError('COUNT must be followed by ()')
  122. cols_str = self._ReplaceParens(self.tokens[1])
  123. cols = cols_str.split(',')
  124. for col in cols:
  125. if not col.strip():
  126. raise ParseError('Extra comma in columns')
  127. col_split = col.split()
  128. if col_split[0] == 'CONCAT':
  129. # found a concat statement, do the same overall steps for those cols
  130. self._ParseColumnsConcatHelper(col_split)
  131. else:
  132. col_split = col.split()
  133. if len(col_split) > 2 and col_split[1] != 'AS':
  134. raise ParseError('multiple columns must be separated by a comma')
  135. elif len(col_split) == 3:
  136. if col_split[1] != 'AS':
  137. raise ParseError('Invalid column alias, use AS')
  138. my_key = self._ReplaceStringLiterals(col_split[2], quotes=True)
  139. my_val = self._ReplaceStringLiterals(col_split[0], quotes=True)
  140. self.aliases[my_key] = [my_val]
  141. self.columns.append(my_key)
  142. elif len(col_split) > 3:
  143. raise ParseError('multiple columns must be separated by a comma')
  144. elif len(col_split) == 2: # alias
  145. my_key = self._ReplaceStringLiterals(col_split[1], quotes=True)
  146. my_val = self._ReplaceStringLiterals(col_split[0], quotes=True)
  147. self.aliases[my_key] = [my_val]
  148. self.columns.append(my_key)
  149. else:
  150. col = self._ReplaceStringLiterals(col, quotes=True).strip()
  151. if not col: # col == ''
  152. raise ParseError('empty column name not allowed')
  153. self.columns.append(col)
  154. # pop off all the columns related junk
  155. self.tokens = self.tokens[col_end + 1:]
  156. Debug('Columns: %s' % self.columns)
  157. Debug('Aliases: %s' % self.aliases)
  158. Debug('self.tokens: %s' % self.tokens)
  159. def _ParseColumnsConcatHelper(self, col_split):
  160. """Handles the columns being CONCAT'd together.
  161. Args:
  162. col_split: ['column', 'column']
  163. Raises:
  164. ParseError: invalid CONCAT()
  165. """
  166. concat_placeholder = '_'
  167. split_len = len(col_split)
  168. if split_len == 1:
  169. raise ParseError('CONCAT() must be followed by column name[s]')
  170. if not col_split[1].startswith(self._paren_placeholder):
  171. raise ParseError('CONCAT must be followed by ()')
  172. if split_len > 2:
  173. if split_len == 4 and col_split[2] != 'AS':
  174. raise ParseError('CONCAT() must be followed by an AS clause')
  175. if split_len > 5:
  176. raise ParseError('CONCAT() AS clause takes exactly 1 arg. '
  177. 'Extra args: [%s]' % (col_split[4:]))
  178. else:
  179. concat_placeholder = self._ReplaceStringLiterals(col_split[-1],
  180. quotes=True)
  181. # make sure this place hodler is unique
  182. while concat_placeholder in self.aliases:
  183. concat_placeholder += '_'
  184. concat_cols_str = self._ReplaceParens(col_split[1])
  185. concat_cols = concat_cols_str.split(',')
  186. concat_col_list = []
  187. for concat_col in concat_cols:
  188. if ' ' in concat_col.strip():
  189. raise ParseError('multiple columns must be separated by a'
  190. ' comma inside CONCAT()')
  191. concat_col = self._ReplaceStringLiterals(concat_col, quotes=True).strip()
  192. if not concat_col:
  193. raise ParseError('Attempting to CONCAT empty set')
  194. concat_col_list.append(concat_col)
  195. self.aliases[concat_placeholder] = concat_col_list
  196. self.columns.append(concat_placeholder)
  197. def _ParseTable(self):
  198. """Parse out the table name (multiple table names not supported).
  199. Raises:
  200. ParseError: Unable to parse table name
  201. """
  202. table_name = ''
  203. if (not self.tokens or # len == 0
  204. (self.tokens[0] in self.RESERVED_WORDS and
  205. self.tokens[0] not in ['FROM', 'INTO'])):
  206. raise ParseError('Missing table name')
  207. # SELECT
  208. if self.command == 'SELECT':
  209. table_name = self.tokens.pop(0)
  210. # INSERT
  211. elif self.command == 'INSERT':
  212. table_name = self.tokens.pop(0)
  213. if table_name == 'INTO':
  214. table_name = self.tokens.pop(0)
  215. # DELETE
  216. elif self.command == 'DELETE':
  217. if self.tokens[0] != 'FROM':
  218. raise ParseError('DELETE command must be followed by FROM')
  219. self.tokens.pop(0) # FROM
  220. table_name = self.tokens.pop(0)
  221. # UPDATE
  222. elif self.command == 'UPDATE':
  223. table_name = self.tokens.pop(0)
  224. if not self.table:
  225. self.table = table_name
  226. else: # multiple queries detected, make sure they're against same table
  227. if self.table != table_name:
  228. raise ParseError('Table changed between queries! %s -> %s' %
  229. (self.table, table_name))
  230. Debug('Table is [%s]' % self.table)
  231. Debug('self.tokens is %s' % self.tokens)
  232. def _ParseTargets(self):
  233. """Parse out name value pairs of columns and their values.
  234. Raises:
  235. ParseError: Unable to parse targets
  236. """
  237. self.targets = {}
  238. # UPDATE
  239. if self.command == 'UPDATE':
  240. if self.tokens.pop(0) != 'SET':
  241. raise ParseError('UPDATE command must be followed by SET')
  242. self.targets = self._ParsePairs(' '.join(self.tokens), ',')
  243. # INSERT
  244. if self.command == 'INSERT':
  245. if self.tokens[0] == 'SET':
  246. self.targets = self._ParsePairs(' '.join(self.tokens[1:]), ',')
  247. elif len(self.tokens) == 3 and self.tokens[1] == 'VALUES':
  248. if not self.tokens[0].startswith(self._paren_placeholder):
  249. raise ParseError('INSERT column names must be inside parens')
  250. if not self.tokens[2].startswith(self._paren_placeholder):
  251. raise ParseError('INSERT values must be inside parens')
  252. cols = self._ReplaceParens(self.tokens[0]).split(',')
  253. vals = self._ReplaceParens(self.tokens[2]).split(',')
  254. if len(cols) != len(vals):
  255. raise ParseError('INSERT column and value numbers must match')
  256. if not cols: # len == 0
  257. raise ParseError('INSERT column number must be greater than 0')
  258. i = 0
  259. while i < len(cols):
  260. val = vals[i].strip()
  261. if not val: # val == ''
  262. raise ParseError('INSERT values cannot be empty')
  263. if ' ' in val:
  264. raise ParseError('INSERT values must be comma separated')
  265. self.targets[cols[i].strip()] = self._ReplaceStringLiterals(val)
  266. i += 1
  267. else:
  268. raise ParseError('Unable to parse INSERT targets')
  269. for target in self.targets:
  270. self.targets[target] = self._EscapeChars(self.targets[target])
  271. Debug('Targets are [%s]' % self.targets)
  272. def _EscapeChars(self, value):
  273. """Escape necessary chars before inserting into dbtext.
  274. Args:
  275. value: 'string'
  276. Returns:
  277. escaped: 'string' with chars escaped appropriately
  278. """
  279. # test that the value is string, if not return it as is
  280. try:
  281. value.find('a')
  282. except:
  283. return value
  284. escaped = value
  285. escaped = escaped.replace('\\', '\\\\').replace('\0', '\\0')
  286. escaped = escaped.replace(':', '\\:').replace('\n', '\\n')
  287. escaped = escaped.replace('\r', '\\r').replace('\t', '\\t')
  288. return escaped
  289. def _UnEscapeChars(self, value):
  290. """Un-escape necessary chars before returning to user.
  291. Args:
  292. value: 'string'
  293. Returns:
  294. escaped: 'string' with chars escaped appropriately
  295. """
  296. # test that the value is string, if not return it as is
  297. try:
  298. value.find('a')
  299. except:
  300. return value
  301. escaped = value
  302. escaped = escaped.replace('\\:', ':').replace('\\n', '\n')
  303. escaped = escaped.replace('\\r', '\r').replace('\\t', '\t')
  304. escaped = escaped.replace('\\0', '\0').replace('\\\\', '\\')
  305. return escaped
  306. def Execute(self, query, writethru=True):
  307. """Parse and execute the query.
  308. Args:
  309. query: e.g. 'select * from table;'
  310. writethru: bool
  311. Returns:
  312. dataset: [{col: val, col: val}, {col: val}, {col: val}]
  313. Raises:
  314. ExecuteError: unable to execute query
  315. """
  316. # parse the query
  317. self.ParseQuery(query)
  318. # get lock and execute the query
  319. self.OpenTable()
  320. Debug('Running ' + self.command)
  321. dataset = []
  322. if self.command == 'SELECT':
  323. dataset = self._RunSelect()
  324. elif self.command == 'UPDATE':
  325. dataset = self._RunUpdate()
  326. elif self.command == 'INSERT':
  327. dataset = self._RunInsert()
  328. elif self.command == 'DELETE':
  329. dataset = self._RunDelete()
  330. if self.command != 'SELECT' and writethru:
  331. self.WriteTempTable()
  332. self.MoveTableIntoPlace()
  333. Debug(dataset)
  334. return dataset
  335. def CleanUp(self):
  336. """Reset the internal variables (for multiple queries)."""
  337. self.tokens = [] # query broken up into tokens
  338. self.conditions = {} # args to the WHERE clause
  339. self.columns = [] # columns requested by SELECT
  340. self.table = '' # name of the table being queried
  341. self.header = {} # table header
  342. self.orig_data = [] # original table data used to diff after updates
  343. self.data = [] # table data as a list of dicts
  344. self.count = False # where or not using COUNT()
  345. self.aliases = {} # column aliases (SELECT AS)
  346. self.targets = {} # target columns-value pairs for INSERT/UPDATE
  347. self.args = '' # query arguments preceeding the ;
  348. self.command = '' # which command are we executing
  349. self.strings = [] # list of string literals parsed from the query
  350. self.parens = [] # list of parentheses parsed from the query
  351. def ParseQuery(self, query):
  352. """External wrapper for the query parsing routines.
  353. Args:
  354. query: string
  355. Raises:
  356. ParseError: Unable to parse query
  357. """
  358. self.args = query.split(';')[0]
  359. self._Tokenize()
  360. self._ParseCommand()
  361. self._ParseOrderBy()
  362. self._ParseConditions()
  363. self._ParseColumns()
  364. self._ParseTable()
  365. self._ParseTargets()
  366. def _ParseCommand(self):
  367. """Determine the command: SELECT, UPDATE, DELETE or INSERT.
  368. Raises:
  369. ParseError: unable to parse command
  370. """
  371. self.command = self.tokens[0]
  372. # Check that command is valid
  373. if self.command not in self.ALL_COMMANDS:
  374. raise ParseError('Unsupported command: ' + self.command)
  375. self.tokens.pop(0)
  376. Debug('Command is: %s' % self.command)
  377. Debug('self.tokens: %s' % self.tokens)
  378. def _Tokenize(self):
  379. """Turn the string query into a list of tokens.
  380. Split on '(', ')', ' ', ';', '=' and ','.
  381. In addition capitalize any SQL keywords found.
  382. """
  383. # horrible hack to handle now()
  384. time_now = '%s' % int(time.time())
  385. time_now = time_now[0:-2] + '00' # round off the seconds for unittesting
  386. while 'now()' in self.args.lower():
  387. start = self.args.lower().find('now()')
  388. self.args = ('%s%s%s' % (self.args[0:start], time_now,
  389. self.args[start + 5:]))
  390. # pad token separators with spaces
  391. pad = self.args.replace('(', ' ( ').replace(')', ' ) ')
  392. pad = pad.replace(',', ' , ').replace(';', ' ; ').replace('=', ' = ')
  393. self.args = pad
  394. # parse out all the blocks (string literals and parens)
  395. self._ParseOutBlocks()
  396. # split remaining into tokens
  397. self.tokens = self.args.split()
  398. # now capitalize
  399. i = 0
  400. while i < len(self.tokens):
  401. if self.tokens[i].upper() in self.RESERVED_WORDS:
  402. self.tokens[i] = self.tokens[i].upper()
  403. i += 1
  404. Debug('Tokens: %s' % self.tokens)
  405. def _ParseOutBlocks(self):
  406. """Parse out string literals and parenthesized values."""
  407. self.strings = []
  408. self.parens = []
  409. # set str placeholder to a value that's not present in the string
  410. while self._str_placeholder in self.args:
  411. self._str_placeholder = '%s_' % self._str_placeholder
  412. # set paren placeholder to a value that's not present in the string
  413. while self._paren_placeholder in self.args:
  414. self._paren_placeholder = '%s_' % self._paren_placeholder
  415. self.strings = self._ParseOutHelper(self._str_placeholder, ["'", '"'],
  416. 'quotes')
  417. self.parens = self._ParseOutHelper(self._paren_placeholder, ['(', ')'],
  418. 'parens')
  419. Debug('Strings: %s' % self.strings)
  420. Debug('Parens: %s' % self.parens)
  421. def _ParseOutHelper(self, placeholder, delims, mode):
  422. """Replace all text within delims with placeholders.
  423. Args:
  424. placeholder: string
  425. delims: list of strings
  426. mode: string
  427. 'parens': if there are 2 delims treat the first as opening
  428. and second as closing, such as with ( and )
  429. 'quotes': treat each delim as either opening or
  430. closing and require the same one to terminate the block,
  431. such as with ' and "
  432. Returns:
  433. list: [value1, value2, ...]
  434. Raises:
  435. ParseError: unable to parse out delims
  436. ExecuteError: Invalid usage
  437. """
  438. if mode not in ['quotes', 'parens']:
  439. raise ExecuteError('_ParseOutHelper: invalid mode ' + mode)
  440. if mode == 'parens' and len(delims) != 2:
  441. raise ExecuteError('_ParseOutHelper: delims must have 2 values '
  442. 'in "parens" mode')
  443. values = []
  444. started = 0
  445. new_args = ''
  446. string = ''
  447. my_id = 0
  448. delim = ''
  449. for c in self.args:
  450. if c in delims:
  451. if not started:
  452. if mode == 'parens' and c != delims[0]:
  453. raise ParseError('Found closing delimeter %s before '
  454. 'corresponding %s' % (c, delims[0]))
  455. started += 1
  456. delim = c
  457. else:
  458. if ((mode == 'parens' and c == delim) or
  459. (mode == 'quotes' and c != delim)):
  460. string = '%s%s' % (string, c)
  461. continue # wait for matching delim
  462. started -= 1
  463. if not started:
  464. values.append(string)
  465. new_args = '%s %s' % (new_args, '%s%d' % (placeholder, my_id))
  466. my_id += 1
  467. string = ''
  468. else:
  469. if not started:
  470. new_args = '%s%s' % (new_args, c)
  471. else:
  472. string = '%s%s' % (string, c)
  473. if started:
  474. if mode == 'parens':
  475. waiting_for = delims[1]
  476. else:
  477. waiting_for = delim
  478. raise ParseError('Unterminated block, waiting for ' + waiting_for)
  479. self.args = new_args
  480. Debug('Values: %s' % values)
  481. return values
  482. def _ReplaceStringLiterals(self, s, quotes=False):
  483. """Replaces string placeholders with real values.
  484. If quotes is set to True surround the returned value with single quotes
  485. Args:
  486. s: string
  487. quotes: bool
  488. Returns:
  489. s: string
  490. """
  491. if s.strip().startswith(self._str_placeholder):
  492. str_index = int(s.split(self._str_placeholder)[1])
  493. s = self.strings[str_index]
  494. if quotes:
  495. s = "'" + s + "'"
  496. return s
  497. def _ReplaceParens(self, s):
  498. """Replaces paren placeholders with real values.
  499. Args:
  500. s: string
  501. Returns:
  502. s: string
  503. """
  504. if s.strip().startswith(self._paren_placeholder):
  505. str_index = int(s.split(self._paren_placeholder)[1])
  506. s = self.parens[str_index].strip()
  507. return s
  508. def _RunDelete(self):
  509. """Run the DELETE command.
  510. Go through the rows in self.data matching them
  511. against the conditions, if they fit delete the row leaving a placeholder
  512. value (in order to keep the iteration process sane). Afterward clean up
  513. any empty values.
  514. Returns:
  515. dataset: [number of affected rows]
  516. """
  517. i = 0
  518. length = len(self.data)
  519. affected = 0
  520. while i < length:
  521. if self._MatchRow(self.data[i]):
  522. self.data[i] = None
  523. affected += 1
  524. i += 1
  525. # clean out the placeholders
  526. while None in self.data:
  527. self.data.remove(None)
  528. return [affected]
  529. def _RunUpdate(self):
  530. """Run the UPDATE command.
  531. Find the matching rows and update based on self.targets
  532. Returns:
  533. affected: [int]
  534. Raises:
  535. ExecuteError: failed to run UPDATE
  536. """
  537. i = 0
  538. length = len(self.data)
  539. affected = 0
  540. while i < length:
  541. if self._MatchRow(self.data[i]):
  542. for target in self.targets:
  543. if target not in self.header:
  544. raise ExecuteError(target + ' is an invalid column name')
  545. if self.header[target]['auto']:
  546. raise ExecuteError(target + ' is type auto and connot be updated')
  547. self.data[i][target] = self._TypeCheck(self.targets[target], target)
  548. affected += 1
  549. i += 1
  550. return [affected]
  551. def _RunInsert(self):
  552. """Run the INSERT command.
  553. Build up the row based on self.targets and table defaults, then append to
  554. self.data
  555. Returns:
  556. affected: [int]
  557. Raises:
  558. ExecuteError: failed to run INSERT
  559. """
  560. new_row = {}
  561. cols = self._SortHeaderColumns()
  562. for col in cols:
  563. if col in self.targets:
  564. if self.header[col]['auto']:
  565. raise ExecuteError(col + ' is type auto: cannot be modified')
  566. new_row[col] = self.targets[col]
  567. elif self.header[col]['null']:
  568. new_row[col] = ''
  569. elif self.header[col]['auto']:
  570. new_row[col] = self._GetNextAuto(col)
  571. else:
  572. raise ExecuteError(col + ' cannot be empty or null')
  573. self.data.append(new_row)
  574. return [1]
  575. def _GetNextAuto(self, col):
  576. """Figure out the next value for col based on existing values.
  577. Scan all the current values and return the highest one + 1.
  578. Args:
  579. col: string
  580. Returns:
  581. next: int
  582. Raises:
  583. ExecuteError: Failed to get auto inc
  584. """
  585. highest = 0
  586. seen = []
  587. for row in self.data:
  588. if row[col] > highest:
  589. highest = row[col]
  590. if row[col] not in seen:
  591. seen.append(row[col])
  592. else:
  593. raise ExecuteError('duplicate value %s in %s' % (row[col], col))
  594. return highest + 1
  595. def _RunSelect(self):
  596. """Run the SELECT command.
  597. Returns:
  598. dataset: []
  599. Raises:
  600. ExecuteError: failed to run SELECT
  601. """
  602. dataset = []
  603. if ['*'] == self.columns:
  604. self.columns = self._SortHeaderColumns()
  605. for row in self.data:
  606. if self._MatchRow(row):
  607. match = []
  608. for col in self.columns:
  609. if col in self.aliases:
  610. concat = ''
  611. for concat_col in self.aliases[col]:
  612. if concat_col.startswith("'") and concat_col.endswith("'"):
  613. concat += concat_col.strip("'")
  614. elif concat_col not in self.header.keys():
  615. raise ExecuteError('Table %s does not have a column %s' %
  616. (self.table, concat_col))
  617. else:
  618. concat = '%s%s' % (concat, row[concat_col])
  619. if not concat.strip():
  620. raise ExecuteError('Empty CONCAT statement')
  621. my_match = concat
  622. elif col.startswith("'") and col.endswith("'"):
  623. my_match = col.strip("'")
  624. elif col not in self.header.keys():
  625. raise ExecuteError('Table %s does not have a column %s' %
  626. (self.table, col))
  627. else:
  628. my_match = row[col]
  629. match.append(self._UnEscapeChars(my_match))
  630. dataset.append(match)
  631. if self.count:
  632. Debug('Dataset: %s' % dataset)
  633. dataset = [len(dataset)]
  634. if self.order_by:
  635. if self.order_by not in self.header.keys():
  636. raise ExecuteError('Unknown column %s in ORDER BY clause' %
  637. self.order_by)
  638. pos = self._PositionByCol(self.order_by)
  639. dataset = self._SortMatrixByCol(dataset, pos)
  640. return dataset
  641. def _SortMatrixByCol(self, dataset, pos):
  642. """Sorts the matrix (array or arrays) based on a given column value.
  643. That is, if given matrix that looks like:
  644. [[1, 2, 3], [6, 5, 4], [3, 2, 1]]
  645. given pos = 0 produce:
  646. [[1, 2, 3], [3, 2, 1], [6, 5, 4]]
  647. given pos = 1 produce:
  648. [[1, 2, 3], [3, 2, 1], [6, 5, 4]]
  649. given pos = 2 produce:
  650. [[3, 2, 1], [1, 2, 3], [6, 5, 4]]
  651. Works for both integer and string values of column.
  652. Args:
  653. dataset: [[], [], ...]
  654. pos: int
  655. Returns:
  656. sorted: [[], [], ...]
  657. """
  658. # prepend value in pos to the beginning of every row
  659. i = 0
  660. while i < len(dataset):
  661. dataset[i].insert(0, dataset[i][pos])
  662. i += 1
  663. # sort the matrix, which is done on the row we just prepended
  664. dataset.sort()
  665. # strip away the first value
  666. i = 0
  667. while i < len(dataset):
  668. dataset[i].pop(0)
  669. i += 1
  670. return dataset
  671. def _MatchRow(self, row):
  672. """Matches the row against self.conditions.
  673. Args:
  674. row: ['val', 'val']
  675. Returns:
  676. Bool
  677. """
  678. match = True
  679. # when there are no conditions we match everything
  680. if not self.conditions:
  681. return match
  682. for condition in self.conditions:
  683. cond_val = self.conditions[condition]
  684. if condition not in self.header.keys():
  685. match = False
  686. break
  687. else:
  688. if cond_val != row[condition]:
  689. match = False
  690. break
  691. return match
  692. def _ProcessHeader(self):
  693. """Parse out the header information.
  694. Returns:
  695. {col_name: {'type': string, 'null': string, 'auto': string, 'pos': int}}
  696. """
  697. header = self.fd.readline().strip()
  698. cols = {}
  699. pos = 0
  700. for col in header.split():
  701. col_name = col.split('(')[0]
  702. col_type = col.split('(')[1].split(')')[0].split(',')[0]
  703. col_null = False
  704. col_auto = False
  705. if ',' in col.split('(')[1].split(')')[0]:
  706. if col.split('(')[1].split(')')[0].split(',')[1].lower() == 'null':
  707. col_null = True
  708. if col.split('(')[1].split(')')[0].split(',')[1].lower() == 'auto':
  709. col_auto = True
  710. cols[col_name] = {}
  711. cols[col_name]['type'] = col_type
  712. cols[col_name]['null'] = col_null
  713. cols[col_name]['auto'] = col_auto
  714. cols[col_name]['pos'] = pos
  715. pos += 1
  716. return cols
  717. def _GetData(self):
  718. """Reads table data into memory as a list of dicts keyed on column names.
  719. Returns:
  720. data: [{row}, {row}, ...]
  721. Raises:
  722. ExecuteError: failed to get data
  723. """
  724. data = []
  725. row_num = 0
  726. for row in self.fd:
  727. row = row.rstrip('\n')
  728. row_dict = {}
  729. i = 0
  730. field_start = 0
  731. field_num = 0
  732. while i < len(row):
  733. if row[i] == ':':
  734. # the following block is executed again after the while is done
  735. val = row[field_start:i]
  736. col = self._ColByPosition(field_num)
  737. val = self._TypeCheck(val, col)
  738. row_dict[col] = val
  739. field_start = i + 1 # skip the colon itself
  740. field_num += 1
  741. if row[i] == '\\':
  742. i += 2 # skip the next char since it's escaped
  743. else:
  744. i += 1
  745. # handle the last field since we won't hit a : at the end
  746. # sucks to duplicate the code outside the loop but I can't think
  747. # of a better way :(
  748. val = row[field_start:i]
  749. col = self._ColByPosition(field_num)
  750. val = self._TypeCheck(val, col)
  751. row_dict[col] = val
  752. # verify that all columns were created
  753. for col in self.header:
  754. if col not in row_dict:
  755. raise ExecuteError('%s is missing from row %d in %s' %
  756. (col, row_num, self.table))
  757. row_num += 1
  758. data.append(row_dict)
  759. return data
  760. def _TypeCheck(self, val, col):
  761. """Verify type of val based on the header.
  762. Make sure the value is returned in quotes if it's a string
  763. and as '' when it's empty and Null
  764. Args:
  765. val: string
  766. col: string
  767. Returns:
  768. val: string
  769. Raises:
  770. ExecuteError: invalid value or column
  771. """
  772. if not val and not self.header[col]['null']:
  773. raise ExecuteError(col + ' cannot be empty or null')
  774. if (self.header[col]['type'].lower() == 'int' or
  775. self.header[col]['type'].lower() == 'double'):
  776. try:
  777. if val:
  778. val = eval(val)
  779. except NameError, e:
  780. raise ExecuteError('Failed to parse %s in %s '
  781. '(unable to convert to type %s): %s' %
  782. (col, self.table, self.header[col]['type'], e))
  783. except SyntaxError, e:
  784. raise ExecuteError('Failed to parse %s in %s '
  785. '(unable to convert to type %s): %s' %
  786. (col, self.table, self.header[col]['type'], e))
  787. return val
  788. def _ColByPosition(self, pos):
  789. """Returns column name based on position.
  790. Args:
  791. pos: int
  792. Returns:
  793. column: string
  794. Raises:
  795. ExecuteError: invalid column
  796. """
  797. for col in self.header:
  798. if self.header[col]['pos'] == pos:
  799. return col
  800. raise ExecuteError('Header does not contain column %d' % pos)
  801. def _PositionByCol(self, col):
  802. """Returns position of the column based on the name.
  803. Args:
  804. col: string
  805. Returns:
  806. pos: int
  807. Raises:
  808. ExecuteError: invalid column
  809. """
  810. if col not in self.header.keys():
  811. raise ExecuteError(col + ' is not a valid column name')
  812. return self.header[col]['pos']
  813. def _SortHeaderColumns(self):
  814. """Sort column names by position.
  815. Returns:
  816. sorted: [col1, col2, ...]
  817. Raises:
  818. ExecuteError: unable to sort header
  819. """
  820. cols = self.header.keys()
  821. sorted_cols = [''] * len(cols)
  822. for col in cols:
  823. pos = self.header[col]['pos']
  824. sorted_cols[pos] = col
  825. if '' in sorted_cols:
  826. raise ExecuteError('Unable to sort header columns: %s' % cols)
  827. return sorted_cols
  828. def OpenTable(self):
  829. """Opens the table file and places its content into memory.
  830. Raises:
  831. ExecuteError: unable to open table
  832. """
  833. # if we already have a header assume multiple queries on same table
  834. # (can't use self.data in case the table was empty to begin with)
  835. if self.header:
  836. return
  837. try:
  838. self.fd = open(os.path.join(self.location, self.table), 'r')
  839. self.header = self._ProcessHeader()
  840. if self.command in ['INSERT', 'DELETE', 'UPDATE']:
  841. fcntl.flock(self.fd, fcntl.LOCK_EX)
  842. self.data = self._GetData()
  843. self.orig_data = self.data[:] # save a copy of the data before modifying
  844. except IOError, e:
  845. raise ExecuteError('Unable to open table %s: %s' % (self.table, e))
  846. Debug('Header is: %s' % self.header)
  847. # type check the conditions
  848. for cond in self.conditions:
  849. if cond not in self.header.keys():
  850. raise ExecuteError('unknown column %s in WHERE clause' % cond)
  851. self.conditions[cond] = self._TypeCheck(self.conditions[cond], cond)
  852. # type check the targets
  853. for target in self.targets:
  854. if target not in self.header.keys():
  855. raise ExecuteError('unknown column in targets: %s' % target)
  856. self.targets[target] = self._TypeCheck(self.targets[target], target)
  857. Debug('Type checked conditions: %s' % self.conditions)
  858. Debug('Data is:')
  859. for row in self.data:
  860. Debug('=======================')
  861. Debug(row)
  862. Debug('=======================')
  863. def WriteTempTable(self):
  864. """Write table header and data.
  865. First write header and data to a temp file,
  866. then move the tmp file to replace the original table file.
  867. """
  868. self.temp_file = tempfile.NamedTemporaryFile()
  869. Debug('temp_file: ' + self.temp_file.name)
  870. # write header
  871. columns = self._SortHeaderColumns()
  872. header = ''
  873. for col in columns:
  874. header = '%s %s' % (header, col)
  875. header = '%s(%s' % (header, self.header[col]['type'])
  876. if self.header[col]['null']:
  877. header = '%s,null)' % header
  878. elif self.header[col]['auto']:
  879. header = '%s,auto)' % header
  880. else:
  881. header = '%s)' % header
  882. self.temp_file.write(header.strip() + '\n')
  883. # write data
  884. for row in self.data:
  885. row_str = ''
  886. for col in columns:
  887. row_str = '%s:%s' % (row_str, row[col])
  888. self.temp_file.write(row_str[1:] + '\n')
  889. self.temp_file.flush()
  890. def MoveTableIntoPlace(self):
  891. """Replace the real table with the temp one.
  892. Diff the new data against the original and replace the table when they are
  893. different.
  894. """
  895. if self.data != self.orig_data:
  896. temp_file = self.temp_file.name
  897. table_file = os.path.join(self.location, self.table)
  898. Debug('Copying %s to %s' % (temp_file, table_file))
  899. shutil.copy(self.temp_file.name, self.location + '/' + self.table)
  900. def _ParsePairs(self, s, delimeter):
  901. """Parses out name value pairs from a string.
  902. String contains name=value pairs
  903. separated by a delimiter (such as "and" or ",")
  904. Args:
  905. s: string
  906. delimeter: string
  907. Returns:
  908. my_dict: dictionary
  909. Raises:
  910. ParseError: unable to parse pairs
  911. """
  912. my_dict = {}
  913. Debug('parse pairs: [%s]' % s)
  914. pairs = s.split(delimeter)
  915. for pair in pairs:
  916. if '=' not in pair:
  917. raise ParseError('Invalid condition pair: ' + pair)
  918. split = pair.split('=')
  919. Debug('split: %s' % split)
  920. if len(split) != 2:
  921. raise ParseError('Invalid condition pair: ' + pair)
  922. col = split[0].strip()
  923. if not col or not split[1].strip() or ' ' in col:
  924. raise ParseError('Invalid condition pair: ' + pair)
  925. val = self._ReplaceStringLiterals(split[1].strip())
  926. my_dict[col] = val
  927. return my_dict
  928. class Error(Exception):
  929. """DBText error."""
  930. class ParseError(Error):
  931. """Parse error."""
  932. class NotSupportedError(Error):
  933. """Not Supported error."""
  934. class ExecuteError(Error):
  935. """Execute error."""
  936. def main(argv):
  937. if len(argv) < 2:
  938. print 'Usage %s query' % argv[0]
  939. sys.exit(1)
  940. if 'DBTEXT_PATH' not in os.environ or not os.environ['DBTEXT_PATH']:
  941. print 'DBTEXT_PATH must be set'
  942. sys.exit(1)
  943. else:
  944. location = os.environ['DBTEXT_PATH']
  945. try:
  946. conn = DBText(location)
  947. dataset = conn.Execute(' '.join(argv[1:]))
  948. if dataset:
  949. for row in dataset:
  950. if conn.command != 'SELECT':
  951. print 'Updated %s, rows affected: %d' % (conn.table, row)
  952. else:
  953. print row
  954. except Error, e:
  955. print e
  956. sys.exit(1)
  957. if __name__ == '__main__':
  958. main(sys.argv)