fastapi-factory-utilities 0.3.6__py3-none-any.whl → 0.9.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (76) hide show
  1. fastapi_factory_utilities/core/api/__init__.py +1 -1
  2. fastapi_factory_utilities/core/api/v1/sys/health.py +1 -1
  3. fastapi_factory_utilities/core/app/__init__.py +4 -4
  4. fastapi_factory_utilities/core/app/application.py +22 -26
  5. fastapi_factory_utilities/core/app/builder.py +19 -35
  6. fastapi_factory_utilities/core/app/config.py +2 -0
  7. fastapi_factory_utilities/core/app/fastapi_builder.py +3 -2
  8. fastapi_factory_utilities/core/exceptions.py +64 -29
  9. fastapi_factory_utilities/core/plugins/__init__.py +2 -31
  10. fastapi_factory_utilities/core/plugins/abstracts.py +40 -0
  11. fastapi_factory_utilities/core/plugins/aiopika/__init__.py +25 -0
  12. fastapi_factory_utilities/core/plugins/aiopika/abstract.py +48 -0
  13. fastapi_factory_utilities/core/plugins/aiopika/configs.py +85 -0
  14. fastapi_factory_utilities/core/plugins/aiopika/depends.py +20 -0
  15. fastapi_factory_utilities/core/plugins/aiopika/exceptions.py +29 -0
  16. fastapi_factory_utilities/core/plugins/aiopika/exchange.py +69 -0
  17. fastapi_factory_utilities/core/plugins/aiopika/listener/__init__.py +7 -0
  18. fastapi_factory_utilities/core/plugins/aiopika/listener/abstract.py +72 -0
  19. fastapi_factory_utilities/core/plugins/aiopika/message.py +86 -0
  20. fastapi_factory_utilities/core/plugins/aiopika/plugins.py +84 -0
  21. fastapi_factory_utilities/core/plugins/aiopika/publisher/__init__.py +7 -0
  22. fastapi_factory_utilities/core/plugins/aiopika/publisher/abstract.py +66 -0
  23. fastapi_factory_utilities/core/plugins/aiopika/queue.py +88 -0
  24. fastapi_factory_utilities/core/plugins/odm_plugin/__init__.py +14 -157
  25. fastapi_factory_utilities/core/plugins/odm_plugin/builder.py +14 -29
  26. fastapi_factory_utilities/core/plugins/odm_plugin/configs.py +1 -1
  27. fastapi_factory_utilities/core/plugins/odm_plugin/depends.py +4 -3
  28. fastapi_factory_utilities/core/plugins/odm_plugin/documents.py +1 -1
  29. fastapi_factory_utilities/core/plugins/odm_plugin/helpers.py +16 -0
  30. fastapi_factory_utilities/core/plugins/odm_plugin/plugins.py +153 -0
  31. fastapi_factory_utilities/core/plugins/odm_plugin/repositories.py +17 -15
  32. fastapi_factory_utilities/core/plugins/opentelemetry_plugin/__init__.py +8 -115
  33. fastapi_factory_utilities/core/plugins/opentelemetry_plugin/instruments/__init__.py +85 -0
  34. fastapi_factory_utilities/core/plugins/opentelemetry_plugin/plugins.py +137 -0
  35. fastapi_factory_utilities/core/plugins/taskiq_plugins/__init__.py +31 -0
  36. fastapi_factory_utilities/core/plugins/taskiq_plugins/configs.py +12 -0
  37. fastapi_factory_utilities/core/plugins/taskiq_plugins/depends.py +51 -0
  38. fastapi_factory_utilities/core/plugins/taskiq_plugins/exceptions.py +13 -0
  39. fastapi_factory_utilities/core/plugins/taskiq_plugins/plugin.py +41 -0
  40. fastapi_factory_utilities/core/plugins/taskiq_plugins/schedulers.py +187 -0
  41. fastapi_factory_utilities/core/protocols.py +1 -54
  42. fastapi_factory_utilities/core/security/__init__.py +5 -0
  43. fastapi_factory_utilities/core/security/abstracts.py +42 -0
  44. fastapi_factory_utilities/core/security/jwt/__init__.py +43 -0
  45. fastapi_factory_utilities/core/security/jwt/configs.py +32 -0
  46. fastapi_factory_utilities/core/security/jwt/decoders.py +130 -0
  47. fastapi_factory_utilities/core/security/jwt/exceptions.py +23 -0
  48. fastapi_factory_utilities/core/security/jwt/objects.py +107 -0
  49. fastapi_factory_utilities/core/security/jwt/services.py +176 -0
  50. fastapi_factory_utilities/core/security/jwt/stores.py +43 -0
  51. fastapi_factory_utilities/core/security/jwt/types.py +9 -0
  52. fastapi_factory_utilities/core/security/jwt/verifiers.py +46 -0
  53. fastapi_factory_utilities/core/security/kratos.py +43 -43
  54. fastapi_factory_utilities/core/services/hydra/__init__.py +20 -0
  55. fastapi_factory_utilities/core/services/hydra/exceptions.py +15 -0
  56. fastapi_factory_utilities/core/services/hydra/objects.py +26 -0
  57. fastapi_factory_utilities/core/services/hydra/services.py +200 -0
  58. fastapi_factory_utilities/core/services/status/__init__.py +2 -2
  59. fastapi_factory_utilities/core/services/status/exceptions.py +1 -1
  60. fastapi_factory_utilities/core/utils/status.py +2 -1
  61. fastapi_factory_utilities/core/utils/yaml_reader.py +1 -1
  62. fastapi_factory_utilities/example/app.py +15 -5
  63. fastapi_factory_utilities/example/entities/books/__init__.py +1 -1
  64. fastapi_factory_utilities/example/models/books/__init__.py +1 -1
  65. {fastapi_factory_utilities-0.3.6.dist-info → fastapi_factory_utilities-0.9.1.dist-info}/METADATA +21 -15
  66. fastapi_factory_utilities-0.9.1.dist-info/RECORD +111 -0
  67. {fastapi_factory_utilities-0.3.6.dist-info → fastapi_factory_utilities-0.9.1.dist-info}/WHEEL +1 -1
  68. fastapi_factory_utilities/core/app/plugin_manager/__init__.py +0 -15
  69. fastapi_factory_utilities/core/app/plugin_manager/exceptions.py +0 -33
  70. fastapi_factory_utilities/core/app/plugin_manager/plugin_manager.py +0 -190
  71. fastapi_factory_utilities/core/plugins/example/__init__.py +0 -31
  72. fastapi_factory_utilities/core/plugins/httpx_plugin/__init__.py +0 -31
  73. fastapi_factory_utilities/core/security/jwt.py +0 -158
  74. fastapi_factory_utilities-0.3.6.dist-info/RECORD +0 -78
  75. {fastapi_factory_utilities-0.3.6.dist-info → fastapi_factory_utilities-0.9.1.dist-info}/entry_points.txt +0 -0
  76. {fastapi_factory_utilities-0.3.6.dist-info → fastapi_factory_utilities-0.9.1.dist-info/licenses}/LICENSE +0 -0
@@ -1,168 +1,25 @@
1
- """Oriented Data Model (ODM) plugin package."""
1
+ """ODM Plugin Module."""
2
2
 
3
- from logging import INFO, Logger, getLogger
4
- from typing import Any
5
-
6
- from beanie import init_beanie # pyright: ignore[reportUnknownVariableType]
7
- from motor.motor_asyncio import AsyncIOMotorClient
8
- from reactivex import Subject
9
- from structlog.stdlib import BoundLogger, get_logger
10
-
11
- from fastapi_factory_utilities.core.plugins import PluginState
12
- from fastapi_factory_utilities.core.protocols import ApplicationAbstractProtocol
13
- from fastapi_factory_utilities.core.services.status.enums import (
14
- ComponentTypeEnum,
15
- HealthStatusEnum,
16
- ReadinessStatusEnum,
17
- )
18
- from fastapi_factory_utilities.core.services.status.services import StatusService
19
- from fastapi_factory_utilities.core.services.status.types import (
20
- ComponentInstanceType,
21
- Status,
22
- )
23
-
24
- from .builder import ODMBuilder
25
3
  from .depends import depends_odm_client, depends_odm_database
26
4
  from .documents import BaseDocument
27
- from .exceptions import OperationError, UnableToCreateEntityDueToDuplicateKeyError
5
+ from .exceptions import (
6
+ ODMPluginBaseException,
7
+ ODMPluginConfigError,
8
+ OperationError,
9
+ UnableToCreateEntityDueToDuplicateKeyError,
10
+ )
11
+ from .helpers import PersistedEntity
12
+ from .plugins import ODMPlugin
28
13
  from .repositories import AbstractRepository
