jupyter-duckdb 1.2.102__py3-none-any.whl → 1.2.103__py3-none-any.whl

duckdb_kernel/db/Connection.py

@@ -10,6 +10,9 @@ class Connection:
     def close(self):
         pass
 
+    def copy(self) -> 'Connection':
+        raise NotImplementedError
+
     @staticmethod
     def plain_explain() -> bool:
         return False

duckdb_kernel/db/implementation/duckdb/Connection.py

@@ -15,6 +15,9 @@ class Connection(Base):
     def close(self):
         self.con.close()
 
+    def copy(self) -> 'Connection':
+        return Connection(self.path)
+
     @staticmethod
     def plain_explain() -> bool:
         return True

duckdb_kernel/db/implementation/postgres/Connection.py

@@ -19,6 +19,7 @@ class Connection(Base):
         self.host: str = host
         self.port: int = port
         self.username: Optional[str] = username
+        self.password: Optional[str] = password
 
         options = {
             'host': host,
@@ -43,6 +44,9 @@ class Connection(Base):
     def close(self):
         self.con.close()
 
+    def copy(self) -> 'Connection':
+        return Connection(self.host, self.port, self.username, self.password, self.database_name)
+
     def __str__(self) -> str:
         user = f'{self.username}@' if self.username is not None else ''
         return f'PostgreSQL: {user}{self.host}:{self.port}/{self.database_name}'

duckdb_kernel/db/implementation/sqlite/Connection.py

@@ -16,6 +16,9 @@ class Connection(Base):
     def close(self):
         self.con.close()
 
+    def copy(self) -> 'Connection':
+        return Connection(self.path)
+
     @staticmethod
     def multiple_statements_per_query() -> bool:
         return False

duckdb_kernel/kernel.py

@@ -10,7 +10,6 @@ from typing import Optional, Dict, List, Tuple
 from ipykernel.kernelbase import Kernel
 
 from .db import Connection, DatabaseError, Table
-from .visualization.lib import *
 from .db.error import *
 from .magics import *
 from .parser import RAParser, DCParser, ParserError
@@ -40,8 +39,11 @@ class DuckDBKernel(Kernel):
         self._magics: MagicCommandHandler = MagicCommandHandler()
 
         self._magics.add(
-            MagicCommand('create').arg('database').opt('of').opt('with_tests').on(self._create_magic),
-            MagicCommand('load').arg('database').opt('with_tests').on(self._load_magic),
+            MagicCommand('create').arg('database').opt('of').opt('name').on(self._create_magic),
+            MagicCommand('load').arg('database').opt('name').on(self._load_magic),
+            MagicCommand('copy').arg('source').arg('target').on(self._copy_magic),
+            MagicCommand('use').arg('name').on(self._use_magic),
+            MagicCommand('load_tests').arg('tests').on(self._load_tests_magic),
             MagicCommand('test').arg('name').result(True).on(self._test_magic),
             MagicCommand('all', 'all_rows').on(self._all_magic),
             MagicCommand('max_rows').arg('count').on(self._max_rows_magic),
@@ -60,8 +62,8 @@ class DuckDBKernel(Kernel):
         )
 
         # create placeholders for database and tests
-        self._db: Optional[Connection] = None
-        self._tests: Optional[Dict] = None
+        self._db: Dict[str, Connection] = {}
+        self._tests: Dict = {}
 
     # output related functions
     def print(self, text: str, name: str = 'stdout'):
@@ -98,68 +100,64 @@ class DuckDBKernel(Kernel):
             })
 
     # database related functions
-    def _load_database(self, path: str):
-        if self._db is None:
-            # If the provided path looks like a postgres url,
-            # we want to use the postgres driver.
-            if path.startswith(('postgresql://', 'postgres://', 'pgsql://', 'psql://', 'pg://')):
-                # pull data from connection string
-                re_expr = r'(postgresql|postgres|pgsql|psql|pg)://((.*?)(:(.*?))?@)?(.*?)(:(\d+))?(/(.*))?'
-                match = re.fullmatch(re_expr, path)
-
-                host = match.group(6)
-                port = int(match.group(8)) if match.group(8) is not None else 5432
-                username = match.group(3)
-                password = match.group(5)
-                database_name = match.group(10)
-
-                # load and create instance
-                try:
-                    from .db.implementation.postgres import Connection as Postgres
-                    self._db = Postgres(host, port, username, password, database_name)
-                except ImportError:
-                    self.print('psycopg could not be found', name='stderr')
-                    return False
-
-            # Otherwise the provided path is used to create an
-            # in-process instance.
-            else:
-                # By default, we try to load DuckDB.
-                try:
-                    from .db.implementation.duckdb import Connection as DuckDB
-                    self._db = DuckDB(path)
+    def _load_database(self, path: str, name: str) -> Connection:
+        if name in self._db:
+            raise ValueError(f'duplicate database name {name}')
+
+        # If the provided path looks like a postgres url,
+        # we want to use the postgres driver.
+        if path.startswith(('postgresql://', 'postgres://', 'pgsql://', 'psql://', 'pg://')):
+            # pull data from connection string
+            re_expr = r'(postgresql|postgres|pgsql|psql|pg)://((.*?)(:(.*?))?@)?(.*?)(:(\d+))?(/(.*))?'
+            match = re.fullmatch(re_expr, path)
+
+            host = match.group(6)
+            port = int(match.group(8)) if match.group(8) is not None else 5432
+            username = match.group(3)
+            password = match.group(5)
+            database_name = match.group(10)
+
+            # load and create instance
+            from .db.implementation.postgres import Connection as Postgres
+            self._db[name] = Postgres(host, port, username, password, database_name)
+
+        # Otherwise the provided path is used to create an
+        # in-process instance.
+        else:
+            # By default, we try to load DuckDB.
+            try:
+                from .db.implementation.duckdb import Connection as DuckDB
+                self._db[name] = DuckDB(path)
 
-                # If DuckDB is not installed or fails to load,
-                # we use SQLite instead.
-                except ImportError:
-                    self.print('DuckDB is not available\n')
+            # If DuckDB is not installed or fails to load,
+            # we use SQLite instead.
+            except ImportError:
+                self.print('DuckDB is not available\n')
 
-                    from .db.implementation.sqlite import Connection as SQLite
-                    self._db = SQLite(path)
+                from .db.implementation.sqlite import Connection as SQLite
+                self._db[name] = SQLite(path)
 
-            return True
-        else:
-            return False
+        return self._db[name]
+
+    def _unload_database(self, name: str):
+        if name in self._db:
+            self._db[name].close()
+            del self._db[name]
 
-    def _unload_database(self):
-        if self._db is not None:
-            self._db.close()
-            self._db = None
             return True
         else:
             return False
 
-    def _execute_stmt(self, name: str, query: str, silent: bool,
-                      column_name_mapping: Dict[str, str],
-                      max_rows: Optional[int]) -> Tuple[Optional[List[str]], Optional[List[List]]]:
-        if self._db is None:
+    def _execute_stmt(self, silent: bool, state: MagicState, name: str, query: str) \
+            -> Tuple[Optional[List[str]], Optional[List[List]]]:
+        if state.db is None:
             raise AssertionError('load a database first')
 
         # execute query and store start and end timestamp
         st = time.time()
 
         try:
-            columns, rows = self._db.execute(query)
+            columns, rows = state.db.execute(query)
         except EmptyResultError:
             columns, rows = None, None
 
@@ -168,32 +166,32 @@ class DuckDBKernel(Kernel):
         # print result if not silent
         if not silent:
             # print EXPLAIN queries as raw text if using DuckDB
-            if query.strip().startswith('EXPLAIN') and self._db.plain_explain():
+            if query.strip().startswith('EXPLAIN') and state.db.plain_explain():
                 for ekey, evalue in rows:
                     html = f'<b>{ekey}</b><br><pre>{evalue}</pre>'
                     break
-
-                return None, None
+                else:
+                    html = ''
 
             # print every other query as a table
             else:
                 if columns is not None:
                     # table header
-                    mapped_columns = (column_name_mapping.get(c, c) for c in columns)
+                    mapped_columns = (state.column_name_mapping.get(c, c) for c in columns)
                     table_header = ''.join(f'<th>{c}</th>' for c in mapped_columns)
 
                     # table data
-                    if max_rows is not None and len(rows) > max_rows:
+                    if state.max_rows is not None and len(rows) > state.max_rows:
                         table_data = f'''
-                            {rows_table(rows[:math.ceil(max_rows / 2)])}
+                            {rows_table(rows[:math.ceil(state.max_rows / 2)])}
                             <tr>
                                 <td colspan="{len(columns)}" 
                                     style="text-align: center"
-                                    title="{row_count(len(rows) - max_rows)} omitted">
+                                    title="{row_count(len(rows) - state.max_rows)} omitted">
                                     ...
                                 </td>
                             </tr>
-                            {rows_table(rows[-math.floor(max_rows // 2):])}
+                            {rows_table(rows[-math.floor(state.max_rows // 2):])}
                         '''
                     else:
                         table_data = rows_table(rows)
@@ -220,15 +218,25 @@ class DuckDBKernel(Kernel):
         return columns, rows
 
     # magic command related functions
-    def _create_magic(self, silent: bool, path: str, of: Optional[str], with_tests: Optional[str]):
-        self._load(silent, path, True, of, with_tests)
+    def _create_magic(self, silent: bool, state: MagicState,
+                      path: str, of: Optional[str], name: Optional[str]):
+        self._load(silent, state, path, True, of, name)
+
+    def _load_magic(self, silent: bool, state: MagicState,
+                    path: str, name: Optional[str]):
+        self._load(silent, state, path, False, None, name)
+
+    def _load(self, silent: bool, state: MagicState,
+              path: str, create: bool, of: Optional[str], name: Optional[str]):
+        # use default name if non provided
+        if name is None:
+            name = 'default'
 
-    def _load_magic(self, silent: bool, path: str, with_tests: Optional[str]):
-        self._load(silent, path, False, None, with_tests)
+        if not silent:
+            self.print(f'--- connection {name} ---\n')
 
