jupyter-duckdb 1.2.0.1__py3-none-any.whl → 1.4.111__py3-none-any.whl

duckdb_kernel/parser/util/QuerySplitter.py

@@ -0,0 +1,87 @@
+from typing import Iterator
+
+
+def split_queries(query: str, split_at: str | None = ';', remove_comments: bool = False) -> Iterator[str]:
+    quotes = '\'"`'
+
+    escaped = False
+    in_quotes = None
+    in_singleline_comment = False
+    in_multiline_comment = False
+
+    previous = None
+    current_query = []
+
+    for symbol in query:
+        keep_symbol = True
+
+        # escaped symbol
+        if escaped:
+            escaped = False
+
+        # backslash (escape)
+        elif symbol == '\\':
+            escaped = True
+
+        # if in quotes
+        elif in_quotes is not None:
+            if symbol == in_quotes:
+                in_quotes = False
+
+        # if in single line comment
+        elif in_singleline_comment:
+            if symbol == '\n':
+                in_singleline_comment = False
+            elif remove_comments:
+                keep_symbol = False
+
+        # if in multiline comment
+        elif in_multiline_comment:
+            if previous == '*' and symbol == '/':
+                in_multiline_comment = False
+
+            if remove_comments:
+                keep_symbol = False
+
+        # start of quotes
+        elif symbol in quotes:
+            in_quotes = symbol
+
+        # start of single line comment
+        elif previous == '-' and symbol == '-':
+            in_singleline_comment = True
+
+            if remove_comments:
+                keep_symbol = False
+                current_query.pop()
+
+        # start of multiline comment
+        elif previous == '/' and symbol == '*':
+            in_multiline_comment = True
+
+            if remove_comments:
+                keep_symbol = False
+                current_query.pop()
+
+        # semicolon
+        elif split_at is not None and symbol == split_at:
+            yield ''.join(current_query)
+
+            current_query = []
+            keep_symbol = False
+
+        # store symbol
+        if keep_symbol:
+            current_query.append(symbol)
+
+        previous = symbol
+
+    # yield remaining symbols
+    yield ''.join(current_query)
+
+
+def get_last_query(query: str, split_at: str | None = ';', remove_comments: bool = False) -> str:
+    for query in split_queries(query, split_at, remove_comments):
+        pass
+
+    return query

duckdb_kernel/parser/util/RenamableColumnList.py

@@ -73,18 +73,26 @@ class RenamableColumnList(list[RenamableColumn]):
 
         return RenamableColumnList(cols.values())
 
-    def intersect(self, other: 'RenamableColumnList') \
+    def intersect(self, other: 'RenamableColumnList', prefer_right: bool = False) \
             -> Tuple[List[Tuple[RenamableColumn, RenamableColumn]], 'RenamableColumnList']:
         self_cols: Dict[str, RenamableColumn] = {col.name: col for col in self}
         other_cols: Dict[str, RenamableColumn] = {col.name: col for col in other}
 
+        replacements: Dict[RenamableColumn, RenamableColumn] = {}
         intersection: List[Tuple[RenamableColumn, RenamableColumn]] = []
+
         for name in tuple(self_cols.keys()):
             if name in other_cols:
+                if prefer_right:
+                    replacements[self_cols[name]] = other_cols[name]
+
                 intersection.append((self_cols[name], other_cols[name]))
                 del other_cols[name]
 
         # if len(intersection) == 0:
         #     raise AssertionError('no common attributes found for join')
 
-        return intersection, RenamableColumnList(self_cols.values(), other_cols.values())
+        return intersection, RenamableColumnList(
+            (replacements.get(x, x) for x in self_cols.values()),
+            other_cols.values()
+        )

duckdb_kernel/tests/__init__.py