29
14
 
30
- _logger: BoundLogger = get_logger()
31
-
32
-
33
- def pre_conditions_check(application: ApplicationAbstractProtocol) -> bool:
34
- """Check the pre-conditions for the OpenTelemetry plugin.
35
-
36
- Args:
37
- application (BaseApplicationProtocol): The application.
38
-
39
- Returns:
40
- bool: True if the pre-conditions are met, False otherwise.
41
- """
42
- del application
43
- return True
44
-
45
-
46
- def on_load(
47
- application: ApplicationAbstractProtocol,
48
- ) -> list["PluginState"] | None:
49
- """Actions to perform on load for the OpenTelemetry plugin.
50
-
51
- Args:
52
- application (BaseApplicationProtocol): The application.
53
- """
54
- del application
55
- # Configure the pymongo logger to INFO level
56
- pymongo_logger: Logger = getLogger("pymongo")
57
- pymongo_logger.setLevel(INFO)
58
- _logger.debug("ODM plugin loaded.")
59
-
60
-
61
- async def on_startup(
62
- application: ApplicationAbstractProtocol,
63
- ) -> list["PluginState"] | None:
64
- """Actions to perform on startup for the ODM plugin.
65
-
66
- Args:
67
- application (BaseApplicationProtocol): The application.
68
- odm_config (ODMConfig): The ODM configuration.
69
-
70
- Returns:
71
- None
72
- """
73
- states: list[PluginState] = []
74
-
75
- status_service: StatusService = application.get_status_service()
76
- component_instance: ComponentInstanceType = ComponentInstanceType(
77
- component_type=ComponentTypeEnum.DATABASE, identifier="MongoDB"
78
- )
79
- monitoring_subject: Subject[Status] = status_service.register_component_instance(
80
- component_instance=component_instance
81
- )
82
-
83
- try:
84
- odm_factory: ODMBuilder = ODMBuilder(application=application).build_all()
85
- await odm_factory.wait_ping()
86
- except Exception as exception: # pylint: disable=broad-except
87
- _logger.error(f"ODM plugin failed to start. {exception}")
88
- # TODO: Report the error to the status_service
89
- # this will report the application as unhealthy
90
- monitoring_subject.on_next(
91
- value=Status(health=HealthStatusEnum.UNHEALTHY, readiness=ReadinessStatusEnum.NOT_READY)
92
- )
93
- return states
94
-
95
- if odm_factory.odm_database is None or odm_factory.odm_client is None:
96
- _logger.error(
97
- f"ODM plugin failed to start. Database: {odm_factory.odm_database} - " f"Client: {odm_factory.odm_client}"
98
- )
99
- # TODO: Report the error to the status_service
100
- # this will report the application as unhealthy
101
- monitoring_subject.on_next(
102
- value=Status(health=HealthStatusEnum.UNHEALTHY, readiness=ReadinessStatusEnum.NOT_READY)
103
- )
104
- return states
105
-
106
- # Add the ODM client and database to the application state
107
- states.append(
108
- PluginState(key="odm_client", value=odm_factory.odm_client),
109
- )
110
- states.append(
111
- PluginState(
112
- key="odm_database",
113
- value=odm_factory.odm_database,
114
- ),
115
- )
116
-
117
- # TODO: Find a better way to initialize beanie with the document models of the concrete application
118
- # through an hook in the application, a dynamis import ?
119
- try:
120
- await init_beanie(
121
- database=odm_factory.odm_database,
122
- document_models=application.ODM_DOCUMENT_MODELS,
123
- )
124
- except Exception as exception: # pylint: disable=broad-except
125
- _logger.error(f"ODM plugin failed to start. {exception}")
126
- # TODO: Report the error to the status_service
127
- # this will report the application as unhealthy
128
- monitoring_subject.on_next(
129
- value=Status(health=HealthStatusEnum.UNHEALTHY, readiness=ReadinessStatusEnum.NOT_READY)
130
- )
131
- return states
132
-
133
- _logger.info(
134
- f"ODM plugin started. Database: {odm_factory.odm_database.name} - "
135
- f"Client: {odm_factory.odm_client.address} - "
136
- f"Document models: {application.ODM_DOCUMENT_MODELS}"
137
- )
138
-
139
- monitoring_subject.on_next(value=Status(health=HealthStatusEnum.HEALTHY, readiness=ReadinessStatusEnum.READY))
140
-
141
- return states
142
-
143
-
144
- async def on_shutdown(application: ApplicationAbstractProtocol) -> None:
145
- """Actions to perform on shutdown for the ODM plugin.
146
-
147
- Args:
148
- application (BaseApplicationProtocol): The application.
149
-
150
- Returns:
151
- None
152
- """
153
- # Skip if the ODM plugin was not started correctly
154
- if not hasattr(application.get_asgi_app().state, "odm_client"):
155
- return
156
-
157
- client: AsyncIOMotorClient[Any] = application.get_asgi_app().state.odm_client
158
- client.close()
159
- _logger.debug("ODM plugin shutdown.")
160
-
161
-
162
15
  __all__: list[str] = [
163
- "BaseDocument",
164
16
  "AbstractRepository",
17
+ "BaseDocument",
18
+ "ODMPlugin",
19
+ "ODMPluginBaseException",
20
+ "ODMPluginConfigError",
165
21
  "OperationError",
22
+ "PersistedEntity",
166
23
  "UnableToCreateEntityDueToDuplicateKeyError",
167
24
  "depends_odm_client",
168
25
  "depends_odm_database",
@@ -3,7 +3,8 @@
3
3
  from typing import Any, ClassVar, Self
4
4
 
5
5
  from bson import CodecOptions
6
- from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
6
+ from pymongo.asynchronous.database import AsyncDatabase
7
+ from pymongo.asynchronous.mongo_client import AsyncMongoClient
7
8
  from pymongo.server_api import ServerApi, ServerApiVersion
8
9
  from structlog.stdlib import get_logger
9
10
 
@@ -45,22 +46,22 @@ class ODMBuilder:
45
46
  self,
46
47
  application: ApplicationAbstractProtocol,
47
48
  odm_config: ODMConfig | None = None,
48
- odm_client: AsyncIOMotorClient[Any] | None = None,
49
- odm_database: AsyncIOMotorDatabase[Any] | None = None,
49
+ odm_client: AsyncMongoClient[Any] | None = None,
50
+ odm_database: AsyncDatabase[Any] | None = None,
50
51
  ) -> None:
51
52
  """Initialize the ODMFactory.
52
53
 
53
54
  Args:
54
55
  application (BaseApplicationProtocol): The application.
55
56
  odm_config (ODMConfig): The ODM configuration for injection. (Default is None)
56
- odm_client (AsyncIOMotorClient): The ODM client for injection. (Default is None)
57
- odm_database (AsyncIOMotorDatabase): The ODM database for injection. (Default is None)
57
+ odm_client (AsyncMongoClient): The ODM client for injection. (Default is None)
58
+ odm_database (AsyncDatabase): The ODM database for injection. (Default is None)
58
59
 
59
60
  """
60
61
  self._application: ApplicationAbstractProtocol = application
61
62
  self._config: ODMConfig | None = odm_config
62
- self._odm_client: AsyncIOMotorClient[Any] | None = odm_client
63
- self._odm_database: AsyncIOMotorDatabase[Any] | None = odm_database
63
+ self._odm_client: AsyncMongoClient[Any] | None = odm_client
64
+ self._odm_database: AsyncDatabase[Any] | None = odm_database
64
65
 
65
66
  @property
66
67
  def config(self) -> ODMConfig | None:
@@ -72,20 +73,20 @@ class ODMBuilder:
72
73
  return self._config
73
74
 
74
75
  @property
