DH-4280 Creates endpoint to add descriptions (#77)

dataherald/api/__init__.py

@@ -13,6 +13,7 @@
     NLQueryResponse,
     QuestionRequest,
     ScannerRequest,
+    TableDescriptionRequest,
     UpdateQueryRequest,
 )
 
@@ -41,6 +42,15 @@ def connect_database(
     ) -> bool:
         pass
 
+    @abstractmethod
+    def add_description(
+        self,
+        db_name: str,
+        table_name: str,
+        table_description_request: TableDescriptionRequest,
+    ) -> bool:
+        pass
+
     @abstractmethod
     def add_golden_records(self, golden_records: List) -> bool:
         pass

dataherald/api/fastapi.py

@@ -32,6 +32,7 @@
     NLQueryResponse,
     QuestionRequest,
     ScannerRequest,
+    TableDescriptionRequest,
     UpdateQueryRequest,
 )
 
@@ -159,6 +160,26 @@ def connect_database(
         )
         return True
 
+    @override
+    def add_description(
+        self,
+        db_name: str,
+        table_name: str,
+        table_description_request: TableDescriptionRequest,
+    ) -> bool:
+        scanner_repository = DBScannerRepository(self.storage)
+        table = scanner_repository.get_table_info(db_name, table_name)
+        if table_description_request.description:
+            table.description = table_description_request.description
+        if table_description_request.columns:
+            for column_request in table_description_request.columns:
+                for column in table.columns:
+                    if column_request.name == column.name:
+                        column.description = column_request.description
+
+        scanner_repository.update(table)
+        return True
+
     @override
     def add_data_definition(
         self, data_definition_request: DataDefinitionRequest

dataherald/db_scanner/models/types.py

@@ -19,6 +19,7 @@ class ColumnDetail(BaseModel):
 
 
 class TableSchemaDetail(BaseModel):
+    id: Any
     db_alias: str
     table_name: str
     description: str | None

dataherald/db_scanner/repository/base.py

@@ -1,7 +1,11 @@
 from typing import List
 
+from bson.objectid import ObjectId
+
 from dataherald.db_scanner.models.types import TableSchemaDetail
 
+DB_COLLECTION = "table_schema_detail"
+
 
 class DBScannerRepository:
     def __init__(self, storage):
@@ -11,19 +15,32 @@ def get_table_info(
         self, db_alias: str, table_name: str
     ) -> TableSchemaDetail | None:
         row = self.storage.find_one(
-            "table_schema_detail", {"db_alias": db_alias, "table_name": table_name}
+            DB_COLLECTION, {"db_alias": db_alias, "table_name": table_name}
         )
+        row["id"] = row["_id"]
         if row:
             return TableSchemaDetail(**row)
         return None
 
     def get_all_tables_by_db(self, db_alias: str) -> List[TableSchemaDetail]:
-        rows = self.storage.find("table_schema_detail", {"db_alias": db_alias})
-        return [TableSchemaDetail(**row) for row in rows]
+        rows = self.storage.find(DB_COLLECTION, {"db_alias": db_alias})
+        tables = []
+        for row in rows:
+            row["id"] = row["_id"]
+            tables.append(TableSchemaDetail(**row))
+        return tables
 
     def save_table_info(self, table_info: TableSchemaDetail) -> None:
         self.storage.update_or_create(
-            "table_schema_detail",
+            DB_COLLECTION,
             {"db_alias": table_info.db_alias, "table_name": table_info.table_name},
             table_info.dict(),
         )
+
+    def update(self, table_info: TableSchemaDetail) -> TableSchemaDetail:
+        self.storage.update_or_create(
+            DB_COLLECTION,
+            {"_id": ObjectId(table_info.id)},
+            table_info.dict(exclude={"id"}),
+        )
+        return table_info

dataherald/repositories/base.py

@@ -22,7 +22,7 @@ def find_one(self, query: dict) -> NLQueryResponse | None:
         row["id"] = row["_id"]
         return NLQueryResponse(**row)
 
-    def update(self, nl_query_response: NLQueryResponse) -> int:
+    def update(self, nl_query_response: NLQueryResponse) -> NLQueryResponse:
         self.storage.update_or_create(
             DB_COLLECTION,
             {"_id": ObjectId(nl_query_response.id)},

dataherald/server/fastapi/__init__.py

@@ -19,6 +19,7 @@
     NLQueryResponse,
     QuestionRequest,
     ScannerRequest,
+    TableDescriptionRequest,
     UpdateQueryRequest,
 )
 
@@ -58,6 +59,12 @@ def __init__(self, settings: Settings):
             "/api/v1/database", self.connect_database, methods=["POST"]
         )
 
+        self.router.add_api_route(
+            "/api/v1/scanned-db/{db_name}/{table_name}",
+            self.add_description,
+            methods=["PATCH"],
+        )
+
         self.router.add_api_route(
             "/api/v1/golden-record", self.add_golden_records, methods=["POST"]
         )
@@ -105,6 +112,15 @@ def connect_database(
         """Connects a database to the Dataherald service"""
         return self._api.connect_database(database_connection_request)
 
+    def add_description(
+        self,
+        db_name: str,
+        table_name: str,
+        table_description_request: TableDescriptionRequest,
+    ) -> bool:
+        """Add descriptions for tables and columns"""
+        return self._api.add_description(db_name, table_name, table_description_request)
+
     def add_golden_records(self, golden_records: List) -> bool:
         """Takes in an English question and answers it based on content from the registered databases"""
         return self._api.add_golden_records(golden_records)

dataherald/types.py

@@ -87,3 +87,13 @@ class DatabaseConnectionRequest(BaseModel):
 class DataDefinitionRequest(BaseModel):
     uri: str
     type: DataDefinitionType
+
+
+class ColumnDescriptionRequest(BaseModel):
+    name: str
+    description: str
+
+
+class TableDescriptionRequest(BaseModel):
+    description: str | None
+    columns: list[ColumnDescriptionRequest] | None