Repository: topoteretes/cognee
File: cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py (Python)

Code before:
        if "sqlite" in connection_string:
            self.engine = create_async_engine(
                connection_string,
                poolclass=NullPool,
                connect_args={"timeout": 30},
            )
        else:
            self.engine = create_async_engine(
                connection_string,
                pool_size=20,
                max_overflow=20,
                pool_recycle=280,
                pool_pre_ping=True,
                pool_timeout=280,
            )

        self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)

    async def push_to_s3(self) -> None:
        if os.getenv("STORAGE_BACKEND", "").lower() == "s3" and hasattr(self, "temp_db_file"):
            from cognee.infrastructure.files.storage.S3FileStorage import S3FileStorage

            s3_file_storage = S3FileStorage("")
            s3_file_storage.s3.put(self.temp_db_file, self.db_path, recursive=True)

    async def pull_from_s3(self) -> None:
        from cognee.infrastructure.files.storage.S3FileStorage import S3FileStorage

        s3_file_storage = S3FileStorage("")
        try:
            s3_file_storage.s3.get(self.db_path, self.temp_db_file, recursive=True)
        except FileNotFoundError:
            pass

    @asynccontextmanager
    async def get_async_session(self) -> AsyncGenerator[AsyncSession, None]:
        """
        Provide an asynchronous context manager for obtaining a database session.
        """
        async_session_maker = self.sessionmaker
        async with async_session_maker() as session:
            try:
                yield session
            finally:
                await session.close()  # Ensure the session is closed

    def get_session(self):
        """
        Provide a context manager for obtaining a database session.
        """
        session_maker = self.sessionmaker

Review comment (suggestion):
There should be no need to unpack two dictionaries here. We can simply put 
`connect_args={"timeout": 30, **final_connect_args}`.
Also worth noting that, with this order in the dictionary, if there is another timeout key in the `final_connect_args`, the first one will be overwritten. Here we probably want that, but just a note.

Code after:
        self.db_uri: str = connection_string

        # Use provided connect_args (already parsed from config)
        final_connect_args = connect_args or {}

        if "sqlite" in connection_string:
            [prefix, db_path] = connection_string.split("///")
            self.db_path = db_path

            if "s3://" in self.db_path:
                db_dir_path = path.dirname(self.db_path)
                file_storage = get_file_storage(db_dir_path)

                run_sync(file_storage.ensure_directory_exists())

                with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
                    self.temp_db_file = temp_file.name
                    connection_string = prefix + "///" + self.temp_db_file

                run_sync(self.pull_from_s3())

        if "sqlite" in connection_string:
            self.engine = create_async_engine(
                connection_string,
                poolclass=NullPool,
                connect_args={**{"timeout": 30}, **final_connect_args},
            )
        else:
            self.engine = create_async_engine(
                connection_string,
                pool_size=20,
                max_overflow=20,
                pool_recycle=280,
                pool_pre_ping=True,
                pool_timeout=280,
                connect_args=final_connect_args,
            )

        self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)

    async def push_to_s3(self) -> None:
        if os.getenv("STORAGE_BACKEND", "").lower() == "s3" and hasattr(self, "temp_db_file"):
            from cognee.infrastructure.files.storage.S3FileStorage import S3FileStorage

            s3_file_storage = S3FileStorage("")
            s3_file_storage.s3.put(self.temp_db_file, self.db_path, recursive=True)

    async def pull_from_s3(self) -> None:
        from cognee.infrastructure.files.storage.S3FileStorage import S3FileStorage

        s3_file_storage = S3FileStorage("")