From 9d1230ad24b244b026aa8a48e999ba17992758e1 Mon Sep 17 00:00:00 2001
From: gabrielmbmb <gmartinbdev@gmail.com>
Date: Tue, 25 Jul 2023 12:00:03 +0200
Subject: [PATCH] feat: remove sync engine, sessiond and `get_db` function
---
.../server/apis/v1/handlers/responses.py | 3 +--
src/argilla/server/database.py | 19 ++++---------------
2 files changed, 5 insertions(+), 17 deletions(-)
diff --git a/src/argilla/server/apis/v1/handlers/responses.py b/src/argilla/server/apis/v1/handlers/responses.py
index cce444a71c..588c67cff7 100644
--- a/src/argilla/server/apis/v1/handlers/responses.py
+++ b/src/argilla/server/apis/v1/handlers/responses.py
@@ -16,10 +16,9 @@
from fastapi import APIRouter, Depends, HTTPException, Security, status
from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import Session
from argilla.server.contexts import datasets
-from argilla.server.database import get_async_db, get_db
+from argilla.server.database import get_async_db
from argilla.server.models import User
from argilla.server.policies import ResponsePolicyV1, authorize
from argilla.server.schemas.v1.responses import Response, ResponseUpdate
diff --git a/src/argilla/server/database.py b/src/argilla/server/database.py
index 33dff80f95..588512474c 100644
--- a/src/argilla/server/database.py
+++ b/src/argilla/server/database.py
@@ -16,16 +16,16 @@
from sqlite3 import Connection as SQLite3Connection
from typing import TYPE_CHECKING, Generator
-from sqlalchemy import create_engine, event
+from sqlalchemy import event
from sqlalchemy.engine import Engine
-from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
import argilla
from argilla.server.settings import settings
if TYPE_CHECKING:
- from sqlalchemy.orm import Session
+ from sqlalchemy.ext.asyncio import AsyncSession
+
ALEMBIC_CONFIG_FILE = os.path.normpath(os.path.join(os.path.dirname(argilla.__file__), "alembic.ini"))
TAGGED_REVISIONS = OrderedDict(
@@ -46,21 +46,10 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
cursor.close()
-engine = create_engine(settings.database_url)
-SessionLocal = sessionmaker(autocommit=False, bind=engine)
-
async_engine = create_async_engine(settings.database_url_async)
AsyncSessionLocal = async_sessionmaker(autocommit=False, expire_on_commit=False, bind=async_engine)
-def get_db() -> Generator["Session", None, None]:
- try:
- db = SessionLocal()
- yield db
- finally:
- db.close()
-
-
async def get_async_db() -> Generator["AsyncSession", None, None]:
try:
db: "AsyncSession" = AsyncSessionLocal()