-    def _load(self, silent: bool, path: str, create: bool, of: Optional[str], with_tests: Optional[str]):
         # unload current database if necessary
-        if self._unload_database():
+        if self._unload_database(name):
             if not silent:
                 self.print('unloaded database\n')
 
@@ -245,10 +253,10 @@ class DuckDBKernel(Kernel):
         if create and os.path.exists(path):
             os.remove(path)
 
-        if self._load_database(path):
-            if not silent:
-                # self.print(f'loaded database "{path}"\n')
-                self.print(str(self._db) + '\n')
+        state.db = self._load_database(path, name)
+        if not silent:
+            # self.print(f'loaded database "{path}"\n')
+            self.print(str(state.db) + '\n')
 
         # copy data from source database
         if of is not None:
@@ -264,18 +272,18 @@ class DuckDBKernel(Kernel):
                     content = file.read()
 
                     # You can only execute one statement at a time using SQLite.
-                    if not self._db.multiple_statements_per_query():
+                    if not state.db.multiple_statements_per_query():
                         statements = re.split(r';\r?\n', content)
                         for statement in statements:
                             try:
-                                self._db.execute(statement)
+                                state.db.execute(statement)
                             except EmptyResultError:
                                 pass
 
                     # Other DBMS can execute multiple statements at a time.
                     else:
                         try:
-                            self._db.execute(content)
+                            state.db.execute(content)
                         except EmptyResultError:
                             pass
 
@@ -289,29 +297,51 @@ class DuckDBKernel(Kernel):
                     of_db.execute('SHOW TABLES')
                     for table, in of_db.fetchall():
                         transfer_df = of_db.query(f'SELECT * FROM {table}').to_df()
-                        self._db.execute(f'CREATE TABLE {table} AS SELECT * FROM transfer_df')
+                        state.db.execute(f'CREATE TABLE {table} AS SELECT * FROM transfer_df')
 
                         if not silent:
                             self.print(f'transferred table {table}\n')
 
-        # load tests
-        if with_tests is None:
-            self._tests = {}
-        else:
-            with open(with_tests, 'r', encoding='utf-8') as tests_file:
-                self._tests = json.load(tests_file)
-                for test in self._tests.values():
-                    if 'attributes' in test:
-                        rows = {k: [] for k in test['attributes']}
-                        for row in test['equals']:
-                            for k, v in zip(test['attributes'], row):
-                                rows[k].append(v)
+    def _copy_magic(self, silent: bool, state: MagicState, source: str, target: str):
+        if source not in self._db:
+            raise ValueError(f'unknown connection {source}')
+
+        if not silent:
+            self.print(f'--- connection {target} ---\n')
+
+        # unload current database if necessary
+        if self._unload_database(target):
+            if not silent:
+                self.print('unloaded database\n')
+
+        # copy connection
+        self._db[target] = self._db[source].copy()
+        state.db = self._db[target]
+
+        if not silent:
+            self.print(str(state.db) + '\n')
+
+    def _use_magic(self, silent: bool, state: MagicState, name: str):
+        if name not in self._db:
+            raise ValueError(f'unknown connection {name}')
+
+        state.db = self._db[name]
+
+    def _load_tests_magic(self, silent: bool, state: MagicState, tests: str):
+        with open(tests, 'r', encoding='utf-8') as tests_file:
+            self._tests = json.load(tests_file)
+            for test in self._tests.values():
+                if 'attributes' in test:
+                    rows = {k: [] for k in test['attributes']}
+                    for row in test['equals']:
+                        for k, v in zip(test['attributes'], row):
+                            rows[k].append(v)
 
-                        test['equals'] = rows
+                    test['equals'] = rows
 
-                self.print(f'loaded tests from {with_tests}\n')
+            self.print(f'loaded tests from {tests}\n')
 
-    def _test_magic(self, silent: bool, result_columns: List[str], result: List[List], name: str):
+    def _test_magic(self, silent: bool, state: MagicState, result_columns: List[str], result: List[List], name: str):
         # If the query was empty, result_columns and result may be None.
         if result_columns is None or result is None:
             self.print_data(wrap_image(False, 'Statement did not return data.'))
@@ -400,31 +430,29 @@ class DuckDBKernel(Kernel):
             elif above > 0:
                 raise TestError(f'{row_count(above)} unnecessary')
 
-    def _all_magic(self, silent: bool):
-        return {
-            'max_rows': None
-        }
+    def _all_magic(self, silent: bool, state: MagicState):
+        state.max_rows = None
 
-    def _max_rows_magic(self, silent: bool, count: str):
+    def _max_rows_magic(self, silent: bool, state: MagicState, count: str):
         if count.lower() != 'none':
             DuckDBKernel.DEFAULT_MAX_ROWS = int(count)
         else:
             DuckDBKernel.DEFAULT_MAX_ROWS = None
 
-    def _query_max_rows_magic(self, silent: bool, count: str):
-        return {
-            'max_rows': int(count) if count.lower() != 'none' else None
-        }
+        state.max_rows = DuckDBKernel.DEFAULT_MAX_ROWS
 
