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
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
15@magics_class
16class MagicSQL(MagicClassWithHelpers):
18 """
19 Defines SQL commands to play with `sqlite3 <https://docs.python.org/3.4/library/sqlite3.html>`_
20 See notebook :ref:`pyensaesqlmagicrst`.
21 """
23 def get_connection(self, name):
24 """
25 returns the connection stored in the workspace
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.")
34 if name not in self.shell.user_ns:
35 raise KeyError(
36 "No opened sqlite3 database called: " + str(name))
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
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
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
67 .. nbref::
68 :tag: SQL
69 :title: SQL_connect_parser
71 The code for magic command ``%SQL_connect_parser`` is equivalent to::
73 from pyense.sql import InterfaceSQL
74 obj = InterfaceSQL.create(args.filename)
75 obj.connect()
77 See notebook :ref:`pyensaesqlmagicrst`.
78 """
79 parser = self.get_parser(MagicSQL.SQL_connect_parser, "SQL_connect")
80 args = self.get_args(line, parser)
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
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
102 @line_magic
103 def SQL_close(self, line=""):
104 """
105 define ``SQL_close`` which closes a database
107 .. nbref::
108 :tag: SQL
109 :title: SQL_close
111 The code for magic command ``%SQL_close`` is equivalent to::
113 db.close()
115 See notebook :ref:`pyensaesqlmagicrst`.
116 """
117 parser = self.get_parser(MagicSQL.SQL_close_parser, "SQL_close")
118 args = self.get_args(line, parser)
120 if args is not None:
121 db = self.get_connection(args.variable)
122 r = db.close()
123 return r
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
139 @line_magic
140 def SQL_tables(self, line=""):
141 """
142 define ``%SQL_tables`` whichs lists the tables in a database
144 .. nbref::
145 :tag: SQL
146 :title: SQL_tables
148 The code for magic command ``%SQL_tables`` is equivalent to::
150 db.get_table_list()
152 See notebook :ref:`pyensaesqlmagicrst`.
153 """
154 parser = self.get_parser(MagicSQL.SQL_tables_parser, "SQL_tables")
155 args = self.get_args(line, parser)
157 if args is not None:
158 db = self.get_connection(args.variable)
159 return db.get_table_list()
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
176 @line_magic
177 def SQL_drop_table(self, line):
178 """
179 defines ``%SQL_drop_table`` which drops a table from a database
181 .. nbref::
182 :tag: SQL
183 :title: SQL_drop_table
185 The code for magic command ``%SQL_drop_table`` is equivalent to::
187 db.drop_table()
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)
195 if args is not None:
196 db = self.get_connection(args.variable)
197 return db.drop_table(args.table)
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
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)
222 if args is not None:
223 db = self.get_connection(args.variable)
224 db.refresh_completion()
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
243 @line_magic
244 def SQL_schema(self, line=""):
245 """
246 define ``SQL_schema``
248 .. nbref::
249 :tag: SQL
250 :title: SQL_schema
252 The code for magic command ``%SQL_schema`` is equivalent to::
254 db.get_table_columns(<table>, as_dict=not <as_list>)
256 See notebook :ref:`pyensaesqlmagicrst`.
257 """
258 parser = self.get_parser(MagicSQL.SQL_schema_parser, "SQL_schema")
259 args = self.get_args(line, parser)
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)
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
285 @line_magic
286 def SQL_import_tsv(self, line):
287 """
288 defines ``%SQL_import_tsv`` whichs import a TSV file into a database
290 .. nbref::
291 :tag: SQL
292 :title: SQL_import_tsv
294 The code for magic command ``%SQL_import_tsv`` is equivalent to::
296 db.import_flat_file(<filename>, <table>)
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)
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)
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
327 @line_magic
328 def SQL_add_function(self, line):
329 """
330 defines ``%SQL_add_function`` which adds a function to the database
332 .. nbref::
333 :tag: SQL
334 :title: SQL_add_function
336 The code for magic command ``%SQL_add_function`` is equivalent to::
338 db.add_function(fu)
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)
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)
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
378 @line_magic
379 def SQL_import_df(self, line):
380 """
381 defines ``%SQL_import_df`` which imports a dataframe into a database
383 .. nbref::
384 :tag: SQL
385 :title: SQL_import_df
387 The code for magic command ``%SQL_import_df`` is equivalent to::
389 db.import_dataframe(<table>, <df>)
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)
397 if args is not None:
398 db = self.get_connection(args.variable)
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")
410 table = df if len(
411 args.table) == 0 or args.table == "-" else args.table
412 return db.import_dataframe(table, odf)
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
435 @line_cell_magic
436 def SQL(self, line, cell=None):
437 """
438 defines command ``%%SQL``
440 .. nbref::
441 :tag: SQL
442 :title: SQL
444 The code for magic command ``%%SQL`` is equivalent to::
446 <variable> = db.execute(<cell>)
448 See notebook :ref:`pyensaesqlmagicrst`.
449 """
450 parser = self.get_parser(MagicSQL.SQL_parser, "SQL")
451 args = self.get_args(line, parser)
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
464 query = query.strip()
465 if len(query) > 0 and query[0] == '"' and query[-1] == '"':
466 query = query[1:-1]
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
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)
484def register_sql_magics(ip=None):
485 """
486 register magics function, can be called from a notebook
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)