cp

src/db/repositories/api_key_repository.py

@@ -15,11 +15,12 @@ class ApiKeyRepository:
     
     @property
     def collection(self):
-        database = db.get_database()
-        if database is None:
+        client = db.get_database()
+        if client is None:
             logger.error("Database connection is None, cannot access collection")
             raise RuntimeError("Database connection is not available")
-        return database.collection(self.collection_name)
+        # Use the Firestore client to get the collection
+        return client.collection(self.collection_name)
     
     async def create(self, api_key: ApiKeyModel) -> ApiKeyModel:
         """
@@ -62,7 +63,7 @@ class ApiKeyRepository:
             API key if found, None otherwise
         """
         try:
-            doc_ref = self.collection.document(key_id)
+            doc_ref = self.collection.document(str(key_id))
             doc = doc_ref.get()
             if doc.exists:
                 data = doc.to_dict()
@@ -86,6 +87,7 @@ class ApiKeyRepository:
         try:
             query = self.collection.where("key_hash", "==", key_hash).limit(1)
             results = query.stream()
+            
             for doc in results:
                 data = doc.to_dict()
                 data['id'] = doc.id
@@ -107,8 +109,9 @@ class ApiKeyRepository:
         """
         try:
             keys = []
-            query = self.collection.where("user_id", "==", user_id)
+            query = self.collection.where("user_id", "==", str(user_id))
             results = query.stream()
+            
             for doc in results:
                 data = doc.to_dict()
                 data['id'] = doc.id
@@ -130,8 +133,9 @@ class ApiKeyRepository:
         """
         try:
             keys = []
-            query = self.collection.where("team_id", "==", team_id)
+            query = self.collection.where("team_id", "==", str(team_id))
             results = query.stream()
+            
             for doc in results:
                 data = doc.to_dict()
                 data['id'] = doc.id
@@ -149,7 +153,7 @@ class ApiKeyRepository:
             key_id: API key ID
         """
         try:
-            doc_ref = self.collection.document(key_id)
+            doc_ref = self.collection.document(str(key_id))
             doc_ref.update({"last_used": datetime.utcnow()})
         except Exception as e:
             logger.error(f"Error updating API key last used: {e}", exc_info=True)
@@ -166,7 +170,7 @@ class ApiKeyRepository:
             True if deactivated, False otherwise
         """
         try:
-            doc_ref = self.collection.document(key_id)
+            doc_ref = self.collection.document(str(key_id))
             doc_ref.update({"is_active": False})
             return True
         except Exception as e:
@@ -184,7 +188,7 @@ class ApiKeyRepository:
             True if deleted, False otherwise
         """
         try:
-            doc_ref = self.collection.document(key_id)
+            doc_ref = self.collection.document(str(key_id))
             doc_ref.delete()
             return True
         except Exception as e:

src/db/repositories/team_repository.py

@@ -15,7 +15,12 @@ class TeamRepository:
     
     @property
     def collection(self):
-        return db.get_database()[self.collection_name]
+        client = db.get_database()
+        if client is None:
+            logger.error("Database connection is None, cannot access collection")
+            raise RuntimeError("Database connection is not available")
+        # Use the Firestore client to get the collection
+        return client.collection(self.collection_name)
     
     async def create(self, team: TeamModel) -> TeamModel:
         """
@@ -28,10 +33,18 @@ class TeamRepository:
             Created team with ID
         """
         try:
-            result = await self.collection.insert_one(team.dict(by_alias=True))
-            created_team = await self.get_by_id(result.inserted_id)
-            logger.info(f"Team created: {result.inserted_id}")
-            return created_team
+            # Convert team to dict for Firestore
+            team_dict = team.dict(by_alias=True)
+            
+            # Add document to Firestore and get reference
+            doc_ref = self.collection.document()
+            doc_ref.set(team_dict)
+            
+            # Update ID and retrieve the created team
+            team.id = doc_ref.id
+            logger.info(f"Team created: {doc_ref.id}")
+            
+            return team
         except Exception as e:
             logger.error(f"Error creating team: {e}")
             raise
@@ -47,9 +60,13 @@ class TeamRepository:
             Team if found, None otherwise
         """
         try:
-            team = await self.collection.find_one({"_id": team_id})
-            if team:
-                return TeamModel(**team)
+            doc_ref = self.collection.document(str(team_id))
+            team_doc = doc_ref.get()
+            
+            if team_doc.exists:
+                team_data = team_doc.to_dict()
+                team_data['id'] = team_doc.id
+                return TeamModel(**team_data)
             return None
         except Exception as e:
             logger.error(f"Error getting team by ID: {e}")
@@ -64,9 +81,13 @@ class TeamRepository:
         """
         try:
             teams = []
-            cursor = self.collection.find()
-            async for document in cursor:
-                teams.append(TeamModel(**document))
+            team_docs = self.collection.stream()
+            
+            for doc in team_docs:
+                team_data = doc.to_dict()
+                team_data['id'] = doc.id
+                teams.append(TeamModel(**team_data))
+                
             return teams
         except Exception as e:
             logger.error(f"Error getting all teams: {e}")
@@ -87,20 +108,23 @@ class TeamRepository:
             # Add updated_at timestamp
             team_data["updated_at"] = datetime.utcnow()
             
-            # Don't allow updating _id
+            # Don't allow updating _id or id
             if "_id" in team_data:
                 del team_data["_id"]
+            if "id" in team_data:
+                del team_data["id"]
                 
-            result = await self.collection.update_one(
-                {"_id": team_id},
-                {"$set": team_data}
-            )
+            # Get document reference and update
+            doc_ref = self.collection.document(str(team_id))
+            doc_ref.update(team_data)
             
-            if result.modified_count == 0:
-                logger.warning(f"No team updated for ID: {team_id}")
+            # Get updated team
+            updated_team = await self.get_by_id(team_id)
+            if not updated_team:
+                logger.warning(f"No team found after update for ID: {team_id}")
                 return None
                 
-            return await self.get_by_id(team_id)
+            return updated_team
         except Exception as e:
             logger.error(f"Error updating team: {e}")
             raise
@@ -116,8 +140,9 @@ class TeamRepository:
             True if team was deleted, False otherwise
         """
         try:
-            result = await self.collection.delete_one({"_id": team_id})
-            return result.deleted_count > 0
+            doc_ref = self.collection.document(str(team_id))
+            doc_ref.delete()
+            return True
         except Exception as e:
             logger.error(f"Error deleting team: {e}")
             raise

src/db/repositories/user_repository.py

@@ -15,7 +15,12 @@ class UserRepository:
     
     @property
     def collection(self):
-        return db.get_database()[self.collection_name]
+        client = db.get_database()
+        if client is None:
+            logger.error("Database connection is None, cannot access collection")
+            raise RuntimeError("Database connection is not available")
+        # Use the Firestore client to get the collection
+        return client.collection(self.collection_name)
     
     async def create(self, user: UserModel) -> UserModel:
         """
@@ -28,10 +33,18 @@ class UserRepository:
             Created user with ID
         """
         try:
-            result = await self.collection.insert_one(user.dict(by_alias=True))
-            created_user = await self.get_by_id(result.inserted_id)
-            logger.info(f"User created: {result.inserted_id}")
-            return created_user
+            # Convert user to dict for Firestore
+            user_dict = user.dict(by_alias=True)
+            
+            # Add document to Firestore and get reference
+            doc_ref = self.collection.document()
+            doc_ref.set(user_dict)
+            
+            # Update ID and retrieve the created user
+            user.id = doc_ref.id
+            logger.info(f"User created: {doc_ref.id}")
+            
+            return user
         except Exception as e:
             logger.error(f"Error creating user: {e}")
             raise
@@ -47,9 +60,13 @@ class UserRepository:
             User if found, None otherwise
         """
         try:
-            user = await self.collection.find_one({"_id": user_id})
-            if user:
-                return UserModel(**user)
+            doc_ref = self.collection.document(str(user_id))
+            user_doc = doc_ref.get()
+            
+            if user_doc.exists:
+                user_data = user_doc.to_dict()
+                user_data['id'] = user_doc.id
+                return UserModel(**user_data)
             return None
         except Exception as e:
             logger.error(f"Error getting user by ID: {e}")
@@ -66,9 +83,13 @@ class UserRepository:
             User if found, None otherwise
         """
         try:
-            user = await self.collection.find_one({"email": email})
-            if user:
-                return UserModel(**user)
+            query = self.collection.where("email", "==", email).limit(1)
+            results = query.stream()
+            
+            for doc in results:
+                user_data = doc.to_dict()
+                user_data['id'] = doc.id
+                return UserModel(**user_data)
             return None
         except Exception as e:
             logger.error(f"Error getting user by email: {e}")
@@ -86,9 +107,13 @@ class UserRepository:
         """
         try:
             users = []
-            cursor = self.collection.find({"team_id": team_id})
-            async for document in cursor:
-                users.append(UserModel(**document))
+            query = self.collection.where("team_id", "==", str(team_id))
+            results = query.stream()
+            
+            for doc in results:
+                user_data = doc.to_dict()
+                user_data['id'] = doc.id
+                users.append(UserModel(**user_data))
             return users
         except Exception as e:
             logger.error(f"Error getting users by team: {e}")
@@ -109,20 +134,23 @@ class UserRepository:
             # Add updated_at timestamp
             user_data["updated_at"] = datetime.utcnow()
             
-            # Don't allow updating _id
+            # Don't allow updating _id or id
             if "_id" in user_data:
                 del user_data["_id"]
+            if "id" in user_data:
+                del user_data["id"]
                 
-            result = await self.collection.update_one(
-                {"_id": user_id},
-                {"$set": user_data}
-            )
+            # Get document reference and update
+            doc_ref = self.collection.document(str(user_id))
+            doc_ref.update(user_data)
             
-            if result.modified_count == 0:
-                logger.warning(f"No user updated for ID: {user_id}")
+            # Get updated user
+            updated_user = await self.get_by_id(user_id)
+            if not updated_user:
+                logger.warning(f"No user found after update for ID: {user_id}")
                 return None
                 
-            return await self.get_by_id(user_id)
+            return updated_user
         except Exception as e:
             logger.error(f"Error updating user: {e}")
             raise
@@ -138,8 +166,9 @@ class UserRepository:
             True if user was deleted, False otherwise
         """
         try:
-            result = await self.collection.delete_one({"_id": user_id})
-            return result.deleted_count > 0
+            doc_ref = self.collection.document(str(user_id))
+            doc_ref.delete()
+            return True
         except Exception as e:
             logger.error(f"Error deleting user: {e}")
             raise
@@ -152,10 +181,8 @@ class UserRepository:
             user_id: User ID
         """
         try:
-            await self.collection.update_one(
-                {"_id": user_id},
-                {"$set": {"last_login": datetime.utcnow()}}
-            )
+            doc_ref = self.collection.document(str(user_id))
+            doc_ref.update({"last_login": datetime.utcnow()})
         except Exception as e:
             logger.error(f"Error updating user last login: {e}")
             raise