-    def _schema_magic(self, silent: bool, td: bool, only: Optional[str]):
-        if self._db is None:
-            raise AssertionError('load a database first')
+    def _query_max_rows_magic(self, silent: bool, state: MagicState, count: str):
+        state.max_rows = int(count) if count.lower() != 'none' else None
 
+    def _schema_magic(self, silent: bool, state: MagicState, td: bool, only: Optional[str]):
         if silent:
             return
 
+        if state.db is None:
+            raise AssertionError('load a database first')
+
         # analyze tables
-        tables = self._db.analyze()
+        tables = state.db.analyze()
 
         # apply filter
         if only is None:
@@ -466,7 +494,9 @@ class DuckDBKernel(Kernel):
 
         self.print_data(svg)
 
-    def _store_magic(self, silent: bool, result_columns: List[str], result: List[List], file: str, noheader: bool):
+    def _store_magic(self, silent: bool, state: MagicState,
+                     result_columns: List[str], result: List[List],
+                     file: str, noheader: bool):
         _, ext = file.rsplit('.', 1)
 
         # csv
@@ -486,44 +516,38 @@ class DuckDBKernel(Kernel):
         else:
             raise ValueError(f'extension {ext} not supported')
 
-    def _ra_magic(self, silent: bool, code: str, analyze: bool):
-        if self._db is None:
-            raise AssertionError('load a database first')
-
+    def _ra_magic(self, silent: bool, state: MagicState, analyze: bool):
         if silent:
             return
 
-        if not code.strip():
+        if not state.code.strip():
             return
 
+        if state.db is None:
+            raise AssertionError('load a database first')
+
         # analyze tables
-        tables = self._db.analyze()
+        tables = state.db.analyze()
 
         # parse ra input
-        root_node = RAParser.parse_query(code)
+        root_node = RAParser.parse_query(state.code)
 
         # create and show visualization
         if analyze:
-            vd = RATreeDrawer(self._db, root_node, tables)
+            vd = RATreeDrawer(state.db, root_node, tables)
+
             svg = vd.to_interactive_svg()
-            data = {
-                'generated_code': {
-                    node_id: node.to_sql_with_renamed_columns(tables)
-                    for node_id, node in vd.nodes.items()
-                }
-            }
+            self.print_data(svg)
 
-        else:
-            svg = ''
-            data = {
-                'generated_code': root_node.to_sql_with_renamed_columns(tables)
+            state.code = {
+                node_id: node.to_sql_with_renamed_columns(tables)
+                for node_id, node in vd.nodes.items()
             }
 
-        # return data
-        self.print_data(svg)
-        return data
+        else:
+            state.code = root_node.to_sql_with_renamed_columns(tables)
 
-    def _all_ra_magic(self, silent: bool, value: str):
+    def _all_ra_magic(self, silent: bool, state: MagicState, value: str):
         if value.lower() in ('1', 'on', 'true'):
             self._magics['ra'].default(True)
             self._magics['dc'].default(False)
@@ -532,31 +556,29 @@ class DuckDBKernel(Kernel):
         else:
             self._magics['ra'].default(False)
 
-    def _dc_magic(self, silent: bool, code: str):
-        if self._db is None:
-            raise AssertionError('load a database first')
-
+    def _dc_magic(self, silent: bool, state: MagicState):
         if silent:
             return
 
-        if not code.strip():
+        if not state.code.strip():
             return
 
+        if state.db is None:
+            raise AssertionError('load a database first')
+
         # analyze tables
-        tables = self._db.analyze()
+        tables = state.db.analyze()
 
         # parse dc input
-        root_node = DCParser.parse_query(code)
+        root_node = DCParser.parse_query(state.code)
 
         # generate sql
         sql, cnm = root_node.to_sql_with_renamed_columns(tables)
 
-        return {
-            'generated_code': sql,
-            'column_name_mapping': cnm
-        }
+        state.code = sql
+        state.column_name_mapping.update(cnm)
 
-    def _all_dc_magic(self, silent: bool, value: str):
+    def _all_dc_magic(self, silent: bool, state: MagicState, value: str):
         if value.lower() in ('1', 'on', 'true'):
             self._magics['dc'].default(True)
             self._magics['ra'].default(False)
@@ -565,41 +587,43 @@ class DuckDBKernel(Kernel):
         else:
             self._magics['dc'].default(False)
 
-    def _guess_parser_magic(self, silent: bool, value: str):
+    def _guess_parser_magic(self, silent: bool, state: MagicState, value: str):
         if value.lower() in ('1', 'on', 'true'):
             self._magics['auto_parser'].default(True)
             self.print('The correct parser is guessed for each subsequently executed cell.\n')
         else:
             self._magics['auto_parser'].default(False)
 
-    def _auto_parser_magic(self, silent: bool, code: str):
+    def _auto_parser_magic(self, silent: bool, state: MagicState):
         # do not handle statements starting with SQL keywords
-        first_word = code.strip().split(maxsplit=1)
+        first_word = state.code.strip().split(maxsplit=1)
         if len(first_word) > 0:
             if first_word[0].upper() in SQL_KEYWORDS:
                 return
 
         # try to parse DC
         try:
-            return self._dc_magic(silent, code)
+            self._dc_magic(silent, state)
         except ParserError as e:
             if e.depth > 0:
                 raise e
 
         # try to parse RA
         try:
-            return self._ra_magic(silent, code, analyze=False)
+            self._ra_magic(silent, state, analyze=False)
         except ParserError as e:
             if e.depth > 0:
                 raise e
 
-    def _plotly_magic(self, silent: bool, cols: List, rows: List[Tuple], type: str, mapping: str, title: str = None):
+    def _plotly_magic(self, silent: bool, state: MagicState,
+                      cols: List, rows: List[Tuple],
+                      type: str, mapping: str, title: str = None):
         # split mapping and handle asterisks
         mapping = [m.strip() for m in mapping.split(',')]
 
         for i in range(len(mapping)):
             if mapping[i] == '*':
-                mapping = mapping[:i] + cols + mapping[i+1:]
+                mapping = mapping[:i] + cols + mapping[i + 1:]
 
         # convert all column names to lower case
         lower_cols = [c.lower() for c in cols]
@@ -668,7 +692,9 @@ class DuckDBKernel(Kernel):
         # finally print the code
         self.print_data(html, mime='text/html')
 
-    def _plotly_raw_magic(self, silent: bool, cols: List, rows: List[Tuple], title: str = None):
+    def _plotly_raw_magic(self, silent: bool, state: MagicState,
+                          cols: List, rows: List[Tuple],
+                          title: str = None):
         if len(cols) != 1 and len(rows) != 1:
             raise ValueError(f'expected exactly one column and one row')
 
@@ -683,34 +709,27 @@ class DuckDBKernel(Kernel):
                    **kwargs):
         try:
             # get magic command
-            clean_code, pre_query_callbacks, post_query_callbacks = self._magics(silent, code)
+            if len(self._db) > 0:
+                init_db = self._db[list(self._db.keys())[0]]
+            else:
+                init_db = None
 
-            # execute magic commands here if it does not depend on query results
-            execution_args = {
-                'max_rows': DuckDBKernel.DEFAULT_MAX_ROWS
-            }
+            magic_state = MagicState(init_db, code, DuckDBKernel.DEFAULT_MAX_ROWS)
+            pre_query_callbacks, post_query_callbacks = self._magics(silent, magic_state)
 
+            # execute magic commands here if it does not depend on query results
             for callback in pre_query_callbacks:
-                execution_args.update(callback())
-
-            # overwrite clean_code with generated code
-            if 'generated_code' in execution_args:
-                clean_code = execution_args['generated_code']
-                del execution_args['generated_code']
-
-            # set default column name mapping if none provided
-            if 'column_name_mapping' not in execution_args:
-                execution_args['column_name_mapping'] = {}
+                callback()
 
             # execute statement if needed
             cols, rows = None, None
 
-            if not isinstance(clean_code, dict):
-                clean_code = {'default': clean_code}
+            if not isinstance(magic_state.code, dict):
+                magic_state.code = {'default': magic_state.code}
 
-            for name, code in clean_code.items():
+            for name, code in magic_state.code.items():
                 if code.strip():
-                    cols, rows = self._execute_stmt(name, code, silent, **execution_args)
+                    cols, rows = self._execute_stmt(silent, magic_state, name, code)
 
             # execute magic command here if it does depend on query results
             for callback in post_query_callbacks:
@@ -734,5 +753,7 @@ class DuckDBKernel(Kernel):
             }
 
     def do_shutdown(self, restart):
-        self._unload_database()
+        for name in list(self._db.keys()):
+            self._unload_database(name)
+
         return super().do_shutdown(restart)

duckdb_kernel/magics/MagicCommandCallback.py

@@ -1,14 +1,13 @@
 from typing import Optional, List
 
-from . import MagicCommand
-from .StringWrapper import StringWrapper
+from . import MagicCommand, MagicState
 
 
 class MagicCommandCallback:
-    def __init__(self, mc: MagicCommand, silent: bool, code: StringWrapper, *args, **kwargs):
+    def __init__(self, mc: MagicCommand, silent: bool, state: MagicState, *args, **kwargs):
         self._mc: MagicCommand = mc
         self._silent: bool = silent
-        self._code: StringWrapper = code
+        self._state: MagicState = state
         self._args = args
         self._kwargs = kwargs
 
@@ -18,14 +17,8 @@ class MagicCommandCallback:
 
     def __call__(self, columns: Optional[List[str]] = None, rows: Optional[List[List]] = None):
         if self._mc.requires_code:
-            result = self._mc(self._silent, self._code.value, *self._args, **self._kwargs)
-            if 'generated_code' in result:
-                self._code.value = result['generated_code']
-
-            return result
-
-        if self._mc.requires_query_result:
-            return self._mc(self._silent, columns, rows, *self._args, **self._kwargs)
-
+            self._mc(self._silent, self._state, *self._args, **self._kwargs)
+        elif self._mc.requires_query_result:
+            self._mc(self._silent, self._state, columns, rows, *self._args, **self._kwargs)
         else:
