11import tempfile
22import shutil
3+ import warnings
34
4- from chdb import query
5+ import chdb
6+ from ..state import sqlitelike as chdb_stateful
7+
8+
9+ g_session = None
10+ g_session_path = None
511
612
713class Session :
814 """
9- Session will keep the state of query. All DDL and DML state will be kept in a dir.
10- Dir path could be passed in as an argument. If not, a temporary dir will be created.
15+ Session will keep the state of query.
16+ If path is None, it will create a temporary directory and use it as the database path
17+ and the temporary directory will be removed when the session is closed.
18+ You can also pass in a path to create a database at that path where will keep your data.
19+
20+ You can also use a connection string to pass in the path and other parameters.
21+ Examples:
22+ - ":memory:" (for in-memory database)
23+ - "test.db" (for relative path)
24+ - "file:test.db" (same as above)
25+ - "/path/to/test.db" (for absolute path)
26+ - "file:/path/to/test.db" (same as above)
27+ - "file:test.db?param1=value1¶m2=value2" (for relative path with query params)
28+ - "file::memory:?verbose&log-level=test" (for in-memory database with query params)
29+ - "///path/to/test.db?param1=value1¶m2=value2" (for absolute path)
1130
12- If path is not specified, the temporary dir will be deleted when the Session object is deleted.
13- Otherwise path will be kept.
31+ Connection string args handling:
32+ Connection string can contain query params like "file:test.db?param1=value1¶m2=value2"
33+ "param1=value1" will be passed to ClickHouse engine as start up args.
1434
15- Note: The default database is "_local" and the default engine is "Memory" which means all data
16- will be stored in memory. If you want to store data in disk, you should create another database.
35+ For more details, see `clickhouse local --help --verbose`
36+ Some special args handling:
37+ - "mode=ro" would be "--readonly=1" for clickhouse (read-only mode)
38+
39+ Important:
40+ - There can be only one session at a time. If you want to create a new session, you need to close the existing one.
41+ - Creating a new session will close the existing one.
1742 """
1843
1944 def __init__ (self , path = None ):
45+ global g_session , g_session_path
46+ if g_session is not None :
47+ warnings .warn (
48+ "There is already an active session. Creating a new session will close the existing one. "
49+ "It is recommended to close the existing session before creating a new one. "
50+ f"Closing the existing session { g_session_path }
51+ )
52+ g_session .close ()
53+ g_session_path = None
2054 if path is None :
2155 self ._cleanup = True
2256 self ._path = tempfile .mkdtemp ()
2357 else :
2458 self ._cleanup = False
2559 self ._path = path
60+ if chdb .g_udf_path != "" :
61+ self ._udf_path = chdb .g_udf_path
62+ # add udf_path to conn_str here.
63+ # - the `user_scripts_path` will be the value of `udf_path`
64+ # - the `user_defined_executable_functions_config` will be `user_scripts_path/*.xml`
65+ # Both of them will be added to the conn_str in the Connection class
66+ if "?" in self ._path :
67+ self ._conn_str = f"{ self ._path } { self ._udf_path }
68+ else :
69+ self ._conn_str = f"{ self ._path } { self ._udf_path }
70+ else :
71+ self ._udf_path = ""
72+ self ._conn_str = f"{ self ._path }
73+ self ._conn = chdb_stateful .Connection (self ._conn_str )
74+ g_session = self
75+ g_session_path = self ._path
2676
2777 def __del__ (self ):
28- if self ._cleanup :
29- self .cleanup ()
78+ self .close ()
3079
3180 def __enter__ (self ):
3281 return self
3382
3483 def __exit__ (self , exc_type , exc_value , traceback ):
35- self .cleanup ()
84+ self .close ()
85+
86+ def close (self ):
87+ if self ._cleanup :
88+ self .cleanup ()
89+ if self ._conn is not None :
90+ self ._conn .close ()
91+ self ._conn = None
3692
3793 def cleanup (self ):
3894 try :
95+ if self ._conn is not None :
96+ self ._conn .close ()
97+ self ._conn = None
3998 shutil .rmtree (self ._path )
4099 except : # noqa
41100 pass
@@ -44,7 +103,14 @@ def query(self, sql, fmt="CSV", udf_path=""):
44103 """
45104 Execute a query.
46105 """
47- return query (sql , fmt , path = self ._path , udf_path = udf_path )
106+ if fmt == "Debug" :
107+ warnings .warn (
108+ """Debug format is not supported in Session.query
109+ Please try use parameters in connection string instead:
110+ Eg: conn = connect(f"db_path?verbose&log-level=test")"""
111+ )
112+ fmt = "CSV"
113+ return self ._conn .query (sql , fmt )
48114
49115 # alias sql = query
50116 sql = query
0 commit comments