Initial commit with all features

open_notebook/repository.py

@@ -0,0 +1,109 @@
+import asyncio
+import os
+from contextlib import asynccontextmanager
+
+from loguru import logger
+from surrealdb import Surreal
+
+from open_notebook.exceptions import InvalidDatabaseSchema
+
+EXPECTED_VERSION = "0.0.1"
+
+
+@asynccontextmanager
+async def db_connection():
+    db = Surreal(os.environ["SURREAL_ADDRESS"])
+    try:
+        await db.connect()
+        await db.signin(
+            {"user": os.environ["SURREAL_USER"], "pass": os.environ["SURREAL_PASS"]}
+        )
+        await db.use(os.environ["SURREAL_NAMESPACE"], os.environ["SURREAL_DATABASE"])
+        yield db
+    finally:
+        await db.close()
+
+
+def repo_query(query_str, vars=None):
+    async def _query():
+        async with db_connection() as db:
+            result = await db.query(query_str, vars)
+        return result
+
+    result = asyncio.run(_query())
+    return result[0]["result"]
+
+
+def check_version():
+    async def _check_version():
+        async with db_connection() as db:
+            result = await db.query("select * from open_notebook:database_info;")
+        return result
+
+    try:
+        result = asyncio.run(_check_version())
+        if len(result) == 0 or len(result[0]["result"]) == 0:
+            raise InvalidDatabaseSchema("Database schema not found")
+        version = result[0]["result"][0]["version"]
+        logger.info(f"Connected to SurrealDB, using schema version {version}")
+        if version != EXPECTED_VERSION:
+            raise InvalidDatabaseSchema(
+                f"Version mismatch. Expected {EXPECTED_VERSION}, got {version}"
+            )
+    except Exception as e:
+        logger.error(e)
+        raise e
+
+
+def repo_create(table, data):
+    async def _create():
+        async with db_connection() as db:
+            result = await db.create(table, data)
+        return result
+
+    result = asyncio.run(_create())
+    return result
+
+
+def repo_update(id, data):
+    async def _update():
+        async with db_connection() as db:
+            result = await db.update(id, data)
+        return result
+
+    result = asyncio.run(_update())
+    return result
+
+
+def repo_delete(id):
+    async def _delete():
+        async with db_connection() as db:
+            result = await db.delete(id)
+        return result
+
+    result = asyncio.run(_delete())
+    return result
+
+
+def repo_relate(source, relationship, target):
+    async def _relate():
+        async with db_connection() as db:
+            query = f"RELATE {source}->{relationship}->{target};"
+            result = await db.query(query)
+        return result
+
+    result = asyncio.run(_relate())
+    return result
+
+
+def execute_migration():
+    async def _query():
+        content = None
+        with open("db_setup.surrealql", "r") as file:
+            content = file.read()
+        async with db_connection() as db:
+            result = await db.query(content)
+        return result
+
+    result = asyncio.run(_query())
+    return result[0]["result"]