Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# -*- coding: utf-8 -*- 

2""" 

3@file 

4@brief Defines SQL commands to play with `sqlite3 <https://docs.python.org/3/library/sqlite3.html>`_. 

5See notebook :ref:`pyensaesqlmagicrst`. 

6""" 

7import os 

8 

9from IPython.core.magic import magics_class, line_magic 

10from IPython.core.magic import line_cell_magic 

11from pyquickhelper.ipythonhelper import MagicCommandParser, MagicClassWithHelpers 

12from .sql_interface import InterfaceSQL, InterfaceSQLException 

13 

14 

15@magics_class 

16class MagicSQL(MagicClassWithHelpers): 

17 

18 """ 

19 Defines SQL commands to play with `sqlite3 <https://docs.python.org/3.4/library/sqlite3.html>`_ 

20 See notebook :ref:`pyensaesqlmagicrst`. 

21 """ 

22 

23 def get_connection(self, name): 

24 """ 

25 returns the connection stored in the workspace 

26 

27 @param name variable name of the database 

28 @return object 

29 """ 

30 if isinstance(name, str): 

31 if self.shell is None: 

32 raise Exception("No detected workspace.") 

33 

34 if name not in self.shell.user_ns: 

35 raise KeyError( 

36 "No opened sqlite3 database called: " + str(name)) 

37 

38 res = self.shell.user_ns[name] 

39 else: 

40 res = name 

41 if not isinstance(res, InterfaceSQL): 

42 res = InterfaceSQL.create(res) 

43 return res 

44 

45 @staticmethod 

46 def SQL_connect_parser(): 

47 """ 

48 defines the way to parse the magic command ``%SQL_connect`` 

49 """ 

50 parser = MagicCommandParser(prog="SQL_connect", 

51 description='connect to a SQL database') 

52 parser.add_argument('filename', type=str, 

53 help='database filename', eval_type=str) 

54 parser.add_argument( 

55 '-v', 

56 '--variable', 

57 default="DB", 

58 help='variable name used to store the database object') 

59 return parser 

60 

61 @line_magic 

62 def SQL_connect(self, line): 

63 """ 

64 define ``SQL_connect`` which connects to a SQL database, 

65 it stores the database object in variable DB by default 

66 

67 .. nbref:: 

68 :tag: SQL 

69 :title: SQL_connect_parser 

70 

71 The code for magic command ``%SQL_connect_parser`` is equivalent to:: 

72 

73 from pyense.sql import InterfaceSQL 

74 obj = InterfaceSQL.create(args.filename) 

75 obj.connect() 

76 

77 See notebook :ref:`pyensaesqlmagicrst`. 

78 """ 

79 parser = self.get_parser(MagicSQL.SQL_connect_parser, "SQL_connect") 

80 args = self.get_args(line, parser) 

81 

82 if args is not None: 

83 obj = InterfaceSQL.create(args.filename) 

84 obj.connect() 

85 self.shell.user_ns[args.variable] = obj 

86 return obj 

87 

88 @staticmethod 

89 def SQL_close_parser(): 

90 """ 

91 defines the way to parse the magic command ``%SQL_close`` 

92 """ 

93 parser = MagicCommandParser(prog="SQL_close", 

94 description='connect to a SQL database') 

95 parser.add_argument( 

96 '-v', 

97 '--variable', 

98 default="DB", 

99 help='variable name used to store the database object (and to close)') 

100 return parser 

101 

102 @line_magic 

103 def SQL_close(self, line=""): 

104 """ 

105 define ``SQL_close`` which closes a database 

106 

107 .. nbref:: 

108 :tag: SQL 

109 :title: SQL_close 

110 

111 The code for magic command ``%SQL_close`` is equivalent to:: 

112 

113 db.close() 

114 

115 See notebook :ref:`pyensaesqlmagicrst`. 

116 """ 

117 parser = self.get_parser(MagicSQL.SQL_close_parser, "SQL_close") 

118 args = self.get_args(line, parser) 

119 

120 if args is not None: 

121 db = self.get_connection(args.variable) 

122 r = db.close() 

123 return r 

124 

125 @staticmethod 

