reflex 0.8.14.post1__py3-none-any.whl → 0.8.15a0__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.

Potentially problematic release.


This version of reflex might be problematic. Click here for more details.

reflex/model.py CHANGED
@@ -5,69 +5,22 @@ from __future__ import annotations
5
5
  import re
6
6
  from collections import defaultdict
7
7
  from contextlib import suppress
8
- from typing import Any, ClassVar
9
-
10
- import alembic.autogenerate
11
- import alembic.command
12
- import alembic.config
13
- import alembic.operations.ops
14
- import alembic.runtime.environment
15
- import alembic.script
16
- import alembic.util
17
- import sqlalchemy
18
- import sqlalchemy.exc
19
- import sqlalchemy.ext.asyncio
20
- import sqlalchemy.orm
21
- from alembic.runtime.migration import MigrationContext
22
- from alembic.script.base import Script
23
-
24
- from reflex.base import Base
8
+ from importlib.util import find_spec
9
+ from typing import TYPE_CHECKING, Any, ClassVar
10
+
25
11
  from reflex.config import get_config
26
12
  from reflex.environment import environment
27
13
  from reflex.utils import console
28
- from reflex.utils.compat import sqlmodel, sqlmodel_field_has_primary_key
29
-
30
- _ENGINE: dict[str, sqlalchemy.engine.Engine] = {}
31
- _ASYNC_ENGINE: dict[str, sqlalchemy.ext.asyncio.AsyncEngine] = {}
32
- _AsyncSessionLocal: dict[str | None, sqlalchemy.ext.asyncio.async_sessionmaker] = {}
33
-
34
- # Import AsyncSession _after_ reflex.utils.compat
35
- from sqlmodel.ext.asyncio.session import AsyncSession # noqa: E402
36
-
37
-
38
- def format_revision(
39
- rev: Script,
40
- current_rev: str | None,
41
- current_reached_ref: list[bool],
42
- ) -> str:
43
- """Format a single revision for display.
44
-
45
- Args:
46
- rev: The alembic script object
47
- current_rev: The currently applied revision ID
48
- current_reached_ref: Mutable reference to track if we've reached current revision
49
-
50
- Returns:
51
- Formatted string for display
52
- """
53
- current = rev.revision
54
- message = rev.doc
55
-
56
- # Determine if this migration is applied
57
- if current_rev is None:
58
- is_applied = False
59
- elif current == current_rev:
60
- is_applied = True
61
- current_reached_ref[0] = True
62
- else:
63
- is_applied = not current_reached_ref[0]
14
+ from reflex.utils.compat import sqlmodel_field_has_primary_key
15
+ from reflex.utils.serializers import serializer
64
16
 
65
- # Show checkmark or X with colors
66
- status_icon = "[green]✓[/green]" if is_applied else "[red]✗[/red]"
67
- head_marker = " (head)" if rev.is_head else ""
17
+ if TYPE_CHECKING:
18
+ import sqlalchemy
19
+ import sqlmodel
68
20
 
69
- # Format output with message
70
- return f" [{status_icon}] {current}{head_marker}, {message}"
21
+ SQLModelOrSqlAlchemy = (
22
+ type[sqlmodel.SQLModel] | type[sqlalchemy.orm.DeclarativeBase]
23
+ )
71
24
 
72
25
 
73
26
  def _safe_db_url_for_logging(url: str) -> str:
@@ -82,536 +35,602 @@ def _safe_db_url_for_logging(url: str) -> str:
82
35
  return re.sub(r"://[^@]+@", "://<username>:<password>@", url)
83
36
 
84
37
 