@@ -0,0 +1,76 @@
+import os
+from typing import Dict, List, Tuple
+
+from duckdb_kernel.db import Connection as DB, Table
+from duckdb_kernel.db.error import EmptyResultError
+from duckdb_kernel.parser.elements import RAElement
+from duckdb_kernel.parser.elements.binary import ConditionalSet
+
+with open('examples/tables.sql', 'r', encoding='utf-8') as file:
+    EXAMPLE_STMTS = [stmt
+                     for stmt in file.read().split(';')
+                     if stmt.strip()]
+
+
+class Connection:
+    def __enter__(self):
+        db_type = os.environ.get('DB_TYPE')
+        if db_type == 'postgres':
+            host = os.environ.get('POSTGRES_HOST', 'localhost')
+            port = int(os.environ.get('POSTGRES_PORT', 5432))
+            username = os.environ.get('POSTGRES_USER', 'postgres')
+            password = os.environ.get('POSTGRES_PASSWORD', 'postgres')
+
+            from duckdb_kernel.db.implementation.postgres import Connection as PostgreSQL
+            self.con: DB = PostgreSQL(host, port, username, password, None)
+        elif db_type == 'sqlite':
+            from duckdb_kernel.db.implementation.sqlite import Connection as SQLite
+            self.con: DB = SQLite(':memory:')
+        else:
+            from duckdb_kernel.db.implementation.duckdb import Connection as DuckDB
+            self.con: DB = DuckDB(':memory:')
+
+        for stmt in EXAMPLE_STMTS:
+            try:
+                self.con.execute(stmt)
+            except EmptyResultError:
+                pass
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.con.close()
+
+    @property
+    def tables(self) -> Dict[str, Table]:
+        return self.con.analyze()
+
+    def execute_sql_return_cols(self, query: str) -> Tuple[List, List]:
+        return self.con.execute(query)
+
+    def execute_sql(self, query: str) -> List:
+        _, rows = self.execute_sql_return_cols(query)
+        return rows
+
+    def execute_ra_return_cols(self, root: RAElement) -> Tuple[List, List]:
+        sql = root.to_sql_with_renamed_columns(self.tables)
+        cols, rows = self.execute_sql_return_cols(sql)
+
+        return cols, rows  # sorted(rows, key=lambda t: tuple(-1 if x is None else x for x in t))
+
+    def execute_ra(self, root: RAElement) -> List:
+        _, rows = self.execute_ra_return_cols(root)
+        return rows
+
+    def execute_dc_return_cols(self, root: ConditionalSet) -> Tuple[List, List]:
+        sql, cnm = root.to_sql_with_renamed_columns(self.tables)
+        cols, rows = self.execute_sql_return_cols(sql)
+
+        return [cnm.get(c, c) for c in cols], rows  # sorted(rows)
+
+    def execute_dc(self, root: ConditionalSet) -> List:
+        _, rows = self.execute_dc_return_cols(root)
+        return rows
+
+
+__all__ = ['Connection']

duckdb_kernel/tests/test_dc.py