126 def SQL_tables_parser(): 

127 """ 

128 defines the way to parse the magic command ``%SQL_tables`` 

129 """ 

130 parser = MagicCommandParser(prog="SQL_tables", 

131 description='list the tables of a database') 

132 parser.add_argument( 

133 '-v', 

134 '--variable', 

135 default="DB", 

136 help='variable name used to store the database object') 

137 return parser 

138 

139 @line_magic 

140 def SQL_tables(self, line=""): 

141 """ 

142 define ``%SQL_tables`` whichs lists the tables in a database 

143 

144 .. nbref:: 

145 :tag: SQL 

146 :title: SQL_tables 

147 

148 The code for magic command ``%SQL_tables`` is equivalent to:: 

149 

150 db.get_table_list() 

151 

152 See notebook :ref:`pyensaesqlmagicrst`. 

153 """ 

154 parser = self.get_parser(MagicSQL.SQL_tables_parser, "SQL_tables") 

155 args = self.get_args(line, parser) 

156 

157 if args is not None: 

158 db = self.get_connection(args.variable) 

159 return db.get_table_list() 

160 

161 @staticmethod 

162 def SQL_drop_table_parser(): 

163 """ 

164 defines the way to parse the magic command ``%SQL_drop_table`` 

165 """ 

166 parser = MagicCommandParser(prog="SQL_drop_table", 

167 description='drop a table from a database') 

168 parser.add_argument('table', type=str, help='table', eval_type=str) 

169 parser.add_argument( 

170 '-v', 

171 '--variable', 

172 default="DB", 

173 help='variable name used to store the database object') 

174 return parser 

175 

176 @line_magic 

177 def SQL_drop_table(self, line): 

178 """ 

179 defines ``%SQL_drop_table`` which drops a table from a database 

180 

181 .. nbref:: 

182 :tag: SQL 

183 :title: SQL_drop_table 

184 

185 The code for magic command ``%SQL_drop_table`` is equivalent to:: 

186 

187 db.drop_table() 

188 

189 See notebook :ref:`pyensaesqlmagicrst`. 

190 """ 

191 parser = self.get_parser( 

192 MagicSQL.SQL_drop_table_parser, "SQL_drop_table") 

193 args = self.get_args(line, parser) 

194 

195 if args is not None: 

196 db = self.get_connection(args.variable) 

197 return db.drop_table(args.table) 

198 

199 @staticmethod 

200 def SQL_refresh_completion_parser(): 

201 """ 

202 defines the way to parse the magic command ``%SQL_refresh_completion`` 

203 """ 

204 parser = MagicCommandParser(prog="SQL_refresh_completion", 

205 description='refresh completion (tables names, ...)') 

206 parser.add_argument( 

207 '-v', 

208 '--variable', 

209 default="DB", 

210 help='variable name used to store the database object') 

211 return parser 

212 

213 @line_magic 

214 def SQL_refresh_completion(self, line=""): 

215 """ 

216 defines ``%SQL_refresh_completion`` 

217 """ 

218 parser = self.get_parser( 

219 MagicSQL.SQL_refresh_completion_parser, "SQL_refresh_completion") 

220 args = self.get_args(line, parser) 

221 

222 if args is not None: 

223 db = self.get_connection(args.variable) 

224 db.refresh_completion() 

225 

226 @staticmethod 

227 def SQL_schema_parser(): 

228 """ 

229 defines the way to parse the magic command ``%SQL_schema`` 

230 """ 

231 parser = MagicCommandParser(prog="SQL_schema", 

232 description='schema of a table') 

233 parser.add_argument('table', type=str, help='table', eval_type=str) 

234 parser.add_argument( 

235 '--as_list', help='as a dictionary (False) or as a list (True)', action="store_true") 

236 parser.add_argument( 

237 '-v', 

238 '--variable', 

239 default="DB", 

240 help='variable name used to store the database object') 

241 return parser 

242 

243 @line_magic 

244 def SQL_schema(self, line=""): 

