From ba9e90cd60f424463810b4e3eebbf63d669e1c9e Mon Sep 17 00:00:00 2001
From: ipu <uwu@ipu.moe>
Date: Thu, 7 Aug 2025 20:38:50 +0300
Subject: [PATCH] add drugs db index

---
 ..._08_07_2030-058739dc7aa6_add_drug_index.py | 36 +++++++++++++++++++
 pyproject.toml                                |  1 +
 src/cache/drug_cache.py                       | 24 +++++++------
 src/config.py                                 |  2 ++
 src/database.py                               | 11 +++++-
 src/services/session_service.py               |  2 +-
 uv.lock                                       | 14 ++++++++
 7 files changed, 78 insertions(+), 12 deletions(-)
 create mode 100644 alembic/versions/2025_08_07_2030-058739dc7aa6_add_drug_index.py

diff --git a/alembic/versions/2025_08_07_2030-058739dc7aa6_add_drug_index.py b/alembic/versions/2025_08_07_2030-058739dc7aa6_add_drug_index.py
new file mode 100644
index 0000000..bbd0e2b
--- /dev/null
+++ b/alembic/versions/2025_08_07_2030-058739dc7aa6_add_drug_index.py
@@ -0,0 +1,36 @@
+"""add drug index
+
+Revision ID: 058739dc7aa6
+Revises: c509dcf806d4
+Create Date: 2025-08-07 20:30:35.962671
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = '058739dc7aa6'
+down_revision: Union[str, Sequence[str], None] = 'c509dcf806d4'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('drugs', schema=None) as batch_op:
+        batch_op.create_index('ix_drugs_name_dosage_dosage_unit', ['name', 'dosage', 'dosage_unit'], unique=True)
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('drugs', schema=None) as batch_op:
+        batch_op.drop_index('ix_drugs_name_dosage_dosage_unit')
+
+    # ### end Alembic commands ###
diff --git a/pyproject.toml b/pyproject.toml
index 1ae8625..795a5a3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,6 +22,7 @@ dependencies = [
     "asyncpg>=0.30.0",
     "alembic>=1.16.4",
     "psycopg2-binary>=2.9.10",
+    "tqdm>=4.67.1",
 ]
 
 [project.optional-dependencies]
diff --git a/src/cache/drug_cache.py b/src/cache/drug_cache.py
index 711a985..ef7f5fd 100644
--- a/src/cache/drug_cache.py
+++ b/src/cache/drug_cache.py
@@ -6,6 +6,7 @@ from src.database import Drug, Session
 from src.drug_price_parser import DrugPriceParser, DrugPriceResponse
 from src.config import settings
 from src.services.session_service import session_service
+from sqlalchemy.exc import IntegrityError
 
 
 class DrugFull(BaseModel):
@@ -36,25 +37,28 @@ async def convert_drug_result(drug: DrugPriceResponse) -> list[DrugFull]:
     return [DrugFull.model_validate(r) for r in response_dict["result"]]
 
 
-async def get_drug(drug_name: str) -> list[Drug]:
+async def get_drug(drug_name: str) -> list[DrugFull]:
     parser = DrugPriceParser()
     result = parser.get_drug_prices(drug_name)
     drugs = await convert_drug_result(result)
-    
     return drugs
 
 
 async def store_drug(drugs: list[DrugFull]):
     with Session() as session:
         for drug in drugs:
-            session.add(Drug(
-                name=drug.name,
-                dosage=drug.dosage,
-                dosage_unit=drug.dosage_unit,
-                unit_price=drug.unit_price,
-                description=drug.description
-            ))
-        session.commit()
+            try:
+                session.add(Drug(
+                    name=drug.name,
+                    dosage=drug.dosage,
+                    dosage_unit=drug.dosage_unit,
+                    unit_price=drug.unit_price,
+                    description=drug.description
+                ))
+                session.commit()
+            except IntegrityError as e:
+                session.rollback()
+                pass
 
 
 async def fetch_drug_with_dosage(drug_name: str, dosage: float) -> DrugFull | None:
diff --git a/src/config.py b/src/config.py
index 95cbe59..11d3cec 100644
--- a/src/config.py
+++ b/src/config.py
@@ -29,6 +29,8 @@ class Settings(BaseSettings):
     POSTGRES_HOST: str = "db"
     POSTGRES_PORT: int = 5432
 
+    DEBUG: bool = False
+
 
     @property
     def DATABASE_URL(self):    # noqa
diff --git a/src/database.py b/src/database.py
index bcf8c7e..7a5fcc4 100644
--- a/src/database.py
+++ b/src/database.py
@@ -1,4 +1,4 @@
-from sqlalchemy import BigInteger, Column, MetaData, String, Float, Text
+from sqlalchemy import BigInteger, Column, MetaData, String, Float, Text, Index
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy import create_engine
@@ -17,5 +17,14 @@ class Drug(Base):
     unit_price = Column(Float, nullable=False)
     description = Column(Text, nullable=True)
 
+Index(
+    "ix_drugs_name_dosage_dosage_unit",
+    Drug.name,
+    Drug.dosage,
+    Drug.dosage_unit,
+    unique=True
+)
+
+
 engine = create_engine(settings.DATABASE_URL)
 Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
diff --git a/src/services/session_service.py b/src/services/session_service.py
index bcafa7a..5515c51 100644
--- a/src/services/session_service.py
+++ b/src/services/session_service.py
@@ -61,7 +61,7 @@ class SessionService:
                 agent_id = self.agent_id or settings.TALESTORM_AGENT_ID
             
             response = await client.post("/sessions/", params={"agent_id": agent_id})
-            print(response.request.url)
+            # print(response.request.url)
             session_data = response.json()
             return str(session_data["id"])
 
diff --git a/uv.lock b/uv.lock
index d6f817d..12e1569 100644
--- a/uv.lock
+++ b/uv.lock
@@ -426,6 +426,7 @@ dependencies = [
     { name = "redis" },
     { name = "requests" },
     { name = "sqlalchemy", extra = ["asyncio"] },
+    { name = "tqdm" },
 ]
 
 [package.optional-dependencies]
@@ -461,6 +462,7 @@ requires-dist = [
     { name = "redis", specifier = ">=6.2.0" },
     { name = "requests", specifier = ">=2.31.0" },
     { name = "sqlalchemy", extras = ["asyncio"], specifier = ">=2.0.42" },
+    { name = "tqdm", specifier = ">=4.67.1" },
 ]
 provides-extras = ["dev"]
 
@@ -1178,6 +1180,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/1f/b876b1f83aef204198a42dc101613fefccb32258e5428b5f9259677864b4/starlette-0.47.2-py3-none-any.whl", hash = "sha256:c5847e96134e5c5371ee9fac6fdf1a67336d5815e09eb2a01fdb57a351ef915b", size = 72984, upload-time = "2025-07-20T17:31:56.738Z" },
 ]
 
+[[package]]
+name = "tqdm"
+version = "4.67.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737, upload-time = "2024-11-24T20:12:22.481Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540, upload-time = "2024-11-24T20:12:19.698Z" },
+]
+
 [[package]]
 name = "typer"
 version = "0.16.0"