11from __future__ import annotations
22
3+ import contextlib
34import logging
45from collections import defaultdict
56from contextlib import asynccontextmanager
@@ -48,10 +49,11 @@ class DataStore:
4849 db_wrapper : DBWrapper2
4950
5051 @classmethod
51- async def create (
52+ @contextlib .asynccontextmanager
53+ async def managed (
5254 cls , database : Union [str , Path ], uri : bool = False , sql_log_path : Optional [Path ] = None
53- ) -> DataStore :
54- db_wrapper = await DBWrapper2 .create (
55+ ) -> AsyncIterator [ DataStore ] :
56+ async with DBWrapper2 .managed (
5557 database = database ,
5658 uri = uri ,
5759 journal_mode = "WAL" ,
@@ -63,100 +65,97 @@ async def create(
6365 foreign_keys = True ,
6466 row_factory = aiosqlite .Row ,
6567 log_path = sql_log_path ,
66- )
67- self = cls (db_wrapper = db_wrapper )
68+ ) as db_wrapper :
69+ self = cls (db_wrapper = db_wrapper )
6870
69- async with db_wrapper .writer () as writer :
70- await writer .execute (
71- f"""
72- CREATE TABLE IF NOT EXISTS node(
73- hash BLOB PRIMARY KEY NOT NULL CHECK(length(hash) == 32),
74- node_type INTEGER NOT NULL CHECK(
75- (
76- node_type == { int (NodeType .INTERNAL )}
77- AND left IS NOT NULL
78- AND right IS NOT NULL
79- AND key IS NULL
80- AND value IS NULL
81- )
82- OR
83- (
84- node_type == { int (NodeType .TERMINAL )}
85- AND left IS NULL
86- AND right IS NULL
87- AND key IS NOT NULL
88- AND value IS NOT NULL
89- )
90- ),
91- left BLOB REFERENCES node,
92- right BLOB REFERENCES node,
93- key BLOB,
94- value BLOB
71+ async with db_wrapper .writer () as writer :
72+ await writer .execute (
73+ f"""
74+ CREATE TABLE IF NOT EXISTS node(
75+ hash BLOB PRIMARY KEY NOT NULL CHECK(length(hash) == 32),
76+ node_type INTEGER NOT NULL CHECK(
77+ (
78+ node_type == { int (NodeType .INTERNAL )}
79+ AND left IS NOT NULL
80+ AND right IS NOT NULL
81+ AND key IS NULL
82+ AND value IS NULL
83+ )
84+ OR
85+ (
86+ node_type == { int (NodeType .TERMINAL )}
87+ AND left IS NULL
88+ AND right IS NULL
89+ AND key IS NOT NULL
90+ AND value IS NOT NULL
91+ )
92+ ),
93+ left BLOB REFERENCES node,
94+ right BLOB REFERENCES node,
95+ key BLOB,
96+ value BLOB
97+ )
98+ """
9599 )
96- """
97- )
98- await writer .execute (
99- """
100- CREATE TRIGGER IF NOT EXISTS no_node_updates
101- BEFORE UPDATE ON node
102- BEGIN
103- SELECT RAISE(FAIL, 'updates not allowed to the node table');
104- END
105- """
106- )
107- await writer .execute (
108- f"""
109- CREATE TABLE IF NOT EXISTS root(
110- tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
111- generation INTEGER NOT NULL CHECK(generation >= 0),
112- node_hash BLOB,
113- status INTEGER NOT NULL CHECK(
114- { " OR " .join (f"status == { status } for status in Status )}
115- ),
116- PRIMARY KEY(tree_id, generation),
117- FOREIGN KEY(node_hash) REFERENCES node(hash)
100+ await writer .execute (
101+ """
102+ CREATE TRIGGER IF NOT EXISTS no_node_updates
103+ BEFORE UPDATE ON node
104+ BEGIN
105+ SELECT RAISE(FAIL, 'updates not allowed to the node table');
106+ END
107+ """
118108 )
119- """
120- )
121- # TODO: Add ancestor -> hash relationship, this might involve temporarily
122- # deferring the foreign key enforcement due to the insertion order
123- # and the node table also enforcing a similar relationship in the
124- # other direction.
125- # FOREIGN KEY(ancestor) REFERENCES ancestors(ancestor)
126- await writer .execute (
127- """
128- CREATE TABLE IF NOT EXISTS ancestors(
129- hash BLOB NOT NULL REFERENCES node,
130- ancestor BLOB CHECK(length(ancestor) == 32),
131- tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
132- generation INTEGER NOT NULL,
133- PRIMARY KEY(hash, tree_id, generation),
134- FOREIGN KEY(ancestor) REFERENCES node(hash)
109+ await writer .execute (
110+ f"""
111+ CREATE TABLE IF NOT EXISTS root(
112+ tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
113+ generation INTEGER NOT NULL CHECK(generation >= 0),
114+ node_hash BLOB,
115+ status INTEGER NOT NULL CHECK(
116+ { " OR " .join (f"status == { status } for status in Status )}
117+ ),
118+ PRIMARY KEY(tree_id, generation),
119+ FOREIGN KEY(node_hash) REFERENCES node(hash)
120+ )
121+ """
135122 )
136- """
137- )
138- await writer .execute (
139- """
140- CREATE TABLE IF NOT EXISTS subscriptions(
141- tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
142- url TEXT,
143- ignore_till INTEGER,
144- num_consecutive_failures INTEGER,
145- from_wallet tinyint CHECK(from_wallet == 0 OR from_wallet == 1),
146- PRIMARY KEY(tree_id, url)
123+ # TODO: Add ancestor -> hash relationship, this might involve temporarily
124+ # deferring the foreign key enforcement due to the insertion order
125+ # and the node table also enforcing a similar relationship in the
126+ # other direction.
127+ # FOREIGN KEY(ancestor) REFERENCES ancestors(ancestor)
128+ await writer .execute (
129+ """
130+ CREATE TABLE IF NOT EXISTS ancestors(
131+ hash BLOB NOT NULL REFERENCES node,
132+ ancestor BLOB CHECK(length(ancestor) == 32),
133+ tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
134+ generation INTEGER NOT NULL,
135+ PRIMARY KEY(hash, tree_id, generation),
136+ FOREIGN KEY(ancestor) REFERENCES node(hash)
137+ )
138+ """
139+ )
140+ await writer .execute (
141+ """
142+ CREATE TABLE IF NOT EXISTS subscriptions(
143+ tree_id BLOB NOT NULL CHECK(length(tree_id) == 32),
144+ url TEXT,
145+ ignore_till INTEGER,
146+ num_consecutive_failures INTEGER,
147+ from_wallet tinyint CHECK(from_wallet == 0 OR from_wallet == 1),
148+ PRIMARY KEY(tree_id, url)
149+ )
150+ """
151+ )
152+ await writer .execute (
153+ """
154+ CREATE INDEX IF NOT EXISTS node_hash ON root(node_hash)
155+ """
147156 )
148- """
149- )
150- await writer .execute (
151- """
152- CREATE INDEX IF NOT EXISTS node_hash ON root(node_hash)
153- """
154- )
155-
156- return self
157157
158- async def close (self ) -> None :
159- await self .db_wrapper .close ()
158+ yield self
160159
161160 @asynccontextmanager
162161 async def transaction (self ) -> AsyncIterator [None ]:
0 commit comments