@@ -0,0 +1,483 @@
+from duckdb_kernel.parser import DCParser
+from . import Connection
+
+
+def test_case_insensitivity():
+    for query in (
+            '{ username | users(id, username) }',
+            '{ username | Users(id, username) }',
+            '{ username | USERS(id, username) }',
+            '{ username | uSers(id, username) }',
+    ):
+        root = DCParser.parse_query(query)
+
+        # execute to test case insensitivity
+        with Connection() as con:
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'username'
+            ]
+            assert rows == [
+                ('Alice',),
+                ('Bob',),
+                ('Charlie',)
+            ]
+
+    for query in (
+            '{ username | users(id, username) }',
+            '{ Username | users(id, username) }',
+            '{ USERNAME | users(id, username) }',
+            '{ userName | users(id, username) }',
+    ):
+        root = DCParser.parse_query(query)
+
+        # execute to test case insensitivity
+        with Connection() as con:
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'username'
+            ]
+            assert rows == [
+                ('Alice',),
+                ('Bob',),
+                ('Charlie',)
+            ]
+
+
+def test_comments():
+    for query in (
+            '{ username | users(id, username) } -- cmmnt',
+            '{ username | users(id, username) } /* cmmnt */',
+            '{ username |  /* cmmnt */ users(id, username) }',
+    ):
+        root = DCParser.parse_query(query)
+
+        # execute to test case insensitivity
+        with Connection() as con:
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'username'
+            ]
+            assert rows == [
+                ('Alice',),
+                ('Bob',),
+                ('Charlie',)
+            ]
+
+    for query in (
+        '-- comment',
+        '/* comment */'
+    ):
+        root = DCParser.parse_query(query)
+        assert root is None
+
+
+def test_simple_queries():
+    with Connection() as con:
+        for query in [
+            '{ id | Users(id, _) }',
+            '{ id | Users (id, _) }',
+            '{ id | Users id,_ }',
+            '{id|Users(id,_)}',
+            'id | Users (id ,_)',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'id'
+            ]
+            assert rows == [
+                (1,),
+                (2,),
+                (3,)
+            ]
+
+        for query in [
+            '{ id, name | Users(id, name) }',
+            '{ id,name | Users(id, name) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'id',
+                'name'
+            ]
+            assert rows == [
+                (1, 'Alice'),
+                (2, 'Bob'),
+                (3, 'Charlie')
+            ]
+
+
+def test_asterisk_projection():
+    with Connection() as con:
+        root = DCParser.parse_query('{ * | Users(id, _) }')
+        cols, rows = con.execute_dc_return_cols(root)
+
+        assert [c.lower() for c in cols] == [
+            'id'
+        ]
+        assert rows == [
+            (1,),
+            (2,),
+            (3,)
+        ]
+
+        root = DCParser.parse_query('{ * | Users(id, name) }')
+        cols, rows = con.execute_dc_return_cols(root)
+
+        assert [c.lower() for c in cols] == [
+            'id',
+            'name'
+        ]
+        assert rows == [
+            (1, 'Alice'),
+            (2, 'Bob'),
+            (3, 'Charlie')
+        ]
+
+
+def test_conditions():
+    with Connection() as con:
+        for query in [
+            '{ name | Users(id, name) ∧ id > 1 }',
+            '{ name | Users(id, name) ∧ id ≠ 1 }',
+            '{ name | Users(id, name) ∧ (id = 2 ∨ id = 3) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'name'
+            ]
+            assert rows == [
+                ('Bob',),
+                ('Charlie',)
+            ]
+
+        for query in [
+            '{ id | Users(id, name) ∧ name = "Bob" }',
+            '{ id | Users(id, name) ∧ name > "B" ∧ name < "C" }',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'id'
+            ]
+            assert rows == [
+                (2,)
+            ]
+
+
+def test_shortcut_conditions():
+    with Connection() as con:
+        # single shortcut conditions
+        for query in [
+            '{ name | Users(1, name) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'name'
+            ]
+            assert rows == [
+                ('Alice',)
+            ]
+
+        for query in [
+            '{ season_name | Seasons(1, 1, season_name) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'season_name'
+            ]
+            assert rows == [
+                ('Show 1 / Season 1',)
+            ]
+
+        # multiple shortcut conditions
+        for query in [
+            '{ sname, ename | Seasons(snum, 2, sname) ∧ Episodes(enum, snum, 2, ename) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'sname',
+                'ename'
+            ]
+            assert rows == [
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 1'),
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 2'),
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 3'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 1'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 2'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 3'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 4')
+            ]
+
+
+def test_joins():
+    with Connection() as con:
+        # with one attribute
+        for query in [
+            '{ sename | Shows(shid, shname) ∧ Seasons(senum, shid, sename) }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'sename'
+            ]
+            assert rows == [
+                ('Show 1 / Season 1',),
+                ('Show 1 / Season 2',),
+                ('Show 2 / Season 1',),
+                ('Show 2 / Season 2',)
+            ]
+
+        for query in [
+            '{ sename | Shows(shid, shname) ∧ Seasons(senum, shid, sename) ∧ shname = "Show 1" }',
+            '{ sename | Seasons(senum, shid, sename) ∧ Shows(shid, "Show 1") }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'sename'
+            ]
+            assert rows == [
+                ('Show 1 / Season 1',),
+                ('Show 1 / Season 2',)
+            ]
+
+        # with multiple attributes
+        for query in [
+            '{ sname, ename | Seasons(snum, shid, sname) ∧ Episodes(enum, snum, shid, ename) ∧ shid = 2 }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'sname',
+                'ename'
+            ]
+            assert rows == [
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 1'),
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 2'),
+                ('Show 2 / Season 1', 'Show 2 / Season 1 / Episode 3'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 1'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 2'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 3'),
+                ('Show 2 / Season 2', 'Show 2 / Season 2 / Episode 4')
+            ]
+
+        # join three relations
+        for query in [
+            '{ s2,c5 | Shows(s1,s2) ∧ Episodes(e1,e2,s1,e4) ∧ Characters(c1,e1,c3,s1,c5) ∧ s1=2 ∧ e4="Show 2 / Season 1 / Episode 2" }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                's2',
+                'c5'
+            ]
+            assert rows == [
+                ('Show 2', 'Actor F')
+            ]
+
+        # cross join
+        root = DCParser.parse_query('{ sename | Shows(shid1, shname) ∧ Seasons(senum, shid2, sename) ∧ shid1 = shid2 }')
+        cols, rows = con.execute_dc_return_cols(root)
+
+        assert [c.lower() for c in cols] == [
+            'sename'
+        ]
+        assert rows == [
+            ('Show 1 / Season 1',),
+            ('Show 1 / Season 2',),
+            ('Show 2 / Season 1',),
+            ('Show 2 / Season 2',)
+        ]
+
+        for query in [
+            '{ s2,c5 | Shows(sa1,s2) ∧ Episodes(e1,e2,sb1,e4) ∧ Characters(c1,e1,c3,sb1,c5) ∧ sa1=2 ∧ sa1 = sb1 ∧ e4="Show 2 / Season 1 / Episode 2" }',
+            '{ s2,c5 | Shows(sa1,s2) ∧ Episodes(e1,e2,sb1,e4) ∧ Characters(c1,e1,c3,sc1,c5) ∧ sa1=2 ∧ sa1 = sb1 ∧ sb1 = sc1 ∧ e4="Show 2 / Season 1 / Episode 2" }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                's2',
+                'c5'
+            ]
+            assert rows == [
+                ('Show 2', 'Actor F')
+            ]
+
+
+def test_disjunction_joins():
+    with Connection() as con:
+        for query in [
+            "{ enum, snum, sid | Episodes(enum, snum, sid, ename) ∧ (Characters('Character B', enum, snum, sid, _) ∨ Characters('Character D', enum, snum, sid, _)) }",
+            "{ enum, snum, sid | Episodes(enum, snum, sid, ename) ∧ (Characters(cname1, enum, snum, sid, _) ∧ cname1 = 'Character B' ∨ Characters(cname2, enum, snum, sid, _) ∧ cname2 = 'Character D') }",
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'enum',
+                'snum',
+                'sid'
+            ]
+            assert rows == [
+                (1, 1, 1),
+                (2, 1, 1)
+            ]
+
+
+def test_cross_join():
+    with Connection() as con:
+        for query in [
+            "{ cname1, cname2 | Characters(cname1, _, _, 2, _) ∧ Characters(cname2, _, _, 2, _) }",
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'cname1',
+                'cname2'
+            ]
+            assert rows == [
+                ('Character E', 'Character E'),
+                ('Character E', 'Character F'),
+                ('Character F', 'Character E'),
+                ('Character F', 'Character F'),
+            ]
+
+
+def test_underscores():
+    with Connection() as con:
+        # distinct underscores
+        for query in [
+            '{ ename | Seasons(snum, shid, sname) ∧ Episodes(_, snum, shid, ename) ∧ shid = 2 }',
+            '{ ename | Seasons(snum, shid, sname) ∧ Episodes(enum, _, shid, ename) ∧ shid = 2 }',
+            '{ ename | Seasons(snum, shid, sname) ∧ Episodes(__, snum, shid, ename) ∧ shid = 2 }',
+            '{ ename | Seasons(snum, shid, sname) ∧ Episodes(_, __, shid, ename) ∧ shid = 2 }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'ename'
+            ]
+            assert rows == [
+                ('Show 2 / Season 1 / Episode 1',),
+                ('Show 2 / Season 1 / Episode 2',),
+                ('Show 2 / Season 1 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 1',),
+                ('Show 2 / Season 2 / Episode 2',),
+                ('Show 2 / Season 2 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 4',)
+            ]
+
+        # reused underscores in a single relation
+        for query in [
+            '{ ename | Seasons(snum, shid, sname) ∧ Episodes(_, _, shid, ename) ∧ shid = 2 }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'ename'
+            ]
+            assert rows == [
+                ('Show 2 / Season 1 / Episode 1',),
+                ('Show 2 / Season 1 / Episode 2',),
+                ('Show 2 / Season 1 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 1',),
+                ('Show 2 / Season 2 / Episode 2',),
+                ('Show 2 / Season 2 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 4',)
+            ]
+
+        # reused underscores in two different relations
+        for query in [
+            '{ ename | Seasons(_, shid, _) ∧ Episodes(_, _, shid, ename) ∧ shid = 2 }',
+            '{ ename | Seasons(_, shid, __) ∧ Episodes(_, __, shid, ename) ∧ shid = 2 }'
+        ]:
+            root = DCParser.parse_query(query)
+            cols, rows = con.execute_dc_return_cols(root)
+
+            assert [c.lower() for c in cols] == [
+                'ename'
+            ]
+            assert rows == [
+                ('Show 2 / Season 1 / Episode 1',),
+                ('Show 2 / Season 1 / Episode 2',),
+                ('Show 2 / Season 1 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 1',),
+                ('Show 2 / Season 2 / Episode 2',),
+                ('Show 2 / Season 2 / Episode 3',),
+                ('Show 2 / Season 2 / Episode 4',)
+            ]
+
+def test_anonymous_column_names():
+    with Connection() as con:
+        for query in [
+            '{ * | Episodes(_, _, 2, ename) }',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ['2', 'ename']
+
+        for query in [
+            "{ * | Episodes(_, _, '2', ename) }",
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ["'2'", 'ename']
+
+        for query in [
+            "{ * | Episodes(_, _, sid, 'ename') }",
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ['sid', "'ename'"]
+
+        for query in [
+            '{ * | Episodes(_, 1, 2, ename) }',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ['1', '2', 'ename']
+
+        for query in [
+            '{ * | Episodes(_, 2, 2, ename) }',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ['2', '2', 'ename']
+
+        for query in [
+            '{ * | Episodes(_, _, sid, ename) ∧ sid = 2 }',
+        ]:
+            root = DCParser.parse_query(query)
+            cols, _ = con.execute_dc_return_cols(root)
+
+            assert cols == ['sid', 'ename']