245 """ 

246 define ``SQL_schema`` 

247 

248 .. nbref:: 

249 :tag: SQL 

250 :title: SQL_schema 

251 

252 The code for magic command ``%SQL_schema`` is equivalent to:: 

253 

254 db.get_table_columns(<table>, as_dict=not <as_list>) 

255 

256 See notebook :ref:`pyensaesqlmagicrst`. 

257 """ 

258 parser = self.get_parser(MagicSQL.SQL_schema_parser, "SQL_schema") 

259 args = self.get_args(line, parser) 

260 

261 if args is not None: 

262 db = self.get_connection(args.variable) 

263 return db.get_table_columns(args.table, as_dict=not args.as_list) 

264 

265 @staticmethod 

266 def SQL_import_tsv_parser(): 

267 """ 

268 defines the way to parse the magic command ``%SQL_import_tsv`` 

269 """ 

270 parser = MagicCommandParser(prog="SQL_import_tsv", 

271 description='import a tsv file into the database') 

272 parser.add_argument('filename', type=str, 

273 help='tsv file name', eval_type=str) 

274 parser.add_argument('-t', '--table', type=str, 

275 help='table name', default="-", eval_type=str) 

276 parser.add_argument( 

277 '--verbose', help='print progress', action="store_true") 

278 parser.add_argument( 

279 '-v', 

280 '--variable', 

281 default="DB", 

282 help='variable name used to store the database object') 

283 return parser 

284 

285 @line_magic 

286 def SQL_import_tsv(self, line): 

287 """ 

288 defines ``%SQL_import_tsv`` whichs import a TSV file into a database 

289 

290 .. nbref:: 

291 :tag: SQL 

292 :title: SQL_import_tsv 

293 

294 The code for magic command ``%SQL_import_tsv`` is equivalent to:: 

295 

296 db.import_flat_file(<filename>, <table>) 

297 

298 See notebook :ref:`pyensaesqlmagicrst`. 

299 """ 

300 parser = self.get_parser( 

301 MagicSQL.SQL_import_tsv_parser, "SQL_import_tsv") 

302 args = self.get_args(line, parser) 

303 

304 if args is not None: 

305 if not os.path.exists(args.filename): 

306 raise FileNotFoundError(args.filename) 

307 db = self.get_connection(args.variable) 

308 table = os.path.splitext(os.path.split(args.filename)[-1])[0] \ 

309 if len(args.table) == 0 or args.table == "-" else args.table 

310 return db.import_flat_file(args.filename, table) 

311 

312 @staticmethod 

313 def SQL_add_function_parser(): 

314 """ 

315 defines the way to parse the magic command ``%SQL_add_function`` 

316 """ 

317 parser = MagicCommandParser(prog="SQL_add_function", 

318 description='add a custom function to the database') 

319 parser.add_argument('funct', type=str, help='function name') 

320 parser.add_argument( 

321 '-v', 

322 '--variable', 

323 default="DB", 

324 help='variable name used to store the database object') 

325 return parser 

326 

327 @line_magic 

328 def SQL_add_function(self, line): 

329 """ 

330 defines ``%SQL_add_function`` which adds a function to the database 

331 

332 .. nbref:: 

333 :tag: SQL 

334 :title: SQL_add_function 

335 

336 The code for magic command ``%SQL_add_function`` is equivalent to:: 

337 

338 db.add_function(fu) 

339 

340 See notebook :ref:`pyensaesqlmagicrst`. 

341 """ 

342 parser = self.get_parser( 

343 MagicSQL.SQL_add_function_parser, "SQL_add_function") 

344 args = self.get_args(line, parser) 

345 

346 if args is not None: 

347 db = self.get_connection(args.variable) 

348 if isinstance(args.funct, str): 

349 if self.shell is not None: 

350 if args.funct not in self.shell.user_ns: 

351 raise KeyError( 

352 "unable to find function %s in your workspace" % 

353 args.funct) 

354 fu = self.shell.user_ns[args.funct] 

355 else: 

356 raise Exception("unable to find IPython workspace") 

357 else: 

358 fu = args.funct 

359 return db.add_function(fu) 

360 

361 @staticmethod 

362 def SQL_import_df_parser(): 

363 """ 

364 defines the way to parse the magic command ``%SQL_import_df`` 

365 """ 

