25class AsyncConnection:
 26    """AsyncConnection enables asynchronous execution of queries with a pool of connections and threads."""
 27
 28    def __init__(
 29        self,
 30        database: Database,
 31        max_concurrent_queries: int = 4,
 32        max_threads_per_query: int = 0,
 33    ) -> None:
 34        """
 35        Initialise the async connection.
 36
 37        Parameters
 38        ----------
 39        database : Database
 40            Database to connect to.
 41
 42        max_concurrent_queries : int
 43            Maximum number of concurrent queries to execute. This corresponds to the
 44            number of connections and thread pool size. Default is 4.
 45
 46        max_threads_per_query : int
 47            Controls the maximum number of threads per connection that can be used
 48            to execute one query. Default is 0, which means no limit.
 49        """
 50        self.database = database
 51        self.connections = [Connection(database) for _ in range(max_concurrent_queries)]
 52        self.connections_counter = [0 for _ in range(max_concurrent_queries)]
 53        self.lock = threading.Lock()
 54
 55        for conn in self.connections:
 56            conn.init_connection()
 57            conn.set_max_threads_for_exec(max_threads_per_query)
 58
 59        self.executor = ThreadPoolExecutor(max_workers=max_concurrent_queries)
 60
 61    def __enter__(self) -> Self:
 62        return self
 63
 64    def __exit__(
 65        self,
 66        exc_type: type[BaseException] | None,
 67        exc_value: BaseException | None,
 68        exc_traceback: TracebackType | None,
 69    ) -> None:
 70        self.close()
 71
 72    def __del__(self) -> None:
 73        self.close()
 74
 75    def __get_connection_with_least_queries(self) -> tuple[Connection, int]:
 76        with self.lock:
 77            conn_index = self.connections_counter.index(min(self.connections_counter))
 78            self.connections_counter[conn_index] += 1
 79        return self.connections[conn_index], conn_index
 80
 81    def __decrement_connection_counter(self, conn_index: int) -> None:
 82        """Decrement the query counter for a connection."""
 83        with self.lock:
 84            self.connections_counter[conn_index] -= 1
 85            if self.connections_counter[conn_index] < 0:
 86                self.connections_counter[conn_index] = 0
 87
 88    def acquire_connection(self) -> Connection:
 89        """
 90        Acquire a connection from the connection pool for temporary synchronous
 91        calls. If the connection pool is oversubscribed, the method will return
 92        the connection with the least number of queued queries. It is required
 93        to release the connection by calling `release_connection` after the
 94        connection is no longer needed.
 95
 96        Returns
 97        -------
 98        Connection
 99            A connection object.
100        """
101        conn, _ = self.__get_connection_with_least_queries()
102        return conn
103
104    def release_connection(self, conn: Connection) -> None:
105        """
106        Release a connection acquired by `acquire_connection` back to the
107        connection pool. Calling this method is required when the connection is
108        no longer needed.
109
110        Parameters
111        ----------
112        conn : Connection
113            Connection object to release.
114
115
116        """
117        for i, existing_conn in enumerate(self.connections):
118            if existing_conn == conn:
119                self.__decrement_connection_counter(i)
120                break
121
122    def set_query_timeout(self, timeout_in_ms: int) -> None:
123        """
124        Set the query timeout value in ms for executing queries.
125
126        Parameters
127        ----------
128        timeout_in_ms : int
129            query timeout value in ms for executing queries.
130
131        """
132        for conn in self.connections:
133            conn.set_query_timeout(timeout_in_ms)
134
135    async def execute(
136        self, query: str | PreparedStatement, parameters: dict[str, Any] | None = None
137    ) -> QueryResult | list[QueryResult]:
138        """
139        Execute a query asynchronously.
140
141        Parameters
142        ----------
143        query : str | PreparedStatement
144            A prepared statement or a query string.
145            If a query string is given, a prepared statement will be created
146            automatically.
147
148        parameters : dict[str, Any]
149            Parameters for the query.
150
151        Returns
152        -------
153        QueryResult
154            Query result.
155
156        """
157        loop = asyncio.get_running_loop()
158        # If the query is a prepared statement, use the connection associated with it
159        if isinstance(query, PreparedStatement):
160            conn = query._connection
161            for i, existing_conn in enumerate(self.connections):
162                if existing_conn == conn:
163                    conn_index = i
164                    with self.lock:
165                        self.connections_counter[conn_index] += 1
166                    break
167        else:
168            conn, conn_index = self.__get_connection_with_least_queries()
169
170        try:
171            return await loop.run_in_executor(self.executor, conn.execute, query, parameters)
172        except asyncio.CancelledError:
173            conn.interrupt()
174        finally:
175            self.__decrement_connection_counter(conn_index)
176
177    async def prepare(self, query: str) -> PreparedStatement:
178        """
179        Create a prepared statement for a query asynchronously.
180
181        Parameters
182        ----------
183        query : str
184            Query to prepare.
185
186        Returns
187        -------
188        PreparedStatement
189            Prepared statement.
190
191        """
192        loop = asyncio.get_running_loop()
193        conn, conn_index = self.__get_connection_with_least_queries()
194
195        try:
196            preparedStatement = await loop.run_in_executor(self.executor, conn.prepare, query)
197            return preparedStatement
198        finally:
199            self.__decrement_connection_counter(conn_index)
200
201    def close(self) -> None:
202        """
203        Close all connections and shutdown the thread pool.
204
205        Note: Call to this method is optional. The connections and thread pool
206        will be closed automatically when the instance is garbage collected.
207        """
208        for conn in self.connections:
209            conn.close()
210
211        self.executor.shutdown(wait=True)

 23class Connection:
 24    """Connection to a database."""
 25
 26    def __init__(self, database: Database, num_threads: int = 0):
 27        """
 28        Initialise kuzu database connection.
 29
 30        Parameters
 31        ----------
 32        database : Database
 33            Database to connect to.
 34
 35        num_threads : int
 36            Maximum number of threads to use for executing queries.
 37
 38        """
 39        self._connection: Any = None  # (type: _kuzu.Connection from pybind11)
 40        self.database = database
 41        self.num_threads = num_threads
 42        self.is_closed = False
 43        self.init_connection()
 44
 45    def __getstate__(self) -> dict[str, Any]:
 46        state = {
 47            "database": self.database,
 48            "num_threads": self.num_threads,
 49            "_connection": None,
 50        }
 51        return state
 52
 53    def init_connection(self) -> None:
 54        """Establish a connection to the database, if not already initalised."""
 55        if self.is_closed:
 56            error_msg = "Connection is closed."
 57            raise RuntimeError(error_msg)
 58        self.database.init_database()
 59        if self._connection is None:
 60            self._connection = _kuzu.Connection(self.database._database, self.num_threads)  # type: ignore[union-attr]
 61
 62    def set_max_threads_for_exec(self, num_threads: int) -> None:
 63        """
 64        Set the maximum number of threads for executing queries.
 65
 66        Parameters
 67        ----------
 68        num_threads : int
 69            Maximum number of threads to use for executing queries.
 70
 71        """
 72        self.init_connection()
 73        self._connection.set_max_threads_for_exec(num_threads)
 74
 75    def close(self) -> None:
 76        """
 77        Close the connection.
 78
 79        Note: Call to this method is optional. The connection will be closed
 80        automatically when the object goes out of scope.
 81        """
 82        if self._connection is not None:
 83            self._connection.close()
 84        self._connection = None
 85        self.is_closed = True
 86
 87    def __enter__(self) -> Self:
 88        return self
 89
 90    def __exit__(
 91        self,
 92        exc_type: type[BaseException] | None,
 93        exc_value: BaseException | None,
 94        exc_traceback: TracebackType | None,
 95    ) -> None:
 96        self.close()
 97
 98    def execute(
 99        self,
100        query: str | PreparedStatement,
101        parameters: dict[str, Any] | None = None,
102    ) -> QueryResult | list[QueryResult]:
103        """
104        Execute a query.
105
106        Parameters
107        ----------
108        query : str | PreparedStatement
109            A prepared statement or a query string.
110            If a query string is given, a prepared statement will be created
111            automatically.
112
113        parameters : dict[str, Any]
114            Parameters for the query.
115
116        Returns
117        -------
118        QueryResult
119            Query result.
120
121        """
122        if parameters is None:
123            parameters = {}
124
125        self.init_connection()
126        if not isinstance(parameters, dict):
127            msg = f"Parameters must be a dict; found {type(parameters)}."
128            raise RuntimeError(msg)  # noqa: TRY004
129
130        if len(parameters) == 0 and isinstance(query, str):
131            _query_result = self._connection.query(query)
132        else:
133            prepared_statement = self.prepare(query) if isinstance(query, str) else query
134            _query_result = self._connection.execute(prepared_statement._prepared_statement, parameters)
135        if not _query_result.isSuccess():
136            raise RuntimeError(_query_result.getErrorMessage())
137        current_query_result = QueryResult(self, _query_result)
138        if not _query_result.hasNextQueryResult():
139            return current_query_result
140        all_query_results = [current_query_result]
141        while _query_result.hasNextQueryResult():
142            _query_result = _query_result.getNextQueryResult()
143            if not _query_result.isSuccess():
144                raise RuntimeError(_query_result.getErrorMessage())
145            all_query_results.append(QueryResult(self, _query_result))
146        return all_query_results
147
148    def prepare(self, query: str) -> PreparedStatement:
149        """
150        Create a prepared statement for a query.
151
152        Parameters
153        ----------
154        query : str
155            Query to prepare.
156
157        Returns
158        -------
159        PreparedStatement
160            Prepared statement.
161
162        """
163        return PreparedStatement(self, query)
164
165    def _get_node_property_names(self, table_name: str) -> dict[str, Any]:
166        LIST_START_SYMBOL = "["
167        LIST_END_SYMBOL = "]"
168        self.init_connection()
169        query_result = self.execute(f"CALL table_info('{table_name}') RETURN *;")
170        results = {}
171        while query_result.has_next():
172            row = query_result.get_next()
173            prop_name = row[1]
174            prop_type = row[2]
175            is_primary_key = row[4] is True
176            dimension = prop_type.count(LIST_START_SYMBOL)
177            splitted = prop_type.split(LIST_START_SYMBOL)
178            shape = []
179            for s in splitted:
180                if LIST_END_SYMBOL not in s:
181                    continue
182                s = s.split(LIST_END_SYMBOL)[0]
183                if s != "":
184                    shape.append(int(s))
185            prop_type = splitted[0]
186            results[prop_name] = {
187                "type": prop_type,
188                "dimension": dimension,
189                "is_primary_key": is_primary_key,
190            }
191            if len(shape) > 0:
192                results[prop_name]["shape"] = tuple(shape)
193        return results
194
195    def _get_node_table_names(self) -> list[Any]:
196        results = []
197        self.init_connection()
198        query_result = self.execute("CALL show_tables() RETURN *;")
199        while query_result.has_next():
200            row = query_result.get_next()
201            if row[2] == "NODE":
202                results.append(row[1])
203        return results
204
205    def _get_rel_table_names(self) -> list[dict[str, Any]]:
206        results = []
207        self.init_connection()
208        tables_result = self.execute("CALL show_tables() RETURN *;")
209        while tables_result.has_next():
210            row = tables_result.get_next()
211            if row[2] == "REL":
212                name = row[1]
213                connections_result = self.execute(f"CALL show_connection({name!r}) RETURN *;")
214                src_dst_row = connections_result.get_next()
215                src_node = src_dst_row[0]
216                dst_node = src_dst_row[1]
217                results.append({"name": name, "src": src_node, "dst": dst_node})
218        return results
219
220    def set_query_timeout(self, timeout_in_ms: int) -> None:
221        """
222        Set the query timeout value in ms for executing queries.
223
224        Parameters
225        ----------
226        timeout_in_ms : int
227            query timeout value in ms for executing queries.
228
229        """
230        self.init_connection()
231        self._connection.set_query_timeout(timeout_in_ms)
232
233    def interrupt(self) -> None:
234        """
235        Interrupts execution of the current query.
236
237        If there is no currently executing query, this function does nothing.
238        """
239        self._connection.interrupt()
240
241    def create_function(
242        self,
243        name: str,
244        udf: Callable[[...], Any],
245        params_type: list[Type | str] | None = None,
246        return_type: Type | str = "",
247        *,
248        default_null_handling: bool = True,
249        catch_exceptions: bool = False,
250    ) -> None:
251        """
252        Sets a User Defined Function (UDF) to use in cypher queries.
253
254        Parameters
255        ----------
256        name: str
257            name of function
258
259        udf: Callable[[...], Any]
260            function to be executed
261
262        params_type: Optional[list[Type]]
263            list of Type enums to describe the input parameters
264
265        return_type: Optional[Type]
266            a Type enum to describe the returned value
267
268        default_null_handling: Optional[bool]
269            if true, when any parameter is null, the resulting value will be null
270
271        catch_exceptions: Optional[bool]
272            if true, when an exception is thrown from python, the function output will be null
273            Otherwise, the exception will be rethrown
274        """
275        if params_type is None:
276            params_type = []
277        parsed_params_type = [x if type(x) is str else x.value for x in params_type]
278        if type(return_type) is not str:
279            return_type = return_type.value
280
281        self._connection.create_function(
282            name=name,
283            udf=udf,
284            params_type=parsed_params_type,
285            return_value=return_type,
286            default_null=default_null_handling,
287            catch_exceptions=catch_exceptions,
288        )
289
290    def remove_function(self, name: str) -> None:
291        """
292        Removes a User Defined Function (UDF).
293
294        Parameters
295        ----------
296        name: str
297            name of function to be removed.
298        """
299        self._connection.remove_function(name)

 26class Database:
 27    """Kuzu database instance."""
 28
 29    def __init__(
 30        self,
 31        database_path: str | Path | None = None,
 32        *,
 33        buffer_pool_size: int = 0,
 34        max_num_threads: int = 0,
 35        compression: bool = True,
 36        lazy_init: bool = False,
 37        read_only: bool = False,
 38        max_db_size: int = (1 << 43),
 39        auto_checkpoint: bool = True,
 40        checkpoint_threshold: int = -1,
 41    ):
 42        """
 43        Parameters
 44        ----------
 45        database_path : str, Path
 46            The path to database files. If the path is not specified, or empty, or equal to `:memory:`, the database
 47            will be created in memory.
 48
 49        buffer_pool_size : int
 50            The maximum size of buffer pool in bytes. Defaults to ~80% of system memory.
 51
 52        max_num_threads : int
 53            The maximum number of threads to use for executing queries.
 54
 55        compression : bool
 56            Enable database compression.
 57
 58        lazy_init : bool
 59            If True, the database will not be initialized until the first query.
 60            This is useful when the database is not used in the main thread or
 61            when the main process is forked.
 62            Default to False.
 63
 64        read_only : bool
 65            If true, the database is opened read-only. No write transactions is
 66            allowed on the `Database` object. Multiple read-only `Database`
 67            objects can be created with the same database path. However, there
 68            cannot be multiple `Database` objects created with the same
 69            database path.
 70            Default to False.
 71
 72        max_db_size : int
 73            The maximum size of the database in bytes. Note that this is introduced
 74            temporarily for now to get around with the default 8TB mmap address
 75             space limit some environment. This will be removed once we implemente
 76             a better solution later. The value is default to 1 << 43 (8TB) under 64-bit
 77             environment and 1GB under 32-bit one.
 78
 79        auto_checkpoint: bool
 80            If true, the database will automatically checkpoint when the size of
 81            the WAL file exceeds the checkpoint threshold.
 82
 83        checkpoint_threshold: int
 84            The threshold of the WAL file size in bytes. When the size of the
 85            WAL file exceeds this threshold, the database will checkpoint if autoCheckpoint is true.
 86
 87        """
 88        if database_path is None:
 89            database_path = ":memory:"
 90        if isinstance(database_path, Path):
 91            database_path = str(database_path)
 92
 93        self.database_path = database_path
 94        self.buffer_pool_size = buffer_pool_size
 95        self.max_num_threads = max_num_threads
 96        self.compression = compression
 97        self.read_only = read_only
 98        self.max_db_size = max_db_size
 99        self.auto_checkpoint = auto_checkpoint