85
- def get_engine_args(url: str | None = None) -> dict[str, Any]:
86
- """Get the database engine arguments.
87
-
88
- Args:
89
- url: The database url.
90
-
91
- Returns:
92
- The database engine arguments as a dict.
93
- """
94
- kwargs: dict[str, Any] = {
95
- # Print the SQL queries if the log level is INFO or lower.
96
- "echo": environment.SQLALCHEMY_ECHO.get(),
97
- # Check connections before returning them.
98
- "pool_pre_ping": environment.SQLALCHEMY_POOL_PRE_PING.get(),
99
- "pool_size": environment.SQLALCHEMY_POOL_SIZE.get(),
100
- "max_overflow": environment.SQLALCHEMY_MAX_OVERFLOW.get(),
101
- "pool_recycle": environment.SQLALCHEMY_POOL_RECYCLE.get(),
102
- "pool_timeout": environment.SQLALCHEMY_POOL_TIMEOUT.get(),
103
- }
104
- conf = get_config()
105
- url = url or conf.db_url
106
- if url is not None and url.startswith("sqlite"):
107
- # Needed for the admin dash on sqlite.
108
- kwargs["connect_args"] = {"check_same_thread": False}
109
- return kwargs
110
-
111
-
112
- def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
113
- """Get the database engine.
114
-
115
- Args:
116
- url: the DB url to use.
117
-
118
- Returns:
119
- The database engine.
120
-
121
- Raises:
122
- ValueError: If the database url is None.
123
- """
124
- conf = get_config()
125
- url = url or conf.db_url
126
- if url is None:
127
- msg = "No database url configured"
128
- raise ValueError(msg)
129
-
130
- global _ENGINE
131
- if url in _ENGINE:
132
- return _ENGINE[url]
133
-
134
- if not environment.ALEMBIC_CONFIG.get().exists():
135
- console.warn(
136
- "Database is not initialized, run [bold]reflex db init[/bold] first."
137
- )
138
- _ENGINE[url] = sqlmodel.create_engine(
139
- url,
140
- **get_engine_args(url),
38
+ def _print_db_not_available(*args, **kwargs):
39
+ msg = (
40
+ "Database is not available. Please install the required packages: "
41
+ "`pip install reflex[db]`."
141
42
  )
142
- return _ENGINE[url]
143
-
43
+ raise ImportError(msg)
144
44
 
145
- def get_async_engine(url: str | None) -> sqlalchemy.ext.asyncio.AsyncEngine:
146
- """Get the async database engine.
147
45
 
148
- Args:
149
- url: The database url.
150
-
151
- Returns:
152
- The async database engine.
153
-
154
- Raises:
155
- ValueError: If the async database url is None.
156
- """
157
- if url is None:
158
- conf = get_config()
159
- url = conf.async_db_url
160
- if url is not None and conf.db_url is not None:
161
- async_db_url_tail = url.partition("://")[2]
162
- db_url_tail = conf.db_url.partition("://")[2]
163
- if async_db_url_tail != db_url_tail:
164
- console.warn(
165
- f"async_db_url `{_safe_db_url_for_logging(url)}` "
166
- "should reference the same database as "
167
- f"db_url `{_safe_db_url_for_logging(conf.db_url)}`."
168
- )
169
- if url is None:
170
- msg = "No async database url configured"
171
- raise ValueError(msg)
172
-
173
- global _ASYNC_ENGINE
174
- if url in _ASYNC_ENGINE:
175
- return _ASYNC_ENGINE[url]
46
+ class _ClassThatErrorsOnInit:
47
+ def __init__(self, *args, **kwargs):
48
+ _print_db_not_available(*args, **kwargs)
176
49
 
177
- if not environment.ALEMBIC_CONFIG.get().exists():
178
- console.warn(
179
- "Database is not initialized, run [bold]reflex db init[/bold] first."
180
- )
181
- _ASYNC_ENGINE[url] = sqlalchemy.ext.asyncio.create_async_engine(
182
- url,
183
- **get_engine_args(url),
184
- )
185
- return _ASYNC_ENGINE[url]
186
50
 
51
+ if find_spec("sqlalchemy"):
52
+ import sqlalchemy
53
+ import sqlalchemy.exc
54
+ import sqlalchemy.ext.asyncio
55
+ import sqlalchemy.orm
187
56
 
188
- async def get_db_status() -> dict[str, bool]:
189
- """Checks the status of the database connection.
57
+ _ENGINE: dict[str, sqlalchemy.engine.Engine] = {}
58
+ _ASYNC_ENGINE: dict[str, sqlalchemy.ext.asyncio.AsyncEngine] = {}
190
59
 
191
- Attempts to connect to the database and execute a simple query to verify connectivity.
60
+ def get_engine_args(url: str | None = None) -> dict[str, Any]:
61
+ """Get the database engine arguments.
192
62
 
193
- Returns:
194
- The status of the database connection.
195
- """
196
- status = True
197
- try:
198
- engine = get_engine()
199
- with engine.connect() as connection:
200
- connection.execute(sqlalchemy.text("SELECT 1"))
201
- except sqlalchemy.exc.OperationalError:
202
- status = False
63
+ Args:
64
+ url: The database url.
203
65
 
204
- return {"db": status}
66
+ Returns:
67
+ The database engine arguments as a dict.
68
+ """
69
+ kwargs: dict[str, Any] = {
70
+ # Print the SQL queries if the log level is INFO or lower.
71
+ "echo": environment.SQLALCHEMY_ECHO.get(),
72
+ # Check connections before returning them.
73
+ "pool_pre_ping": environment.SQLALCHEMY_POOL_PRE_PING.get(),
74
+ }
75
+ conf = get_config()
76
+ url = url or conf.db_url
77
+ if url is not None and url.startswith("sqlite"):
78
+ # Needed for the admin dash on sqlite.
79
+ kwargs["connect_args"] = {"check_same_thread": False}
80
+ return kwargs
205
81
 
82
+ def get_engine(url: str | None = None) -> sqlalchemy.engine.Engine:
83
+ """Get the database engine.
206
84
 
207
- SQLModelOrSqlAlchemy = type[sqlmodel.SQLModel] | type[sqlalchemy.orm.DeclarativeBase]
85
+ Args:
86
+ url: the DB url to use.
208
87
 
88
+ Returns:
89
+ The database engine.
209
90
 
210
- class ModelRegistry:
211
- """Registry for all models."""
91
+ Raises:
92
+ ValueError: If the database url is None.
93
+ """
94
+ conf = get_config()
95
+ url = url or conf.db_url
96
+ if url is None:
97
+ msg = "No database url configured"
98
+ raise ValueError(msg)
212
99
 
213
- models: ClassVar[set[SQLModelOrSqlAlchemy]] = set()
100
+ global _ENGINE
101
+ if url in _ENGINE:
102
+ return _ENGINE[url]
214
103
 
215
- # Cache the metadata to avoid re-creating it.
216
- _metadata: ClassVar[sqlalchemy.MetaData | None] = None
104
+ if not environment.ALEMBIC_CONFIG.get().exists():
105
+ console.warn(
106
+ "Database is not initialized, run [bold]reflex db init[/bold] first."
107
+ )
108
+ _ENGINE[url] = sqlalchemy.engine.create_engine(
109
+ url,
110
+ **get_engine_args(url),
111
+ )
112
+ return _ENGINE[url]
217
113
 
218
- @classmethod
219
- def register(cls, model: SQLModelOrSqlAlchemy):
220
- """Register a model. Can be used directly or as a decorator.
114
+ def get_async_engine(url: str | None) -> sqlalchemy.ext.asyncio.AsyncEngine:
115
+ """Get the async database engine.
221
116
 
222
117
  Args:
223
- model: The model to register.
118
+ url: The database url.
224
119
 
225
120
  Returns:
226
- The model passed in as an argument (Allows decorator usage)
227
- """
228
- cls.models.add(model)
229
- return model
230
-
231
- @classmethod
232
- def get_models(cls, include_empty: bool = False) -> set[SQLModelOrSqlAlchemy]:
233
- """Get registered models.
234
-
235
- Args:
236
- include_empty: If True, include models with empty metadata.
121
+ The async database engine.
237
122
 
238
- Returns:
239
- The registered models.
123
+ Raises:
124
+ ValueError: If the async database url is None.
240
125
  """
241
- if include_empty:
242
- return cls.models
243
- return {
244
- model for model in cls.models if not cls._model_metadata_is_empty(model)
245
- }
126
+ if url is None:
127
+ conf = get_config()
128
+ url = conf.async_db_url
129
+ if url is not None and conf.db_url is not None:
130
+ async_db_url_tail = url.partition("://")[2]
131
+ db_url_tail = conf.db_url.partition("://")[2]
132
+ if async_db_url_tail != db_url_tail:
133
+ console.warn(
134
+ f"async_db_url `{_safe_db_url_for_logging(url)}` "
135
+ "should reference the same database as "
136
+ f"db_url `{_safe_db_url_for_logging(conf.db_url)}`."
137
+ )
138
+ if url is None:
139
+ msg = "No async database url configured"
140
+ raise ValueError(msg)
141
+
142
+ global _ASYNC_ENGINE
143
+ if url in _ASYNC_ENGINE:
144
+ return _ASYNC_ENGINE[url]
145
+
146
+ if not environment.ALEMBIC_CONFIG.get().exists():
147
+ console.warn(
148
+ "Database is not initialized, run [bold]reflex db init[/bold] first."
149
+ )
150
+ _ASYNC_ENGINE[url] = sqlalchemy.ext.asyncio.create_async_engine(
151
+ url,
152
+ **get_engine_args(url),
153
+ )
154
+ return _ASYNC_ENGINE[url]
246
155
 
247
- @staticmethod
248
- def _model_metadata_is_empty(model: SQLModelOrSqlAlchemy) -> bool:
249
- """Check if the model metadata is empty.
156
+ def sqla_session(url: str | None = None) -> sqlalchemy.orm.Session:
157
+ """Get a bare sqlalchemy session to interact with the database.
250
158
 
251
159
  Args:
252
- model: The model to check.
160
+ url: The database url.
253
161
 
254
162
  Returns:
255
- True if the model metadata is empty, False otherwise.
163
+ A database session.
256
164
  """
257
- return len(model.metadata.tables) == 0
165
+ return sqlalchemy.orm.Session(get_engine(url))
166
+
167
+ class ModelRegistry:
168
+ """Registry for all models."""
169
+
170
+ models: ClassVar[set[SQLModelOrSqlAlchemy]] = set()
171
+
172
+ # Cache the metadata to avoid re-creating it.
173
+ _metadata: ClassVar[sqlalchemy.MetaData | None] = None
174
+
175
+ @classmethod
176
+ def register(cls, model: SQLModelOrSqlAlchemy):
177
+ """Register a model. Can be used directly or as a decorator.
178
+
179
+ Args:
180
+ model: The model to register.
181
+
182
+ Returns:
183
+ The model passed in as an argument (Allows decorator usage)
184
+ """
185
+ cls.models.add(model)
186
+ return model
187
+
188
+ @classmethod
189
+ def get_models(cls, include_empty: bool = False) -> set[SQLModelOrSqlAlchemy]:
190
+ """Get registered models.
191
+
192
+ Args:
193
+ include_empty: If True, include models with empty metadata.
194
+
195
+ Returns:
196
+ The registered models.
197
+ """
198
+ if include_empty:
199
+ return cls.models
200
+ return {
201
+ model for model in cls.models if not cls._model_metadata_is_empty(model)
202
+ }
203
+
204
+ @staticmethod
205
+ def _model_metadata_is_empty(model: SQLModelOrSqlAlchemy) -> bool:
206
+ """Check if the model metadata is empty.
207
+
208
+ Args:
209
+ model: The model to check.
210
+
211
+ Returns:
212
+ True if the model metadata is empty, False otherwise.
213
+ """
214
+ return len(model.metadata.tables) == 0
215
+
216
+ @classmethod
217
+ def get_metadata(cls) -> sqlalchemy.MetaData:
218
+ """Get the database metadata.
219
+
220
+ Returns:
221
+ The database metadata.
222
+ """
223
+ if cls._metadata is not None:
224
+ return cls._metadata
225
+
226
+ models = cls.get_models(include_empty=False)
227
+
228
+ if len(models) == 1:
229
+ metadata = next(iter(models)).metadata
230
+ else:
231
+ # Merge the metadata from all the models.
232
+ # This allows mixing bare sqlalchemy models with sqlmodel models in one database.
233
+ metadata = sqlalchemy.MetaData()
234
+ for model in cls.get_models():
235
+ for table in model.metadata.tables.values():
236
+ table.to_metadata(metadata)
237
+
238
+ # Cache the metadata
239
+ cls._metadata = metadata
240
+
241
+ return metadata
242
+
243
+ else:
244
+ get_engine_args = _print_db_not_available
245
+ get_engine = _print_db_not_available
246
+ get_async_engine = _print_db_not_available
247
+ sqla_session = _print_db_not_available
248
+ ModelRegistry = _ClassThatErrorsOnInit # pyright: ignore [reportAssignmentType]
249
+
250
+ if find_spec("sqlmodel") and find_spec("sqlalchemy") and find_spec("pydantic"):
251
+ import alembic.autogenerate
252
+ import alembic.command
253
+ import alembic.config
254
+ import alembic.operations.ops
255
+ import alembic.runtime.environment
256
+ import alembic.script
257
+ import sqlmodel
258
+ from alembic.runtime.migration import MigrationContext
259
+ from alembic.script.base import Script
260
+ from sqlmodel.ext.asyncio.session import AsyncSession
261
+
262
+ _AsyncSessionLocal: dict[str | None, sqlalchemy.ext.asyncio.async_sessionmaker] = {}
263
+
264
+ def format_revision(
265
+ rev: Script,
266
+ current_rev: str | None,
267
+ current_reached_ref: list[bool],
268
+ ) -> str:
269
+ """Format a single revision for display.
258
270
 
259
- @classmethod
260
- def get_metadata(cls) -> sqlalchemy.MetaData:
261
- """Get the database metadata.
271
+ Args:
272
+ rev: The alembic script object
273
+ current_rev: The currently applied revision ID
274
+ current_reached_ref: Mutable reference to track if we've reached current revision
262
275
 
263
276
  Returns:
264
- The database metadata.
277
+ Formatted string for display
265
278
  """
266
- if cls._metadata is not None:
267
- return cls._metadata
268
-
269
- models = cls.get_models(include_empty=False)
270
-
271
- if len(models) == 1:
272
- metadata = next(iter(models)).metadata
279
+ current = rev.revision
280
+ message = rev.doc
281
+
282
+ # Determine if this migration is applied
283
+ if current_rev is None:
284
+ is_applied = False
285
+ elif current == current_rev:
286
+ is_applied = True
287
+ current_reached_ref[0] = True
273
288
  else:
274
- # Merge the metadata from all the models.
275
- # This allows mixing bare sqlalchemy models with sqlmodel models in one database.
276
- metadata = sqlalchemy.MetaData()
277
- for model in cls.get_models():
278
- for table in model.metadata.tables.values():
279
- table.to_metadata(metadata)
280
-
281
- # Cache the metadata
282
- cls._metadata = metadata
283
-
284
- return metadata
289
+ is_applied = not current_reached_ref[0]
285
290
 
291
+ # Show checkmark or X with colors
292
+ status_icon = "[green]✓[/green]" if is_applied else "[red]✗[/red]"
293
+ head_marker = " (head)" if rev.is_head else ""
286
294
 
287
- class Model(Base, sqlmodel.SQLModel): # pyright: ignore [reportGeneralTypeIssues,reportIncompatibleVariableOverride]
288
- """Base class to define a table in the database."""
295
+ # Format output with message
296
+ return f" [{status_icon}] {current}{head_marker}, {message}"
289
297
 
290
- # The primary key for the table.
291
- id: int | None = sqlmodel.Field(default=None, primary_key=True)
298
+ async def get_db_status() -> dict[str, bool]:
299
+ """Checks the status of the database connection.
292
300
 
293
- def __init_subclass__(cls):
294
- """Drop the default primary key field if any primary key field is defined."""
295
- non_default_primary_key_fields = [
296
- field_name
297
- for field_name, field in cls.__fields__.items()
298
- if field_name != "id" and sqlmodel_field_has_primary_key(field)
299
- ]
300
- if non_default_primary_key_fields:
301
- cls.__fields__.pop("id", None)
302
-
303
- super().__init_subclass__()
304
-
305
- @classmethod
306
- def _dict_recursive(cls, value: Any):
307
- """Recursively serialize the relationship object(s).
308
-
309
- Args:
310
- value: The value to serialize.
301
+ Attempts to connect to the database and execute a simple query to verify connectivity.
311
302
 
312
303
  Returns:
313
- The serialized value.
304
+ The status of the database connection.
314
305
  """
315
- if hasattr(value, "dict"):
316
- return value.dict()
317
- if isinstance(value, list):
318
- return [cls._dict_recursive(item) for item in value]
319
- return value
306
+ status = True
307
+ try:
308
+ engine = get_engine()
309
+ with engine.connect() as connection:
310
+ connection.execute(sqlalchemy.text("SELECT 1"))
311
+ except sqlalchemy.exc.OperationalError:
312
+ status = False
313
+
314
+ return {"db": status}
320
315
 
321
- def dict(self, **kwargs):
322
- """Convert the object to a dictionary.
316
+ @serializer
317
+ def serialize_sqlmodel(m: sqlmodel.SQLModel) -> dict[str, Any]:
318
+ """Serialize a SQLModel object to a dictionary.
323
319
 
324
320
  Args:
325
- kwargs: Ignored but needed for compatibility.
321
+ m: The SQLModel object to serialize.
326
322
 
327
323
  Returns:
328
- The object as a dictionary.
324
+ The serialized object as a dictionary.
329
325
  """
330
- base_fields = {name: getattr(self, name) for name in self.__fields__}
326
+ base_fields = m.model_dump()
331
327
  relationships = {}
332
328
  # SQLModel relationships do not appear in __fields__, but should be included if present.
333
- for name in self.__sqlmodel_relationships__:
329
+ for name in m.__sqlmodel_relationships__:
334
330
  with suppress(
335
331
  sqlalchemy.orm.exc.DetachedInstanceError # This happens when the relationship was never loaded and the session is closed.
336
332
  ):
337
- relationships[name] = self._dict_recursive(getattr(self, name))
333
+ relationships[name] = getattr(m, name)
338
334
  return {
339
335
  **base_fields,
340
336
  **relationships,
341
337
  }
342
338
 
343
- @staticmethod
344
- def create_all():
345
- """Create all the tables."""
346
- engine = get_engine()
347
- ModelRegistry.get_metadata().create_all(engine)
339
+ class Model(sqlmodel.SQLModel):
340
+ """Base class to define a table in the database."""
348
341
 
349
- @staticmethod
350
- def get_db_engine():
351
- """Get the database engine.
352
-
353
- Returns:
354
- The database engine.
355
- """
356
- return get_engine()
342
+ # The primary key for the table.
343
+ id: int | None = sqlmodel.Field(default=None, primary_key=True)
357
344
 
358
- @staticmethod
359
- def _alembic_config():
360
- """Get the alembic configuration and script_directory.
345
+ model_config = { # pyright: ignore [reportAssignmentType]
346
+ "arbitrary_types_allowed": True,
347
+ "use_enum_values": True,
348
+ "extra": "allow",
349
+ }
361
350
 
362
- Returns:
363
- tuple of (config, script_directory)
364
- """
365
- config = alembic.config.Config(environment.ALEMBIC_CONFIG.get())
366
- if not config.get_main_option("script_location"):
367
- config.set_main_option("script_location", "version")
368
- return config, alembic.script.ScriptDirectory.from_config(config)
351
+ @classmethod
352
+ def __pydantic_init_subclass__(cls):
353
+ """Drop the default primary key field if any primary key field is defined."""
354
+ non_default_primary_key_fields = [
355
+ field_name
356
+ for field_name, field_info in cls.model_fields.items()
357
+ if field_name != "id" and sqlmodel_field_has_primary_key(field_info)
358
+ ]
359
+ if non_default_primary_key_fields:
360
+ cls.model_fields.pop("id", None)
361
+ console.deprecate(
362
+ feature_name="Overriding default primary key",
363
+ reason=(
364
+ "Register sqlmodel.SQLModel classes with `@rx.ModelRegistry.register`"
365
+ ),
366
+ deprecation_version="0.8.0",
367
+ removal_version="0.9.0",
368
+ )
369
+ super().__pydantic_init_subclass__()
370
+
371
+ @staticmethod
372
+ def create_all():
373
+ """Create all the tables."""
374
+ engine = get_engine()
375
+ ModelRegistry.get_metadata().create_all(engine)
376
+
377
+ @staticmethod
378
+ def get_db_engine():
379
+ """Get the database engine.
380
+
381
+ Returns:
382
+ The database engine.
383
+ """
384
+ return get_engine()
385
+
386
+ @staticmethod
387
+ def _alembic_config():
388
+ """Get the alembic configuration and script_directory.
389
+
390
+ Returns:
391
+ tuple of (config, script_directory)
392
+ """
393
+ config = alembic.config.Config(environment.ALEMBIC_CONFIG.get())
394
+ if not config.get_main_option("script_location"):
395
+ config.set_main_option("script_location", "version")
396
+ return config, alembic.script.ScriptDirectory.from_config(config)
397
+
398
+ @staticmethod
399
+ def _alembic_render_item(
400
+ type_: str,
401
+ obj: Any,
402
+ autogen_context: alembic.autogenerate.api.AutogenContext,
403
+ ):
404
+ """Alembic render_item hook call.
369
405
 
370
- @staticmethod
371
- def _alembic_render_item(
372
- type_: str,
373
- obj: Any,
374
- autogen_context: alembic.autogenerate.api.AutogenContext,
375
- ):
376
- """Alembic render_item hook call.
406
+ This method is called to provide python code for the given obj,
407
+ but currently it is only used to add `sqlmodel` to the import list
408
+ when generating migration scripts.
377
409
 
378
- This method is called to provide python code for the given obj,
379
- but currently it is only used to add `sqlmodel` to the import list
380
- when generating migration scripts.
410
+ See https://alembic.sqlalchemy.org/en/latest/api/runtime.html
381
411
 
382
- See https://alembic.sqlalchemy.org/en/latest/api/runtime.html
412
+ Args:
413
+ type_: One of "schema", "table", "column", "index",
414
+ "unique_constraint", or "foreign_key_constraint".
415
+ obj: The object being rendered.
416
+ autogen_context: Shared AutogenContext passed to each render_item call.
383
417
 
384
- Args:
385
- type_: One of "schema", "table", "column", "index",
386
- "unique_constraint", or "foreign_key_constraint".
387
- obj: The object being rendered.
388
- autogen_context: Shared AutogenContext passed to each render_item call.
418
+ Returns:
419
+ False - Indicating that the default rendering should be used.
420
+ """
421
+ autogen_context.imports.add("import sqlmodel")
422
+ return False
389
423
 
390
- Returns:
391
- False - Indicating that the default rendering should be used.
392
- """
393
- autogen_context.imports.add("import sqlmodel")
394
- return False
395
-
396
- @classmethod
397
- def alembic_init(cls):
398
- """Initialize alembic for the project."""
399
- alembic.command.init(
400
- config=alembic.config.Config(environment.ALEMBIC_CONFIG.get()),
401
- directory=str(environment.ALEMBIC_CONFIG.get().parent / "alembic"),
402
- )
424
+ @classmethod
425
+ def alembic_init(cls):
426
+ """Initialize alembic for the project."""
427
+ alembic.command.init(
428
+ config=alembic.config.Config(environment.ALEMBIC_CONFIG.get()),
429
+ directory=str(environment.ALEMBIC_CONFIG.get().parent / "alembic"),
430
+ )
403
431
 
404
- @classmethod
405
- def get_migration_history(cls):
406
- """Get migration history with current database state.
432
+ @classmethod
433
+ def get_migration_history(cls):
434
+ """Get migration history with current database state.
435
+
436
+ Returns:
437
+ tuple: (current_revision, revisions_list) where revisions_list is in chronological order
438
+ """
439
+ # Get current revision from database
440
+ with cls.get_db_engine().connect() as connection:
441
+ context = MigrationContext.configure(connection)
442
+ current_rev = context.get_current_revision()
443
+
444
+ # Get all revisions from base to head
445
+ _, script_dir = cls._alembic_config()
446
+ revisions = list(script_dir.walk_revisions())
447
+ revisions.reverse() # Reverse to get chronological order (base first)
448
+
449
+ return current_rev, revisions
450
+
451
+ @classmethod
452
+ def alembic_autogenerate(
453
+ cls,
454
+ connection: sqlalchemy.engine.Connection,
455
+ message: str | None = None,
456
+ write_migration_scripts: bool = True,
457
+ ) -> bool:
458
+ """Generate migration scripts for alembic-detectable changes.
459
+
460
+ Args:
461
+ connection: SQLAlchemy connection to use when detecting changes.
462
+ message: Human readable identifier describing the generated revision.
463
+ write_migration_scripts: If True, write autogenerated revisions to script directory.
464
+
465
+ Returns:
466
+ True when changes have been detected.
467
+ """
468
+ if not environment.ALEMBIC_CONFIG.get().exists():
469
+ return False
470
+
471
+ config, script_directory = cls._alembic_config()
472
+ revision_context = alembic.autogenerate.api.RevisionContext(
473
+ config=config,
474
+ script_directory=script_directory,
475
+ command_args=defaultdict(
476
+ lambda: None,
477
+ autogenerate=True,
478
+ head="head",
479
+ message=message,
480
+ ),
481
+ )
482
+ writer = alembic.autogenerate.rewriter.Rewriter()
407
483
 
408
- Returns:
409
- tuple: (current_revision, revisions_list) where revisions_list is in chronological order
410
- """
411
- # Get current revision from database
412
- with cls.get_db_engine().connect() as connection:
413
- context = MigrationContext.configure(connection)
414
- current_rev = context.get_current_revision()
415
-
416
- # Get all revisions from base to head
417
- _, script_dir = cls._alembic_config()
418
- revisions = list(script_dir.walk_revisions())
419
- revisions.reverse() # Reverse to get chronological order (base first)
420
-
421
- return current_rev, revisions
422
-
423
- @classmethod
424
- def alembic_autogenerate(
425
- cls,
426
- connection: sqlalchemy.engine.Connection,
427
- message: str | None = None,
428
- write_migration_scripts: bool = True,
429
- ) -> bool:
430
- """Generate migration scripts for alembic-detectable changes.
484
+ @writer.rewrites(alembic.operations.ops.AddColumnOp)
485
+ def render_add_column_with_server_default(
486
+ context: MigrationContext,
487
+ revision: str | None,
488
+ op: Any,
489
+ ):
490
+ # Carry the sqlmodel default as server_default so that newly added
491
+ # columns get the desired default value in existing rows.
492
+ if op.column.default is not None and op.column.server_default is None:
493
+ op.column.server_default = sqlalchemy.DefaultClause(
494
+ sqlalchemy.sql.expression.literal(op.column.default.arg),
495
+ )
496
+ return op
497
+
498
+ def run_autogenerate(rev: str, context: MigrationContext):
499
+ revision_context.run_autogenerate(rev, context)
500
+ return []
501
+
502
+ with alembic.runtime.environment.EnvironmentContext(
503
+ config=config,
504
+ script=script_directory,
505
+ fn=run_autogenerate,
506
+ ) as env:
507
+ env.configure(
508
+ connection=connection,
509
+ target_metadata=ModelRegistry.get_metadata(),
510
+ render_item=cls._alembic_render_item,
511
+ process_revision_directives=writer,
512
+ compare_type=False,
513
+ render_as_batch=True, # for sqlite compatibility
514
+ )
515
+ env.run_migrations()
516
+ changes_detected = False
517
+ if revision_context.generated_revisions:
518
+ upgrade_ops = revision_context.generated_revisions[-1].upgrade_ops
519
+ if upgrade_ops is not None:
520
+ changes_detected = bool(upgrade_ops.ops)
521
+ if changes_detected and write_migration_scripts:
522
+ # Must iterate the generator to actually write the scripts.
523
+ _ = tuple(revision_context.generate_scripts())
524
+ return changes_detected
525
+
526
+ @classmethod
527
+ def _alembic_upgrade(
528
+ cls,
529
+ connection: sqlalchemy.engine.Connection,
530
+ to_rev: str = "head",
531
+ ) -> None:
532
+ """Apply alembic migrations up to the given revision.
533
+
534
+ Args:
535
+ connection: SQLAlchemy connection to use when performing upgrade.
536
+ to_rev: Revision to migrate towards.
537
+ """
538
+ config, script_directory = cls._alembic_config()
539
+
540
+ def run_upgrade(rev: str, context: MigrationContext):
541
+ return script_directory._upgrade_revs(to_rev, rev)
542
+
543
+ with alembic.runtime.environment.EnvironmentContext(
544
+ config=config,
545
+ script=script_directory,
546
+ fn=run_upgrade,
547
+ ) as env:
548
+ env.configure(connection=connection)
549
+ env.run_migrations()
550
+
551
+ @classmethod
552
+ def migrate(cls, autogenerate: bool = False) -> bool | None:
553
+ """Execute alembic migrations for all sqlmodel Model classes.
554
+
555
+ If alembic is not installed or has not been initialized for the project,
556
+ then no action is performed.
557
+
558
+ If there are no revisions currently tracked by alembic, then
559
+ an initial revision will be created based on sqlmodel metadata.
560
+
561
+ If models in the app have changed in incompatible ways that alembic
562
+ cannot automatically generate revisions for, the app may not be able to
563
+ start up until migration scripts have been corrected by hand.
564
+
565
+ Args:
566
+ autogenerate: If True, generate migration script and use it to upgrade schema
567
+ (otherwise, just bring the schema to current "head" revision).
568
+
569
+ Returns:
570
+ True - indicating the process was successful.
571
+ None - indicating the process was skipped.
572
+ """
573
+ if not environment.ALEMBIC_CONFIG.get().exists():
574
+ return None
575
+
576
+ with cls.get_db_engine().connect() as connection:
577
+ cls._alembic_upgrade(connection=connection)
578
+ if autogenerate:
579
+ changes_detected = cls.alembic_autogenerate(connection=connection)
580
+ if changes_detected:
581
+ cls._alembic_upgrade(connection=connection)
582
+ connection.commit()
583
+ return True
584
+
585
+ @classmethod
586
+ def select(cls):
587
+ """Select rows from the table.
588
+
589
+ Returns:
590
+ The select statement.
591
+ """
592
+ return sqlmodel.select(cls)
593
+
594
+ ModelRegistry.register(Model)
595
+
596
+ def session(url: str | None = None) -> sqlmodel.Session:
597
+ """Get a sqlmodel session to interact with the database.
431
598
 
432
599
  Args:
433
- connection: SQLAlchemy connection to use when detecting changes.
434
- message: Human readable identifier describing the generated revision.
435
- write_migration_scripts: If True, write autogenerated revisions to script directory.
600
+ url: The database url.
436
601
 
437
602
  Returns:
438
- True when changes have been detected.
439
- """
440
- if not environment.ALEMBIC_CONFIG.get().exists():
441
- return False
442
-
443
- config, script_directory = cls._alembic_config()
444
- revision_context = alembic.autogenerate.api.RevisionContext(
445
- config=config,
446
- script_directory=script_directory,
447
- command_args=defaultdict(
448
- lambda: None,
449
- autogenerate=True,
450
- head="head",
451
- message=message,
452
- ),
453
- )
454
- writer = alembic.autogenerate.rewriter.Rewriter()
455
-
456
- @writer.rewrites(alembic.operations.ops.AddColumnOp)
457
- def render_add_column_with_server_default(
458
- context: MigrationContext,
459
- revision: str | None,
460
- op: Any,
461
- ):
462
- # Carry the sqlmodel default as server_default so that newly added
463
- # columns get the desired default value in existing rows.
464
- if op.column.default is not None and op.column.server_default is None:
465
- op.column.server_default = sqlalchemy.DefaultClause(
466
- sqlalchemy.sql.expression.literal(op.column.default.arg),
467
- )
468
- return op
469
-
470
- def run_autogenerate(rev: str, context: MigrationContext):
471
- revision_context.run_autogenerate(rev, context)
472
- return []
473
-
474
- with alembic.runtime.environment.EnvironmentContext(
475
- config=config,
476
- script=script_directory,
477
- fn=run_autogenerate,
478
- ) as env:
479
- env.configure(
480
- connection=connection,
481
- target_metadata=ModelRegistry.get_metadata(),
482
- render_item=cls._alembic_render_item,
483
- process_revision_directives=writer,
484
- compare_type=False,
485
- render_as_batch=True, # for sqlite compatibility
486
- )
487
- env.run_migrations()
488
- changes_detected = False
489
- if revision_context.generated_revisions:
490
- upgrade_ops = revision_context.generated_revisions[-1].upgrade_ops
491
- if upgrade_ops is not None:
492
- changes_detected = bool(upgrade_ops.ops)
493
- if changes_detected and write_migration_scripts:
494
- # Must iterate the generator to actually write the scripts.
495
- _ = tuple(revision_context.generate_scripts())
496
- return changes_detected
497
-
498
- @classmethod
499
- def _alembic_upgrade(
500
- cls,
501
- connection: sqlalchemy.engine.Connection,
502
- to_rev: str = "head",
503
- ) -> None:
504
- """Apply alembic migrations up to the given revision.
505
-
506
- Args:
507
- connection: SQLAlchemy connection to use when performing upgrade.
508
- to_rev: Revision to migrate towards.
603
+ A database session.
509
604
  """
510
- config, script_directory = cls._alembic_config()
511
-
512
- def run_upgrade(rev: str, context: MigrationContext):
513
- return script_directory._upgrade_revs(to_rev, rev)
514
-
515
- with alembic.runtime.environment.EnvironmentContext(
516
- config=config,
517
- script=script_directory,
518
- fn=run_upgrade,
519
- ) as env:
520
- env.configure(connection=connection)
521
- env.run_migrations()
605
+ return sqlmodel.Session(get_engine(url))
522
606
 
523
- @classmethod
524
- def migrate(cls, autogenerate: bool = False) -> bool | None:
525
- """Execute alembic migrations for all sqlmodel Model classes.
607
+ def asession(url: str | None = None) -> AsyncSession:
608
+ """Get an async sqlmodel session to interact with the database.
526
609
 
527
- If alembic is not installed or has not been initialized for the project,
528
- then no action is performed.
610
+ async with rx.asession() as asession:
611
+ ...
529
612
 
530
- If there are no revisions currently tracked by alembic, then
531
- an initial revision will be created based on sqlmodel metadata.
532
-
533
- If models in the app have changed in incompatible ways that alembic
534
- cannot automatically generate revisions for, the app may not be able to
535
- start up until migration scripts have been corrected by hand.
613
+ Most operations against the `asession` must be awaited.
536
614
 
537
615
  Args:
538
- autogenerate: If True, generate migration script and use it to upgrade schema
539
- (otherwise, just bring the schema to current "head" revision).
540
-
541
- Returns:
542
- True - indicating the process was successful.
543
- None - indicating the process was skipped.
544
- """
545
- if not environment.ALEMBIC_CONFIG.get().exists():
546
- return None
547
-
548
- with cls.get_db_engine().connect() as connection:
549
- cls._alembic_upgrade(connection=connection)
550
- if autogenerate:
551
- changes_detected = cls.alembic_autogenerate(connection=connection)
552
- if changes_detected:
553
- cls._alembic_upgrade(connection=connection)
554
- connection.commit()
555
- return True
556
-
557
- @classmethod
558
- def select(cls):
559
- """Select rows from the table.
616
+ url: The database url.
560
617
 
561
618
  Returns:
562
- The select statement.
619
+ An async database session.
563
620
  """
564
- return sqlmodel.select(cls)
565
-
566
-
567
- ModelRegistry.register(Model)
568
-
569
-
570
- def session(url: str | None = None) -> sqlmodel.Session:
571
- """Get a sqlmodel session to interact with the database.
572
-
573
- Args:
574
- url: The database url.
575
-
576
- Returns:
577
- A database session.
578
- """
579
- return sqlmodel.Session(get_engine(url))
580
-
581
-
582
- def asession(url: str | None = None) -> AsyncSession:
583
- """Get an async sqlmodel session to interact with the database.
584
-
585
- async with rx.asession() as asession:
586
- ...
587
-
588
- Most operations against the `asession` must be awaited.
589
-
590
- Args:
591
- url: The database url.
592
-
593
- Returns:
594
- An async database session.
595
- """
596
- global _AsyncSessionLocal
597
- if url not in _AsyncSessionLocal:
598
- _AsyncSessionLocal[url] = sqlalchemy.ext.asyncio.async_sessionmaker(
599
- bind=get_async_engine(url),
600
- class_=AsyncSession,
601
- expire_on_commit=False,
602
- autocommit=False,
603
- autoflush=False,
604
- )
605
- return _AsyncSessionLocal[url]()
606
-
607
-
608
- def sqla_session(url: str | None = None) -> sqlalchemy.orm.Session:
609
- """Get a bare sqlalchemy session to interact with the database.
610
-
611
- Args:
612
- url: The database url.
621
+ global _AsyncSessionLocal
622
+ if url not in _AsyncSessionLocal:
623
+ _AsyncSessionLocal[url] = sqlalchemy.ext.asyncio.async_sessionmaker(
624
+ bind=get_async_engine(url),
625
+ class_=AsyncSession,
626
+ expire_on_commit=False,
627
+ autocommit=False,
628
+ autoflush=False,
629
+ )
630
+ return _AsyncSessionLocal[url]()
613
631
 
614
- Returns:
615
- A database session.
616
- """
617
- return sqlalchemy.orm.Session(get_engine(url))
632
+ else:
633
+ get_db_status = _print_db_not_available
634
+ session = _print_db_not_available
635
+ asession = _print_db_not_available
636
+ Model = _ClassThatErrorsOnInit # pyright: ignore [reportAssignmentType]