366 parser = MagicCommandParser(prog="SQL_import_df", 

367 description='import a dataframe into the database') 

368 parser.add_argument('df', type=str, help='dataframe', no_eval=True) 

369 parser.add_argument('-t', '--table', type=str, 

370 help='table name', default="-", eval_type=str) 

371 parser.add_argument( 

372 '-v', 

373 '--variable', 

374 default="DB", 

375 help='variable name used to store the database object') 

376 return parser 

377 

378 @line_magic 

379 def SQL_import_df(self, line): 

380 """ 

381 defines ``%SQL_import_df`` which imports a dataframe into a database 

382 

383 .. nbref:: 

384 :tag: SQL 

385 :title: SQL_import_df 

386 

387 The code for magic command ``%SQL_import_df`` is equivalent to:: 

388 

389 db.import_dataframe(<table>, <df>) 

390 

391 See notebook :ref:`pyensaesqlmagicrst`. 

392 """ 

393 parser = self.get_parser( 

394 MagicSQL.SQL_import_df_parser, "SQL_import_df") 

395 args = self.get_args(line, parser) 

396 

397 if args is not None: 

398 db = self.get_connection(args.variable) 

399 

400 df = args.df 

401 if self.shell is not None: 

402 if df not in self.shell.user_ns: 

403 raise KeyError( 

404 "unable to find dataframe %s in your workspace" % 

405 df) 

406 odf = self.shell.user_ns[df] 

407 else: 

408 raise Exception("unable to find IPython workspace") 

409 

410 table = df if len( 

411 args.table) == 0 or args.table == "-" else args.table 

412 return db.import_dataframe(table, odf) 

413 

414 @staticmethod 

415 def SQL_parser(): 

416 """ 

417 defines the way to parse the magic command ``%%SQL`` 

418 """ 

419 parser = MagicCommandParser(prog="SQL", 

420 description='query the database') 

421 parser.add_argument( 

422 '--df', type=str, help='output dataframe', default="temp_view", no_eval=True) 

423 parser.add_argument('-n', '--n', type=int, 

424 help='number of first lines to display', default=10, eval_type=int) 

425 parser.add_argument('-q', '--query', type=str, 

426 help='when used in a single line (no cell), query is the SQL query, the command ' + 

427 'returns the full dataframe', default="", eval_type=str) 

428 parser.add_argument( 

429 '-v', 

430 '--variable', 

431 default="DB", 

432 help='variable name used to store the database object') 

433 return parser 

434 

435 @line_cell_magic 

436 def SQL(self, line, cell=None): 

437 """ 

438 defines command ``%%SQL`` 

439 

440 .. nbref:: 

441 :tag: SQL 

442 :title: SQL 

443 

444 The code for magic command ``%%SQL`` is equivalent to:: 

445 

446 <variable> = db.execute(<cell>) 

447 

448 See notebook :ref:`pyensaesqlmagicrst`. 

449 """ 

450 parser = self.get_parser(MagicSQL.SQL_parser, "SQL") 

451 args = self.get_args(line, parser) 

452 

453 if args is not None: 

454 full = False 

455 if cell is None or len(cell) == 0: 

456 cell = args.query 

457 if cell is None or len(cell) == 0: 

458 raise ValueError("no SQL query is defined") 

459 query = cell 

460 full = True 

461 else: 

462 query = cell 

463 

464 query = query.strip() 

465 if len(query) > 0 and query[0] == '"' and query[-1] == '"': 

466 query = query[1:-1] 

467 

468 db = self.get_connection(args.variable) 

469 try: 

470 df = db.execute(query) 

471 ok = True 

472 except InterfaceSQLException as e: 

473 print(str(e)) 

474 ok = False 

475 

476 if ok: 

477 self.shell.user_ns[args.df] = df 

478 if full: 

479 return df 

480 else: 

481 return df.head(n=args.n) 

482 

483 

484def register_sql_magics(ip=None): 

485 """ 

486 register magics function, can be called from a notebook 

487 

488 @param ip from ``get_ipython()`` 

489 """ 

490 if ip is None: 

491 from IPython import get_ipython 

492 ip = get_ipython() 

493 ip.register_magics(MagicSQL)