75
- def odm_client(self) -> AsyncIOMotorClient[Any] | None:
76
+ def odm_client(self) -> AsyncMongoClient[Any] | None:
76
77
  """Provide the ODM client.
77
78
 
78
79
  Returns:
79
- AsyncIOMotorClient | None: The ODM client.
80
+ AsyncMongoClient | None: The ODM client.
80
81
  """
81
82
  return self._odm_client
82
83
 
83
84
  @property
84
- def odm_database(self) -> AsyncIOMotorDatabase[Any] | None:
85
+ def odm_database(self) -> AsyncDatabase[Any] | None:
85
86
  """Provide the ODM database.
86
87
 
87
88
  Returns:
88
- AsyncIOMotorDatabase | None: The ODM database.
89
+ AsyncDatabase | None: The ODM database.
89
90
  """
90
91
  return self._odm_database
91
92
 
@@ -174,7 +175,7 @@ class ODMBuilder:
174
175
  "build_odm_config method or through parameter."
175
176
  )
176
177
 
177
- self._odm_client = AsyncIOMotorClient(
178
+ self._odm_client = AsyncMongoClient(
178
179
  host=self._config.uri,
179
180
  connect=True,
180
181
  connectTimeoutMS=self._config.connection_timeout_ms,
@@ -215,7 +216,7 @@ class ODMBuilder:
215
216
 
216
217
  if self._odm_client is None:
217
218
  raise ODMPluginConfigError(
218
- "ODM client is not set. Provide the ODM client using " "build_client method or through parameter."
219
+ "ODM client is not set. Provide the ODM client using build_client method or through parameter."
219
220
  )
220
221
 
221
222
  self._odm_database = self._odm_client.get_database(
@@ -241,19 +242,3 @@ class ODMBuilder:
241
242
  self.build_database()
242
243
 
243
244
  return self
244
-
245
- async def wait_ping(self):
246
- """Wait for the ODM client to be ready.
247
-
248
- Returns:
249
- Self: The ODM factory.
250
- """
251
- if self._odm_client is None:
252
- raise ODMPluginConfigError(
253
- "ODM client is not set. Provide the ODM client using " "build_client method or through parameter."
254
- )
255
-
256
- try:
257
- await self._odm_client.admin.command("ping")
258
- except Exception as exception: # pylint: disable=broad-except
259
- raise ODMPluginConfigError("Unable to ping the ODM client.") from exception
@@ -12,4 +12,4 @@ class ODMConfig(BaseModel):
12
12
 
13
13
  database: str = "test"
14
14
 
15
- connection_timeout_ms: int = 1000
15
+ connection_timeout_ms: int = 4000
@@ -3,10 +3,11 @@
3
3
  from typing import Any
4
4
 
5
5
  from fastapi import Request
6
- from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase
6
+ from pymongo.asynchronous.database import AsyncDatabase
7
+ from pymongo.asynchronous.mongo_client import AsyncMongoClient
7
8
 
8
9
 
9
- def depends_odm_client(request: Request) -> AsyncIOMotorClient[Any]:
10
+ def depends_odm_client(request: Request) -> AsyncMongoClient[Any]:
10
11
  """Acquire the ODM client from the request.
11
12
 
12
13
  Args:
@@ -18,7 +19,7 @@ def depends_odm_client(request: Request) -> AsyncIOMotorClient[Any]:
18
19
  return request.app.state.odm_client
19
20
 
20
21
 
21
- def depends_odm_database(request: Request) -> AsyncIOMotorDatabase[Any]:
22
+ def depends_odm_database(request: Request) -> AsyncDatabase[Any]:
22
23
  """Acquire the ODM database from the request.
23
24
 
24
25
  Args:
@@ -13,7 +13,7 @@ class BaseDocument(Document):
13
13
  """Base document class."""
14
14
 
15
15
  # To be agnostic of MongoDN, we use UUID as the document ID.
16
- id: UUID = Field( # pyright: ignore[reportIncompatibleVariableOverride]
16
+ id: UUID = Field( # type: ignore
17
17
  default_factory=uuid4, description="The document ID."
18
18
  )
19
19
 
@@ -0,0 +1,16 @@
1
+ """Helper functions for ODM plugins."""
2
+
3
+ import datetime
4
+ import uuid
5
+
6
+ from pydantic import BaseModel, Field
7
+
8
+
9
+ class PersistedEntity(BaseModel):
10
+ """Base class for persisted entities."""
11
+
12
+ id: uuid.UUID = Field(default_factory=uuid.uuid4)
13
+
14
+ revision_id: uuid.UUID | None = Field(default=None)
15
+ created_at: datetime.datetime = Field(default_factory=datetime.datetime.now)
16
+ updated_at: datetime.datetime = Field(default_factory=datetime.datetime.now)
@@ -0,0 +1,153 @@
1
+ """Oriented Data Model (ODM) plugin package."""
2
+
3
+ from logging import INFO, Logger, getLogger
4
+ from typing import Any, Self, cast
5
+
6
+ from beanie import Document, init_beanie # pyright: ignore[reportUnknownVariableType]
7
+ from pymongo.asynchronous.database import AsyncDatabase
8
+ from pymongo.asynchronous.mongo_client import AsyncMongoClient
9
+ from reactivex import Subject
10
+ from structlog.stdlib import BoundLogger, get_logger
11
+
12
+ from fastapi_factory_utilities.core.plugins.abstracts import PluginAbstract
13
+ from fastapi_factory_utilities.core.protocols import ApplicationAbstractProtocol
14
+ from fastapi_factory_utilities.core.services.status.enums import (
15
+ ComponentTypeEnum,
16
+ HealthStatusEnum,
17
+ ReadinessStatusEnum,
18
+ )
19
+ from fastapi_factory_utilities.core.services.status.services import StatusService
20
+ from fastapi_factory_utilities.core.services.status.types import (
21
+ ComponentInstanceType,
22
+ Status,
23
+ )
24
+
25
+ from .builder import ODMBuilder
26
+ from .configs import ODMConfig
27
+ from .depends import depends_odm_client, depends_odm_database
28
+ from .documents import BaseDocument
29
+ from .exceptions import OperationError, UnableToCreateEntityDueToDuplicateKeyError
30
+ from .helpers import PersistedEntity
31
+ from .repositories import AbstractRepository
32
+
33
+ _logger: BoundLogger = get_logger()
34
+
35
+
36
+ class ODMPlugin(PluginAbstract):
37
+ """ODM plugin."""
38
+
39
+ def __init__(
40
+ self, document_models: list[type[Document]] | None = None, odm_config: ODMConfig | None = None
41
+ ) -> None:
42
+ """Initialize the ODM plugin."""
43
+ super().__init__()
44
+ self._component_instance: ComponentInstanceType | None = None
45
+ self._monitoring_subject: Subject[Status] | None = None
46
+ self._document_models: list[type[Document]] | None = document_models
47
+ self._odm_config: ODMConfig | None = odm_config
48
+ self._odm_client: AsyncMongoClient[Any] | None = None
49
+ self._odm_database: AsyncDatabase[Any] | None = None
50
+
51
+ def set_application(self, application: ApplicationAbstractProtocol) -> Self:
52
+ """Set the application."""
53
+ self._document_models = self._document_models or application.ODM_DOCUMENT_MODELS
54
+ return super().set_application(application)
55
+
56
+ def on_load(self) -> None:
57
+ """Actions to perform on load for the ODM plugin."""
58
+ # Configure the pymongo logger to INFO level
59
+
60
+ pymongo_logger: Logger = getLogger("pymongo")
61
+ pymongo_logger.setLevel(INFO)
62
+ _logger.debug("ODM plugin loaded.")
63
+
64
+ def _setup_status(self) -> None:
65
+ assert self._application is not None
66
+ status_service: StatusService = self._application.get_status_service()
67
+ self._component_instance = ComponentInstanceType(
68
+ component_type=ComponentTypeEnum.DATABASE, identifier="MongoDB"
69
+ )
70
+ self._monitoring_subject = status_service.register_component_instance(
71
+ component_instance=self._component_instance
72
+ )
73
+
74
+ async def _setup_beanie(self) -> None:
75
+ assert self._application is not None
76
+ assert self._odm_database is not None
77
+ assert self._document_models is not None
78
+ assert self._monitoring_subject is not None
79
+ # TODO: Find a better way to initialize beanie with the document models of the concrete application
80
+ # through an hook in the application, a dynamis import ?
81
+ try:
82
+ await init_beanie(
83
+ database=self._odm_database,
84
+ document_models=self._document_models,
85
+ )
86
+ except Exception as exception: # pylint: disable=broad-except
87
+ _logger.error(f"ODM plugin failed to start. {exception}")
88
+ # TODO: Report the error to the status_service
89
+ # this will report the application as unhealthy
90
+ self._monitoring_subject.on_next(
91
+ value=Status(health=HealthStatusEnum.UNHEALTHY, readiness=ReadinessStatusEnum.NOT_READY)
92
+ )
93
+
94
+ async def on_startup(self) -> None:
95
+ """Actions to perform on startup for the ODM plugin."""
96
+ host: str
97
+ port: int
98
+ assert self._application is not None
99
+ self._setup_status()
100
+ assert self._monitoring_subject is not None
101
+ assert self._component_instance is not None
102
+
103
+ try:
104
+ odm_factory: ODMBuilder = ODMBuilder(application=self._application, odm_config=self._odm_config).build_all()
105
+ assert odm_factory.odm_client is not None
106
+ assert odm_factory.odm_database is not None
107
+ assert (await odm_factory.odm_client.address) is not None
108
+ host, port = cast(tuple[str, int], await odm_factory.odm_client.address)
109
+ await odm_factory.odm_client.aconnect()
110
+ self._odm_database = odm_factory.odm_database
111
+ self._odm_client = odm_factory.odm_client
112
+ except Exception as exception: # pylint: disable=broad-except
113
+ _logger.error(f"ODM plugin failed to start. {exception}")
114
+ # TODO: Report the error to the status_service
115
+ # this will report the application as unhealthy
116
+ self._monitoring_subject.on_next(
117
+ value=Status(health=HealthStatusEnum.UNHEALTHY, readiness=ReadinessStatusEnum.NOT_READY)
118
+ )
119
+ return None
120
+
121
+ self._add_to_state(key="odm_client", value=odm_factory.odm_client)
122
+ self._add_to_state(key="odm_database", value=odm_factory.odm_database)
123
+
124
+ await self._setup_beanie()
125
+
126
+ assert self._odm_client is not None
127
+
128
+ _logger.info(
129
+ f"ODM plugin started. Database: {self._odm_database.name} - "
130
+ f"Client: {host}:{port} - "
131
+ f"Document models: {self._application.ODM_DOCUMENT_MODELS}"
132
+ )
133
+
134
+ self._monitoring_subject.on_next(
135
+ value=Status(health=HealthStatusEnum.HEALTHY, readiness=ReadinessStatusEnum.READY)
136
+ )
137
+
138
+ async def on_shutdown(self) -> None:
139
+ """Actions to perform on shutdown for the ODM plugin."""
140
+ if self._odm_client is not None:
141
+ await self._odm_client.close()
142
+ _logger.debug("ODM plugin shutdown.")
143
+
144
+
145
+ __all__: list[str] = [
146
+ "AbstractRepository",
147
+ "BaseDocument",
148
+ "OperationError",
149
+ "PersistedEntity",
150
+ "UnableToCreateEntityDueToDuplicateKeyError",
151
+ "depends_odm_client",
152
+ "depends_odm_database",
153
+ ]
@@ -8,8 +8,9 @@ from typing import Any, Generic, TypeVar, get_args
8
8
  from uuid import UUID
9
9
 
10
10
  from beanie import SortDirection
11
- from motor.motor_asyncio import AsyncIOMotorClientSession, AsyncIOMotorDatabase
12
11
  from pydantic import BaseModel
12
+ from pymongo.asynchronous.client_session import AsyncClientSession
13
+ from pymongo.asynchronous.database import AsyncDatabase
13
14
  from pymongo.errors import DuplicateKeyError, PyMongoError
14
15
  from pymongo.results import DeleteResult
15
16
 
@@ -43,28 +44,30 @@ def managed_session() -> Callable[[Callable[..., Any]], Callable[..., Any]]:
43
44
  class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
44
45
  """Abstract class for the repository."""
45
46
 
46
- def __init__(self, database: AsyncIOMotorDatabase[Any]) -> None:
47
+ def __init__(self, database: AsyncDatabase[Any]) -> None:
47
48
  """Initialize the repository."""
48
49
  super().__init__()
49
- self._database: AsyncIOMotorDatabase[Any] = database
50
+ self._database: AsyncDatabase[Any] = database
50
51
  # Retrieve the generic concrete types
51
52
  generic_args: tuple[Any, ...] = get_args(self.__orig_bases__[0]) # type: ignore
52
53
  self._document_type: type[DocumentGenericType] = generic_args[0]
53
54
  self._entity_type: type[EntityGenericType] = generic_args[1]
54
55
 
55
56
  @asynccontextmanager
56
- async def get_session(self) -> AsyncGenerator[AsyncIOMotorClientSession, None]:
57
+ async def get_session(self) -> AsyncGenerator[AsyncClientSession, None]:
57
58
  """Yield a new session."""
59
+ session: AsyncClientSession | None = None
58
60
  try:
59
- async with await self._database.client.start_session() as session:
60
- yield session
61
+ session = self._database.client.start_session()
62
+ yield session
61
63
  except PyMongoError as error:
62
64
  raise OperationError(f"Failed to create session: {error}") from error
65
+ finally:
66
+ if session is not None:
67
+ await session.end_session()
63
68
 
64
69
  @managed_session()
65
- async def insert(
66
- self, entity: EntityGenericType, session: AsyncIOMotorClientSession | None = None
67
- ) -> EntityGenericType:
70
+ async def insert(self, entity: EntityGenericType, session: AsyncClientSession | None = None) -> EntityGenericType:
68
71
  """Insert the entity into the database.
69
72
 
70
73
  Args:
@@ -104,9 +107,7 @@ class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
104
107
  return entity_created
105
108
 
106
109
  @managed_session()
107
- async def update(
108
- self, entity: EntityGenericType, session: AsyncIOMotorClientSession | None = None
109
- ) -> EntityGenericType:
110
+ async def update(self, entity: EntityGenericType, session: AsyncClientSession | None = None) -> EntityGenericType:
110
111
  """Update the entity in the database.
111
112
 
112
113
  Args:
@@ -145,7 +146,7 @@ class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
145
146
  async def get_one_by_id(
146
147
  self,
147
148
  entity_id: UUID,
148
- session: AsyncIOMotorClientSession | None = None,
149
+ session: AsyncClientSession | None = None,
149
150
  ) -> EntityGenericType | None:
150
151
  """Get the entity by its ID.
151
152
 
@@ -179,7 +180,7 @@ class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
179
180
 
180
181
  @managed_session()
181
182
  async def delete_one_by_id(
182
- self, entity_id: UUID, raise_if_not_found: bool = False, session: AsyncIOMotorClientSession | None = None
183
+ self, entity_id: UUID, raise_if_not_found: bool = False, session: AsyncClientSession | None = None
183
184
  ) -> None:
184
185
  """Delete a document by its ID.
185
186
 
@@ -224,7 +225,7 @@ class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
224
225
  skip: int | None = None,
225
226
  limit: int | None = None,
226
227
  sort: None | str | list[tuple[str, SortDirection]] = None,
227
- session: AsyncIOMotorClientSession | None = None,
228
+ session: AsyncClientSession | None = None,
228
229
  ignore_cache: bool = False,
229
230
  fetch_links: bool = False,
230
231
  lazy_parse: bool = False,
@@ -268,6 +269,7 @@ class AbstractRepository(ABC, Generic[DocumentGenericType, EntityGenericType]):
268
269
  lazy_parse=lazy_parse,
269
270
  nesting_depth=nesting_depth,
270
271
  nesting_depths_per_field=nesting_depths_per_field,
272
+ **pymongo_kwargs,
271
273
  ).to_list()
272
274
  except PyMongoError as error:
273
275
  raise OperationError(f"Failed to find documents: {error}") from error