-            return self._mc(self._silent, *self._args, **self._kwargs)
+            self._mc(self._silent, self._state, *self._args, **self._kwargs)

duckdb_kernel/magics/MagicCommandHandler.py

@@ -2,7 +2,8 @@ import re
 from typing import Dict, Tuple, List
 
 from . import MagicCommand, MagicCommandException, MagicCommandCallback
-from .StringWrapper import StringWrapper
+from .MagicState import MagicState
+from ..db import Connection
 
 
 class MagicCommandHandler:
@@ -18,8 +19,8 @@ class MagicCommandHandler:
     def __getitem__(self, key: str) -> MagicCommand:
         return self._magics[key.lower()]
 
-    def __call__(self, silent: bool, code: str) -> Tuple[str, List[MagicCommandCallback], List[MagicCommandCallback]]:
-        code_wrapper = StringWrapper()
+    def __call__(self, silent: bool, state: MagicState) \
+            -> Tuple[List[MagicCommandCallback], List[MagicCommandCallback]]:
         enabled_callbacks: List[MagicCommandCallback] = []
 
         # enable commands with default==True
@@ -27,21 +28,21 @@ class MagicCommandHandler:
             if magic.is_default:
                 flags = {name: False for name, _ in magic.flags}
                 optionals = {name: default for name, default, _ in magic.optionals}
-                callback = MagicCommandCallback(magic, silent, code_wrapper, **flags, **optionals)
+                callback = MagicCommandCallback(magic, silent, state, **flags, **optionals)
 
                 enabled_callbacks.append(callback)
 
         # search for magic commands in code
         while True:
             # ensure code starts with '%' or '%%' but not with '%%%'
-            match = re.match(r'^%{1,2}([^% ]+?)([ \t]*$| .+?$)', code, re.MULTILINE | re.IGNORECASE)
+            match = re.match(r'^%{1,2}([^% ]+?)([ \t]*$| .+?$)', state.code, re.MULTILINE | re.IGNORECASE)
 
             if match is None:
                 break
 
             # remove magic command from code
             start, end = match.span()
-            code = code[:start] + code[end + 1:]
+            state.code = state.code[:start] + state.code[end + 1:]
 
             # extract command
             command = match.group(1).lower()
@@ -99,7 +100,7 @@ class MagicCommandHandler:
                     optionals[name.lower()] = value
 
             # add to callbacks
-            callback = MagicCommandCallback(magic, silent, code_wrapper, *args, **flags, **optionals)
+            callback = MagicCommandCallback(magic, silent, state, *args, **flags, **optionals)
             enabled_callbacks.append(callback)
 
         # disable overwritten callbacks
@@ -129,5 +130,4 @@ class MagicCommandHandler:
                 post_query_callbacks.append(callback)
 
         # return callbacks
-        code_wrapper.value = code
-        return code, pre_query_callbacks, post_query_callbacks
+        return pre_query_callbacks, post_query_callbacks

duckdb_kernel/magics/MagicState.py

@@ -0,0 +1,11 @@
+from typing import Union, Dict, Optional
+
+from ..db import Connection
+
+
+class MagicState:
+    def __init__(self, db: Connection, code: str, max_rows: Optional[int]):
+        self.db: Connection = db
+        self.code: Union[str, Dict] = code
+        self.max_rows: Optional[int] = max_rows
+        self.column_name_mapping: Dict[str, str] = {}

duckdb_kernel/magics/__init__.py

@@ -2,3 +2,4 @@ from .MagicCommand import MagicCommand
 from .MagicCommandCallback import MagicCommandCallback
 from .MagicCommandException import MagicCommandException
 from .MagicCommandHandler import MagicCommandHandler
+from .MagicState import MagicState