100        self.checkpoint_threshold = checkpoint_threshold
101        self.is_closed = False
102
103        self._database: Any = None  # (type: _kuzu.Database from pybind11)
104        if not lazy_init:
105            self.init_database()
106
107    def __enter__(self) -> Self:
108        return self
109
110    def __exit__(
111        self,
112        exc_type: type[BaseException] | None,
113        exc_value: BaseException | None,
114        exc_traceback: TracebackType | None,
115    ) -> None:
116        self.close()
117
118    @staticmethod
119    def get_version() -> str:
120        """
121        Get the version of the database.
122
123        Returns
124        -------
125        str
126            The version of the database.
127        """
128        return _kuzu.Database.get_version()  # type: ignore[union-attr]
129
130    @staticmethod
131    def get_storage_version() -> int:
132        """
133        Get the storage version of the database.
134
135        Returns
136        -------
137        int
138            The storage version of the database.
139        """
140        return _kuzu.Database.get_storage_version()  # type: ignore[union-attr]
141
142    def __getstate__(self) -> dict[str, Any]:
143        state = {
144            "database_path": self.database_path,
145            "buffer_pool_size": self.buffer_pool_size,
146            "compression": self.compression,
147            "read_only": self.read_only,
148            "_database": None,
149        }
150        return state
151
152    def init_database(self) -> None:
153        """Initialize the database."""
154        self.check_for_database_close()
155        if self._database is None:
156            self._database = _kuzu.Database(  # type: ignore[union-attr]
157                self.database_path,
158                self.buffer_pool_size,
159                self.max_num_threads,
160                self.compression,
161                self.read_only,
162                self.max_db_size,
163                self.auto_checkpoint,
164                self.checkpoint_threshold,
165            )
166
167    def get_torch_geometric_remote_backend(
168        self, num_threads: int | None = None
169    ) -> tuple[KuzuFeatureStore, KuzuGraphStore]:
170        """
171        Use the database as the remote backend for torch_geometric.
172
173        For the interface of the remote backend, please refer to
174        https://pytorch-geometric.readthedocs.io/en/latest/advanced/remote.html.
175        The current implementation is read-only and does not support edge
176        features. The IDs of the nodes are based on the internal IDs (i.e., node
177        offsets). For the remote node IDs to be consistent with the positions in
178        the output tensors, please ensure that no deletion has been performed
179        on the node tables.
180
181        The remote backend can also be plugged into the data loader of
182        torch_geometric, which is useful for mini-batch training. For example:
183
184        ```python
185            loader_kuzu = NeighborLoader(
186                data=(feature_store, graph_store),
187                num_neighbors={('paper', 'cites', 'paper'): [12, 12, 12]},
188                batch_size=LOADER_BATCH_SIZE,
189                input_nodes=('paper', input_nodes),
190                num_workers=4,
191                filter_per_worker=False,
192            )
193        ```
194
195        Please note that the database instance is not fork-safe, so if more than
196        one worker is used, `filter_per_worker` must be set to False.
197
198        Parameters
199        ----------
200        num_threads : int
201            Number of threads to use for data loading. Default to None, which
202            means using the number of CPU cores.
203
204        Returns
205        -------
206        feature_store : KuzuFeatureStore
207            Feature store compatible with torch_geometric.
208        graph_store : KuzuGraphStore
209            Graph store compatible with torch_geometric.
210        """
211        self.check_for_database_close()
212        from .torch_geometric_feature_store import KuzuFeatureStore
213        from .torch_geometric_graph_store import KuzuGraphStore
214
215        return (
216            KuzuFeatureStore(self, num_threads),
217            KuzuGraphStore(self, num_threads),
218        )
219
220    def _scan_node_table(
221        self,
222        table_name: str,
223        prop_name: str,
224        prop_type: str,
225        dim: int,
226        indices: IndexType,
227        num_threads: int,
228    ) -> NDArray[Any]:
229        self.check_for_database_close()
230        import numpy as np
231
232        """
233        Scan a node table from storage directly, bypassing query engine.
234        Used internally by torch_geometric remote backend only.
235        """
236        self.init_database()
237        indices_cast = np.array(indices, dtype=np.uint64)
238        result = None
239
240        if prop_type == Type.INT64.value:
241            result = np.empty(len(indices) * dim, dtype=np.int64)
242            self._database.scan_node_table_as_int64(table_name, prop_name, indices_cast, result, num_threads)
243        elif prop_type == Type.INT32.value:
244            result = np.empty(len(indices) * dim, dtype=np.int32)
245            self._database.scan_node_table_as_int32(table_name, prop_name, indices_cast, result, num_threads)
246        elif prop_type == Type.INT16.value:
247            result = np.empty(len(indices) * dim, dtype=np.int16)
248            self._database.scan_node_table_as_int16(table_name, prop_name, indices_cast, result, num_threads)
249        elif prop_type == Type.DOUBLE.value:
250            result = np.empty(len(indices) * dim, dtype=np.float64)
251            self._database.scan_node_table_as_double(table_name, prop_name, indices_cast, result, num_threads)
252        elif prop_type == Type.FLOAT.value:
253            result = np.empty(len(indices) * dim, dtype=np.float32)
254            self._database.scan_node_table_as_float(table_name, prop_name, indices_cast, result, num_threads)
255
256        if result is not None:
257            return result
258
259        msg = f"Unsupported property type: {prop_type}"
260        raise ValueError(msg)
261
262    def close(self) -> None:
263        """
264        Close the database. Once the database is closed, the lock on the database
265        files is released and the database can be opened in another process.
266
267        Note: Call to this method is not required. The Python garbage collector
268        will automatically close the database when no references to the database
269        object exist. It is recommended not to call this method explicitly. If you
270        decide to manually close the database, make sure that all the QueryResult
271        and Connection objects are closed before calling this method.
272        """
273        if self.is_closed:
274            return
275        self.is_closed = True
276        if self._database is not None:
277            self._database.close()
278            self._database: Any = None  # (type: _kuzu.Database from pybind11)
279
280    def check_for_database_close(self) -> None:
281        """
282        Check if the database is closed and raise an exception if it is.
283
284        Raises
285        ------
286        Exception
287            If the database is closed.
288
289        """
290        if not self.is_closed:
291            return
292        msg = "Database is closed"
293        raise RuntimeError(msg)