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

duckdb_kernel/kernel.py

@@ -39,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),
@@ -59,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'):
@@ -97,131 +100,143 @@ 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')
-
-            # 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, 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
 
         et = time.time()
 
-        # return result if silent
-        if silent:
-            return columns, rows
-
-        # print EXPLAIN queries as raw text if using DuckDB
-        if query.strip().startswith('EXPLAIN') and self._db.plain_explain():
-            for ekey, evalue in rows:
-                self.print_data(f'<b>{ekey}</b><br><pre>{evalue}</pre>')
-
-            return None, None
-
-        # 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)
-                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:
-                    table_data = f'''
-                        {rows_table(rows[:math.ceil(max_rows / 2)])}
-                        <tr>
-                            <td colspan="{len(columns)}" 
-                                style="text-align: center"
-                                title="{row_count(len(rows) - max_rows)} omitted">
-                                ...
-                            </td>
-                        </tr>
-                        {rows_table(rows[-math.floor(max_rows // 2):])}
-                    '''
+        # print result if not silent
+        if not silent:
+            # print EXPLAIN queries as raw text if using DuckDB
+            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
                 else:
-                    table_data = rows_table(rows)
+                    html = ''
 
-                # send to client
-                self.print_data(f'''
-                    <table class="duckdb-query-result">
-                        {table_header}
-                        {table_data}
-                    </table>
-                ''')
+            # print every other query as a table
+            else:
+                if columns is not None:
+                    # table header
+                    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 state.max_rows is not None and len(rows) > state.max_rows:
+                        table_data = f'''
+                            {rows_table(rows[:math.ceil(state.max_rows / 2)])}
+                            <tr>
+                                <td colspan="{len(columns)}" 
+                                    style="text-align: center"
+                                    title="{row_count(len(rows) - state.max_rows)} omitted">
+                                    ...
+                                </td>
+                            </tr>
+                            {rows_table(rows[-math.floor(state.max_rows // 2):])}
+                        '''
+                    else:
+                        table_data = rows_table(rows)
+
+                    # send to client
+                    html = (f'''
+                        <table class="duckdb-query-result-table">
+                            {table_header}
+                            {table_data}
+                        </table>
+                        
+                        {row_count(len(rows))} in {et - st:.3f}s
+                    ''')
 
-                self.print_data(f'{row_count(len(rows))} in {et - st:.3f}s')
+                else:
+                    html = f'statement executed without result in {et - st:.3f}s'
 
-            else:
-                self.print_data(f'statement executed without result in {et - st:.3f}s')
+            self.print_data(f'''
+                <div class="duckdb-query-result {name}">
+                    {html}
+                </div>
+            ''')
 
         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')
 
@@ -238,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:
@@ -257,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
 
@@ -282,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}')
 
-                        test['equals'] = rows
+        state.db = self._db[name]
 
-                self.print(f'loaded tests from {with_tests}\n')
+    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)
 
-    def _test_magic(self, silent: bool, result_columns: List[str], result: List[List], name: str):
+                    test['equals'] = rows
+
+            self.print(f'loaded tests from {tests}\n')
+
+    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.'))
@@ -393,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:
@@ -459,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
@@ -479,37 +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)
-            svg = vd.to_svg(True)
+            vd = RATreeDrawer(state.db, root_node, tables)
 
+            svg = vd.to_interactive_svg()
             self.print_data(svg)
 
-        # generate sql
-        sql = 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 {
-            'generated_code': sql
-        }
+        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)
@@ -518,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)
@@ -551,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]
@@ -654,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')
 
@@ -669,30 +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())
+                callback()
 
-            # overwrite clean_code with generated code
-            if 'generated_code' in execution_args:
-                clean_code = execution_args['generated_code']
-                del execution_args['generated_code']
+            # execute statement if needed
+            cols, rows = None, None
 
-            # set default column name mapping if none provided
-            if 'column_name_mapping' not in execution_args:
-                execution_args['column_name_mapping'] = {}
+            if not isinstance(magic_state.code, dict):
+                magic_state.code = {'default': magic_state.code}
 
-            # execute statement if needed
-            if clean_code.strip():
-                cols, rows = self._execute_stmt(clean_code, silent, **execution_args)
-            else:
-                cols, rows = None, None
+            for name, code in magic_state.code.items():
+                if code.strip():
+                    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:
@@ -716,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)