{jupyter_duckdb-1.2.102.dist-info → jupyter_duckdb-1.2.103.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: jupyter-duckdb
-Version: 1.2.102
+Version: 1.2.103
 Summary: a basic wrapper kernel for DuckDB
 Home-page: https://github.com/erictroebs/jupyter-duckdb
 Author: Eric Tröbs
@@ -118,23 +118,24 @@ another DuckDB file or a file with SQL statements. In the first case the
 included tables will be copied to the new database, while in the second case the
 SQL statements are just executed. We find this feature very useful to work in a
 temporary copy of the data and therefore be able to restart at any time. The
-last optional parameter `WITH_TESTS` is described in
-detail [below](#ship-tests-with-your-notebooks).
+optional parameter `NAME` may be used to name a connection and reference it
+later by using the magic command `USE`.
 
 ```
 %CREATE data.duckdb OF my_statements.sql
 ```
 
 `LOAD` on the other hand loads an existing database and returns an error if it
-does not exist. (That is why `OF` cannot be used with `LOAD`! `WITH_TESTS` on
-the other hand is available also with this magic command.)
+does not exist. (That is why `OF` cannot be used with `LOAD`! `NAME` on the
+other hand is available also with this magic command.)
 
 ```
 %LOAD data.duckdb
 ```
 
-Only one database can be open at any time. If a new database is created or
-loaded, the current one is closed first and saved to disk if necessary.
+Multiple databases can be open at any time. If a new database with the same
+name is created or loaded, the current one is closed first and saved to disk
+if necessary.
 
 Please note that `:memory:` is also a valid file path for DuckDB. The data is
 then stored exclusively in the main memory. In combination with `CREATE`
@@ -213,10 +214,10 @@ FROM bar
 
 ### Ship Tests With Your Notebooks
 
-Simple tests can be loaded together with the database with the help of
-the `WITH_TESTS` parameter. These tests are stored as a JSON file. Each test is
-assigned a unique name, a result set and whether the test should check the order
-of the result. A very simple test file looks like the following JSON object:
+Simple tests can be loaded from json files with the help of magic command
+`LOAD_TESTS`. These tests are stored as a JSON file. Each test is assigned a
+unique name, a result set and whether the test should check the order of the
+result. A very simple test file looks like the following JSON object:
 
 ```json
 {

{jupyter_duckdb-1.2.102.dist-info → jupyter_duckdb-1.2.103.dist-info}/RECORD

@@ -1,9 +1,9 @@
 duckdb_kernel/__init__.py,sha256=6auU6zeJrsA4fxPSr2PYamS8fG-SMXTn5YQFXF2cseo,33
 duckdb_kernel/__main__.py,sha256=Z3GwHEBWoQjNm2Y84ijnbA0Lk66L7nsFREuqhZ_ptk0,165
 duckdb_kernel/kernel.json,sha256=_7E8Ci2FSdCvnzCjsOaue8QE8AvpS5JLQuxORO5IGtA,127
-duckdb_kernel/kernel.py,sha256=DIixbtvXldvVlmf_6jgUWWC_Mu5UY82oebHaJCrH90Y,27236
+duckdb_kernel/kernel.py,sha256=qb2dEj2XhNE_9-xlpQBky7Cd4km_dqRLTXfNaNQNdpM,28189
 duckdb_kernel/db/Column.py,sha256=GM5P6sFdlYK92hiKln5-6038gIDOTxh1AYbR4kiga_w,559
-duckdb_kernel/db/Connection.py,sha256=5pH-CwGh-r9Q2QwJKGSxvoINBU-sqmvZyG4Q1digfeE,599
+duckdb_kernel/db/Connection.py,sha256=tBXQBYt9c52RLbpl9sakNuAm0Z84--fhZ4efo8ACz-U,670
 duckdb_kernel/db/Constraint.py,sha256=1YgUHk7s8mHCVedbcuJKyXDykj7_ybbwT3Dk9p2VMis,287
 duckdb_kernel/db/DatabaseError.py,sha256=43zl8yym1f-fxH_UtGIbWnDnBE_TRwr9aCziY9t40QY,41
 duckdb_kernel/db/ForeignKey.py,sha256=iurUAXwTwSIpLXsL0B7BA8jqDTfW4_wkeHxoqQbZwiU,470
@@ -11,19 +11,19 @@ duckdb_kernel/db/Table.py,sha256=HfvGX54kD_XvmLApYSmxtTQNvz2YYaaUNpm4e8dSOVY,934
 duckdb_kernel/db/__init__.py,sha256=PKQYQDCW7VQYxmzhQK6A0Qloka9FdMfeFQMfY-CjBSA,198
 duckdb_kernel/db/error/EmptyResultError.py,sha256=N9Oxi2HDZBKaRQsfRsWpJJGOYX4BjdQqWOU-XvzUzNY,92
 duckdb_kernel/db/error/__init__.py,sha256=oHfhfbcfyTJ3pAPN835omdQcebvJTauuULFx5gm9rq4,47
-duckdb_kernel/db/implementation/duckdb/Connection.py,sha256=uBL-UVrSo8EJNXnnyAc2AwAYJ1QQQmH0noV6pnngZHs,6977
+duckdb_kernel/db/implementation/duckdb/Connection.py,sha256=nFzj2wjHORFuRn27zfKhHtwu1Hulph7X4ImCaEbQtbQ,7051
 duckdb_kernel/db/implementation/duckdb/__init__.py,sha256=HKogB1es4wOiQUoh7_eT32xnUFLmzoCyR_0LuY9r8YQ,35
-duckdb_kernel/db/implementation/postgres/Connection.py,sha256=39wv-mvKHdu4u_ADFiSbAvGMVEs3FtuzRYIH4uzJ-pw,9307
+duckdb_kernel/db/implementation/postgres/Connection.py,sha256=WCCdkM4V-Y2Pg-ofeTLHpyT7jgbQaw7qnaW5M1bYbX4,9490
 duckdb_kernel/db/implementation/postgres/__init__.py,sha256=HKogB1es4wOiQUoh7_eT32xnUFLmzoCyR_0LuY9r8YQ,35
 duckdb_kernel/db/implementation/postgres/util.py,sha256=4nr1mqXhlwkMVXbJSfJ7dRlUm6UskpvgKApe7GRwmBI,281
-duckdb_kernel/db/implementation/sqlite/Connection.py,sha256=SSLpdigWjO_3lUXdsx3eH5OXOTBlZiXOgIyPCf6L8D0,7044
+duckdb_kernel/db/implementation/sqlite/Connection.py,sha256=W_7Eb2u2dHl54h5FXGQVc9oIFa6czYMhkNc2UWbTrkE,7118
 duckdb_kernel/db/implementation/sqlite/__init__.py,sha256=HKogB1es4wOiQUoh7_eT32xnUFLmzoCyR_0LuY9r8YQ,35
 duckdb_kernel/magics/MagicCommand.py,sha256=l0EmnqgGZ0HyqQhdTljAaftflVo_RYp-U5UiDftYxAA,3180
-duckdb_kernel/magics/MagicCommandCallback.py,sha256=9XVidNtzs8Lwq_T59VrH89t5LUgJcc5C7grusElXVW4,1041
+duckdb_kernel/magics/MagicCommandCallback.py,sha256=Otl7Sa53eJtkMU5gBup3VTVCQfr7p00kt-xpnlETVe4,877
 duckdb_kernel/magics/MagicCommandException.py,sha256=MwuWkpA6NoCqz437urdI0RVXhbSbVdziuRoi7slYFPc,49
-duckdb_kernel/magics/MagicCommandHandler.py,sha256=noRm22EsfTZI_FL_N15f6NYbZSPvejvS24rlA5hkDkQ,4730
-duckdb_kernel/magics/StringWrapper.py,sha256=W6qIfeHU51do1edd6yXgw6K7orzjwSHU4oWAI5DuKEE,96
-duckdb_kernel/magics/__init__.py,sha256=DA8gnQeRCUt1Scy3_NQ9w5CPmMEY9i8YwB-g392pN1U,204
+duckdb_kernel/magics/MagicCommandHandler.py,sha256=hKNHRfa0BSPei4QXdVlLlbpvZIhJJ91aaT_HTMiK28E,4700
+duckdb_kernel/magics/MagicState.py,sha256=Vt_KwUwQP3446c1snSxS68Skl5AZQzrJd2Q1ETpKuKI,344
+duckdb_kernel/magics/__init__.py,sha256=ggxzDzDEsKMZzYsWw9JqYVJhciJPvPVYGV7oNo9Yp-E,239
 duckdb_kernel/parser/DCParser.py,sha256=16c1mxa494KP9OreUKQHsSQKoDGZ7NNp2u_gi_D-dkw,2293
 duckdb_kernel/parser/LogicParser.py,sha256=_vZwE5OPRUEN8aEC_fSZAYKR_dpexqNthXog9OFHYRY,1233
 duckdb_kernel/parser/ParserError.py,sha256=qJQVloFtID1HgVDQ1Io247bODT1ic3oO9Z1ZrWR-2Mk,321
@@ -93,7 +93,7 @@ duckdb_kernel/visualization/lib/__init__.py,sha256=LYi0YPtn5fXOejbLIqbt_3KzP-Xrw
 duckdb_kernel/visualization/lib/plotly-3.0.1.min.js,sha256=oy6Be7Eh6eiQFs5M7oXuPxxm9qbJXEtTpfSI93dW16Q,4653932
 duckdb_kernel/visualization/lib/ra.css,sha256=foz1v69EQ117BDduB9QyHH978PbRs2TG1kBS4VGqZbI,57
 duckdb_kernel/visualization/lib/ra.js,sha256=VzMRn55ztcd5Kfu2B6gdRPARpi8n-fvs8oNFnfp55Ec,1845
-jupyter_duckdb-1.2.102.dist-info/METADATA,sha256=020EZdonW3kvcFer9t5LsfTlgDI5s1WhgJyvzgAYJjk,9115
-jupyter_duckdb-1.2.102.dist-info/WHEEL,sha256=beeZ86-EfXScwlR_HKu4SllMC9wUEj_8Z_4FJ3egI2w,91
-jupyter_duckdb-1.2.102.dist-info/top_level.txt,sha256=KvRRPMnmkQNuhyBsXoPmwyt26LRDp0O-0HN6u0Dm5jA,14
-jupyter_duckdb-1.2.102.dist-info/RECORD,,
+jupyter_duckdb-1.2.103.dist-info/METADATA,sha256=KF4yFiFhppasCLN85BH8KVdKBOudeEbFIkplYy08rMg,9132
+jupyter_duckdb-1.2.103.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+jupyter_duckdb-1.2.103.dist-info/top_level.txt,sha256=KvRRPMnmkQNuhyBsXoPmwyt26LRDp0O-0HN6u0Dm5jA,14
+jupyter_duckdb-1.2.103.dist-info/RECORD,,

{jupyter_duckdb-1.2.102.dist-info → jupyter_duckdb-1.2.103.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (76.1.0)
+Generator: setuptools (77.0.3)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

duckdb_kernel/magics/StringWrapper.py

@@ -1,3 +0,0 @@
-class StringWrapper:
-    def __init__(self, value: str = None):
-        self.value: str = value