Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
open-webui
Commits
981866eb
Unverified
Commit
981866eb
authored
Jun 17, 2024
by
perf3ct
Browse files
use autoconnect and stop using the mixin
parent
10fa887e
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
19 additions
and
19 deletions
+19
-19
backend/apps/webui/internal/wrappers.py
backend/apps/webui/internal/wrappers.py
+19
-19
No files found.
backend/apps/webui/internal/wrappers.py
View file @
981866eb
from
contextvars
import
ContextVar
from
peewee
import
*
from
playhouse.db_url
import
connect
from
playhouse.pool
import
PooledPostgresqlDatabase
from
playhouse.
shortcuts
import
ReconnectMixin
from
playhouse.pool
import
PooledPostgresql
Ext
Database
from
playhouse.
pool
import
PooledSqliteDatabase
db_state_default
=
{
"closed"
:
None
,
"conn"
:
None
,
"ctx"
:
None
,
"transactions"
:
None
}
db_state
=
ContextVar
(
"db_state"
,
default
=
db_state_default
.
copy
())
...
...
@@ -19,26 +19,26 @@ class PeeweeConnectionState(object):
value
=
self
.
_state
.
get
()[
name
]
return
value
class
ReconnectingPostgresqlDatabase
(
ReconnectMixin
,
PostgresqlDatabase
):
pass
class
ReconnectingPooledPostgresqlDatabase
(
ReconnectMixin
,
PooledPostgresqlDatabase
):
pass
class
ReconnectingSqliteDatabase
(
ReconnectMixin
,
SqliteDatabase
):
pass
def
register_connection
(
db_url
):
# Connect using the playhouse.db_url module, which supports multiple
# database types, then wrap the connection in a ReconnectMixin to handle dropped connections
db
=
connect
(
db_url
)
if
isinstance
(
db
,
PostgresqlDatabase
):
db
=
ReconnectingPostgresqlDatabase
(
db
.
database
,
**
db
.
connect_params
)
elif
isinstance
(
db
,
PooledPostgresqlDatabase
):
db
=
ReconnectingPooledPostgresqlDatabase
(
db
.
database
,
**
db
.
connect_params
)
db
=
PooledPostgresqlExtDatabase
(
db
.
database
,
max_connections
=
8
,
stale_timeout
=
300
,
timeout
=
None
,
autoconnect
=
True
,
**
db
.
connect_params
)
elif
isinstance
(
db
,
SqliteDatabase
):
db
=
ReconnectingSqliteDatabase
(
db
.
database
,
**
db
.
connect_params
)
db
=
PooledSqliteDatabase
(
db
.
database
,
max_connections
=
8
,
stale_timeout
=
300
,
timeout
=
None
,
autoconnect
=
True
,
**
db
.
connect_params
)
else
:
raise
ValueError
(
'Unsupported database connection'
)
return
db
return
db
\ No newline at end of file
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment