Merge branch 'origindev' into sparkwallet

This commit is contained in:
Arc
2026-01-29 04:57:29 +00:00
23 changed files with 1745 additions and 1326 deletions
+26 -12
View File
@@ -3,6 +3,7 @@ import io
from uuid import uuid4
from fastapi import UploadFile
from loguru import logger
from PIL import Image
from lnbits.core.crud.assets import create_asset, get_user_assets_count
@@ -29,17 +30,7 @@ async def create_user_asset(user_id: str, file: UploadFile, is_public: bool) ->
f"File limit of {settings.lnbits_max_asset_size_mb}MB exceeded."
)
image = Image.open(io.BytesIO(contents))
thumbnail_width = min(256, settings.lnbits_asset_thumbnail_width)
thumbnail_height = min(256, settings.lnbits_asset_thumbnail_height)
image.thumbnail((thumbnail_width, thumbnail_height))
# Save thumbnail to an in-memory buffer
thumb_buffer = io.BytesIO()
thumbnail_format = settings.lnbits_asset_thumbnail_format or "PNG"
image.save(thumb_buffer, format=thumbnail_format)
thumb_buffer.seek(0)
thumb_buffer = thumbnail_from_bytes(contents)
asset = Asset(
id=uuid4().hex,
@@ -48,9 +39,32 @@ async def create_user_asset(user_id: str, file: UploadFile, is_public: bool) ->
is_public=is_public,
name=file.filename or "unnamed",
size_bytes=len(contents),
thumbnail_base64=base64.b64encode(thumb_buffer.getvalue()).decode("utf-8"),
thumbnail_base64=(
base64.b64encode(thumb_buffer.getvalue()).decode("utf-8")
if thumb_buffer
else None
),
data=contents,
)
await create_asset(asset)
return asset
def thumbnail_from_bytes(contents: bytes) -> io.BytesIO | None:
try:
image = Image.open(io.BytesIO(contents))
thumbnail_width = min(256, settings.lnbits_asset_thumbnail_width)
thumbnail_height = min(256, settings.lnbits_asset_thumbnail_height)
image.thumbnail((thumbnail_width, thumbnail_height))
# Save thumbnail to an in-memory buffer
thumb_buffer = io.BytesIO()
thumbnail_format = settings.lnbits_asset_thumbnail_format or "PNG"
image.save(thumb_buffer, format=thumbnail_format)
thumb_buffer.seek(0)
return thumb_buffer
except Exception as exc:
logger.warning(f"Failed to create thumbnail: {exc}")
return None
+4 -4
View File
@@ -57,11 +57,11 @@ async def api_get_asset(
@asset_router.get(
"/{asset_id}/binary",
name="Get user asset binary",
summary="Get user asset binary data by ID",
"/{asset_id}/data",
name="Get user asset data",
summary="Get user asset data data by ID",
)
async def api_get_asset_binary(
async def api_get_asset_data(
asset_id: str,
user_id: str | None = Depends(optional_user_id),
) -> Response:
+60 -1
View File
@@ -4,9 +4,10 @@ import json
from collections.abc import Callable
from http import HTTPStatus
from time import time
from typing import Annotated
from uuid import uuid4
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi import APIRouter, Cookie, Depends, HTTPException, Request
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi_sso.sso.base import OpenID, SSOBase
from loguru import logger
@@ -29,6 +30,7 @@ from lnbits.core.services.users import update_user_account
from lnbits.decorators import (
access_token_payload,
check_account_exists,
check_admin,
check_user_exists,
)
from lnbits.helpers import (
@@ -120,6 +122,63 @@ async def login_usr(data: LoginUsr) -> JSONResponse:
return _auth_success_response(account.username, account.id, account.email)
@auth_router.post("/impersonate", description="Login via the User ID of another user")
async def impersonate_user(
data: LoginUsr,
user: User = Depends(check_admin),
cookie_access_token: Annotated[str | None, Cookie()] = None,
) -> JSONResponse:
if not cookie_access_token:
raise HTTPException(
HTTPStatus.UNAUTHORIZED, "Only cookie based impersonation is allowed."
)
if data.usr == user.id:
raise HTTPException(HTTPStatus.FORBIDDEN, "You cannot impersonate yourself.")
if settings.is_admin_user(data.usr):
# this check includes the superuser
raise HTTPException(
HTTPStatus.FORBIDDEN, "You cannot impersonate another admin user."
)
account = await get_account(data.usr)
if not account:
raise HTTPException(HTTPStatus.UNAUTHORIZED, "User ID does not exist.")
response = _auth_success_response(account.username, account.id, account.email)
max_age = settings.auth_token_expire_minutes * 60
response.set_cookie(
"admin_access_token", cookie_access_token, httponly=True, max_age=max_age
)
response.set_cookie("is_lnbits_user_impersonated", "true", max_age=max_age)
return response
@auth_router.delete(
"/impersonate", description="Stop impersonation and go back to admin"
)
async def stop_impersonate_user(
user: User = Depends(check_user_exists),
admin_access_token: Annotated[str | None, Cookie()] = None,
) -> JSONResponse:
if not admin_access_token:
raise HTTPException(
HTTPStatus.UNAUTHORIZED,
"No admin access token found to stop impersonation.",
)
response = JSONResponse(
{"access_token": admin_access_token, "token_type": "bearer"}
)
max_age = settings.auth_token_expire_minutes * 60
response.set_cookie(
"cookie_access_token", admin_access_token, httponly=True, max_age=max_age
)
response.delete_cookie("admin_access_token")
response.delete_cookie("is_access_token_expired")
response.delete_cookie("is_lnbits_user_impersonated")
return response
@auth_router.get("/acl")
async def api_get_user_acls(
request: Request,
+4
View File
@@ -306,6 +306,10 @@ class AssetSettings(LNbitsSettings):
"heic",
"heif",
"heics",
"text/plain",
"text/json" "text/xml",
"application/json",
"application/pdf",
]
)
lnbits_asset_thumbnail_width: int = Field(default=128, ge=0)
File diff suppressed because one or more lines are too long
+10 -10
View File
File diff suppressed because one or more lines are too long
+2
View File
@@ -695,6 +695,8 @@ window.localisation.en = {
allow_creation_user: 'Allow creation of new users',
allow_creation_user_desc: 'Allow creation of new users on the index page',
new_user_not_allowed: 'Registration is disabled.',
start_user_impersonation: 'Impersonate this user',
stop_user_impersonation: 'Stop User Impersonation',
components: 'Components',
long_running_endpoints: 'Top 5 Long Running Endpoints',
http_request_methods: 'HTTP Request Methods',
+13
View File
@@ -123,6 +123,19 @@ window._lnbitsApi = {
url: '/api/v1/auth/logout'
})
},
impersonateUser(usr) {
return axios({
method: 'POST',
url: '/api/v1/auth/impersonate',
data: {usr}
})
},
stopImpersonation() {
return axios({
method: 'DELETE',
url: '/api/v1/auth/impersonate'
})
},
getAuthenticatedUser() {
return this.request('get', '/api/v1/auth')
},
@@ -16,6 +16,7 @@ window.app.component('lnbits-admin-assets-config', {
this.newAllowedAssetMimeType
)
this.newAllowedAssetMimeType = ''
this.formData.touch = null
}
},
removeAllowedAssetMimeType(type) {
@@ -23,18 +24,21 @@ window.app.component('lnbits-admin-assets-config', {
if (index !== -1) {
this.formData.lnbits_assets_allowed_mime_types.splice(index, 1)
}
this.formData.touch = null
},
addNewNoLimitUser() {
if (this.newNoLimitUser) {
this.removeNoLimitUser(this.newNoLimitUser)
this.formData.lnbits_assets_no_limit_users.push(this.newNoLimitUser)
this.newNoLimitUser = ''
this.formData.touch = null
}
},
removeNoLimitUser(user) {
if (user) {
this.formData.lnbits_assets_no_limit_users =
this.formData.lnbits_assets_no_limit_users.filter(u => u !== user)
this.formData.touch = null
}
}
}
@@ -60,5 +60,17 @@ window.app.component('lnbits-header', {
return 'User'
}
}
},
methods: {
async stopImpersonation() {
try {
await LNbits.api.stopImpersonation()
LNbits.utils.restoreLocalStorage('impersonation')
window.location = '/users'
} catch (e) {
console.warn(e)
LNbits.utils.notifyApiError(e)
}
}
}
})
+1
View File
@@ -15,6 +15,7 @@ window.g = Vue.reactive({
wallet: null,
isPublicPage: true,
isUserAuthorized: !!Quasar.Cookies.get('is_lnbits_user_authorized'),
isUserImpersonated: !!Quasar.Cookies.get('is_lnbits_user_impersonated'),
offline: !navigator.onLine,
hasCamera: false,
visibleDrawer: false,
+1 -1
View File
@@ -609,7 +609,7 @@ window.PageAccount = {
}
},
copyAssetLinkToClipboard(asset) {
const assetUrl = `${window.location.origin}/api/v1/assets/${asset.id}/binary`
const assetUrl = `${window.location.origin}/api/v1/assets/${asset.id}/data`
this.utils.copyText(assetUrl)
},
addUserLabel() {
+14
View File
@@ -438,6 +438,20 @@ window.PageUsers = {
this.activeUser.show = false
}
},
async impersonateUser(user_id) {
try {
await LNbits.api.impersonateUser(user_id)
LNbits.utils.backupLocalStorage('impersonation', true)
this.$q.localStorage.setItem('lnbits.disclaimerShown', true)
window.location = '/wallet'
} catch (error) {
console.warn(error)
Quasar.Notify.create({
type: 'warning',
message: 'Failed to impersonate user!'
})
}
},
async showWalletPayments(walletId) {
this.activeUser.show = false
await this.fetchWallets(this.users[0].id)
+23
View File
@@ -67,6 +67,29 @@ window._lnbitsUtils = {
}
})
},
backupLocalStorage(backupKey, cleanup = false) {
const lnbitsEntries =
Object.entries(Quasar.LocalStorage.getAll()).filter(
([k, v]) => k.startsWith('lnbits.') && k !== `lnbits.${backupKey}`
) || []
Quasar.LocalStorage.setItem(`lnbits.${backupKey}`, lnbitsEntries)
if (cleanup) {
lnbitsEntries.forEach(([k, v]) => Quasar.LocalStorage.remove(k))
}
},
restoreLocalStorage(backupKey) {
Object.entries(Quasar.LocalStorage.getAll())
.filter(
([k, v]) => k.startsWith('lnbits.') && k !== `lnbits.${backupKey}`
)
.forEach(([k, v]) => Quasar.LocalStorage.remove(k))
const lnbitsEntries =
Quasar.LocalStorage.getItem(`lnbits.${backupKey}`) || []
lnbitsEntries.forEach(([k, v]) => Quasar.LocalStorage.setItem(k, v))
Quasar.LocalStorage.remove(`lnbits.${backupKey}`)
},
async digestMessage(message) {
const msgUint8 = new TextEncoder().encode(message)
const hashBuffer = await crypto.subtle.digest('SHA-256', msgUint8)
+14 -1
View File
@@ -72,7 +72,6 @@
</q-badge>
<lnbits-language-dropdown></lnbits-language-dropdown>
<q-btn-dropdown v-if="g.user" flat rounded size="sm" class="q-pl-sm">
<template v-slot:label>
<q-avatar
@@ -136,6 +135,20 @@
</q-item>
</q-list>
</q-btn-dropdown>
<q-btn
v-if="g.isUserImpersonated"
@click="stopImpersonation"
rounded
size="sm"
class="q-pl-sm"
color="negative"
icon="face_retouching_off"
label="Stop"
>
<q-tooltip
><span v-text="$t('stop_user_impersonation')"></span
></q-tooltip>
</q-btn>
</q-toolbar>
</q-header>
</template>
+1 -1
View File
@@ -941,7 +941,7 @@
v-if="props.row.thumbnail_base64"
target="_blank"
style="color: inherit"
:href="`/api/v1/assets/${props.row.id}/binary`"
:href="`/api/v1/assets/${props.row.id}/data`"
>
<q-img
:src="
+14 -1
View File
@@ -656,7 +656,20 @@
</q-btn>
<span v-text="shortify(props.row.id)"></span>
</q-td>
<q-td v-text="props.row.username"></q-td>
<q-td>
<q-btn
icon="face"
size="sm"
flat
class="cursor-pointer q-mr-xs"
@click="impersonateUser(props.row.id)"
>
<q-tooltip
><span v-text="$t('start_user_impersonation')"></span
></q-tooltip>
</q-btn>
<span v-text="props.row.username"></span>
</q-td>
<q-td v-text="props.row.email"></q-td>
Generated
+823 -772
View File
File diff suppressed because it is too large Load Diff
+37 -58
View File
@@ -8,31 +8,31 @@ urls = { Homepage = "https://lnbits.com", Repository = "https://github.com/lnbit
readme = "README.md"
dependencies = [
"bech32==1.2.0",
"click==8.2.1",
"click==8.3.1",
"ecdsa==0.19.1",
"fastapi==0.116.1",
"starlette==0.47.1",
"httpx==0.27.0",
"httpx==0.27.2",
"jinja2==3.1.6",
"lnurl==0.8.3",
"pydantic==1.10.22",
"pydantic==1.10.26",
"pyqrcode==1.2.1",
"shortuuid==1.0.13",
"sse-starlette==2.3.6",
"typing-extensions==4.14.0",
"uvicorn==0.34.3",
"typing-extensions==4.15.0",
"uvicorn==0.40.0",
"sqlalchemy==1.4.54",
"aiosqlite==0.21.0",
"asyncpg==0.30.0",
"uvloop==0.21.0",
"aiosqlite==0.22.1",
"asyncpg==0.31.0",
"uvloop==0.22.1",
"websockets==15.0.1",
"loguru==0.7.3",
"grpcio==1.69.0",
"protobuf==5.29.5",
"pyln-client==25.5",
"pywebpush==2.0.3",
"grpcio==1.76.0",
"protobuf==6.33.2",
"pyln-client==25.12",
"pywebpush==2.2.0",
"slowapi==0.1.9",
"websocket-client==1.8.0",
"websocket-client==1.9.0",
"pycryptodomex==3.23.0",
"packaging==25.0",
"bolt11==2.1.1",
@@ -42,14 +42,14 @@ dependencies = [
# needed for boltz, lnurldevice, watchonly extensions
"embit==0.8.0",
# needed for scheduler extension
"python-crontab==3.2.0",
"pynostr==0.6.2",
"python-multipart==0.0.20",
"python-crontab==3.3.0",
"pynostr==0.7.0",
"python-multipart==0.0.21",
"filetype==1.2.0",
"nostr-sdk==0.42.1",
"bcrypt==4.3.0",
"nostr-sdk==0.44.0",
"bcrypt==5.0.0",
"jsonpath-ng==1.7.0",
"pillow>=12.0.0",
"pillow>=12.1.0",
"python-dotenv>=1.2.1",
"greenlet (>=3.3.0,<4.0.0)",
]
@@ -60,30 +60,30 @@ lnbits-cli = "lnbits.commands:main"
[project.optional-dependencies]
breez = ["breez-sdk==0.8.0", "breez-sdk-liquid==0.11.11"]
liquid = ["wallycore==1.4.0"]
migration = ["psycopg2-binary==2.9.10"]
liquid = ["wallycore==1.5.1"]
migration = ["psycopg2-binary==2.9.11"]
[dependency-groups]
dev = [
"black>=25.1.0,<26.0.0",
"black>=25.12.0,<26.0.0",
"mypy==1.17.1",
"types-protobuf>=6.30.2.20250516,<7.0.0",
"pre-commit>=4.2.0,<5.0.0",
"openapi-spec-validator>=0.7.1,<1.0.0",
"ruff>=0.12.0,<1.0.0",
"types-passlib>=1.7.7.20240327,<2.0.0",
"openai>=1.39.0,<2.0.0",
"json5>=0.12.0,<1.0.0",
"types-protobuf>=6.32.1.20251210,<7.0.0",
"pre-commit>=4.5.1,<5.0.0",
"openapi-spec-validator>=0.7.2,<1.0.0",
"ruff>=0.14.10,<1.0.0",
"types-passlib>=1.7.7.20250602,<2.0.0",
"openai>=2.14.0",
"json5>=0.13.0,<1.0.0",
"asgi-lifespan>=2.1.0,<3.0.0",
"anyio>=4.7.0,<5.0.0",
"pytest>=8.3.4,<9.0.0",
"pytest-cov>=6.0.0,<7.0.0",
"anyio>=4.12.1",
"pytest>=9.0.2",
"pytest-cov>=7.0.0",
"pytest-md>=0.2.0,<0.3.0",
"pytest-httpserver>=1.1.0,<2.0.0",
"pytest-mock>=3.14.0,<4.0.0",
"types-mock>=5.1.0.20240425,<6.0.0",
"mock>=5.1.0,<6.0.0",
"grpcio-tools>=1.69.0,<2.0.0"
"pytest-httpserver>=1.1.3,<2.0.0",
"pytest-mock>=3.15.1,<4.0.0",
"types-mock>=5.2.0.20250924,<6.0.0",
"mock>=5.2.0,<6.0.0",
"grpcio-tools>=1.76.0,<2.0.0"
]
[tool.poetry]
@@ -92,27 +92,6 @@ packages = [
{include = "lnbits/py.typed"},
]
[tool.poetry.group.dev.dependencies]
black = "^25.1.0"
mypy = "^1.17.1"
types-protobuf = "^6.30.2.20250516"
pre-commit = "^4.2.0"
openapi-spec-validator = "^0.7.1"
ruff = "^0.12.0"
types-passlib = "^1.7.7.20240327"
openai = "^1.39.0"
json5 = "^0.12.0"
asgi-lifespan = "^2.1.0"
anyio = "^4.7.0"
pytest = "^8.3.4"
pytest-cov = "^6.0.0"
pytest-md = "^0.2.0"
pytest-httpserver = "^1.1.0"
pytest-mock = "^3.14.0"
types-mock = "^5.1.0.20240425"
mock = "^5.1.0"
grpcio-tools = "^1.69.0"
[tool.pyright]
include = [
"lnbits",
+190
View File
@@ -2047,3 +2047,193 @@ async def test_api_update_user_labels(http_client: AsyncClient):
"""string does not match regex "([A-Za-z0-9 ._-]{1,100}$)"""
in data["detail"][0]["msg"]
)
@pytest.mark.anyio
async def test_impersonate_user_success(http_client: AsyncClient, admin_user: User):
tiny_id = shortuuid.uuid()[:8]
user_id = uuid4().hex
account = Account(
id=user_id,
username=f"u_{tiny_id}",
email=f"u_{tiny_id}@lnbits.com",
)
account.hash_password("secret1234")
await create_user_account(account)
# Login as admin to get access token
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
assert response.status_code == 200
admin_token = response.json()["access_token"]
# Impersonate the user
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": user_id},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 200
data = response.json()
assert "access_token" in data
assert data["token_type"] == "bearer"
# Check impersonation cookies
assert "cookie_access_token" in response.cookies
assert "admin_access_token" in response.cookies
assert response.cookies.get("is_lnbits_user_impersonated") == "true"
response = await http_client.delete("/api/v1/auth/impersonate")
assert "cookie_access_token" in response.cookies
assert response.cookies.get("cookie_access_token") == admin_token
assert "admin_access_token" not in response.cookies
assert "is_lnbits_user_impersonated" not in response.cookies
@pytest.mark.anyio
async def test_impersonate_user_no_cookie(http_client: AsyncClient, admin_user: User):
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
user_id = uuid4().hex
http_client.cookies.clear()
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": user_id},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 401
assert response.json()["detail"] == "Only cookie based impersonation is allowed."
@pytest.mark.anyio
async def test_impersonate_user_self(http_client: AsyncClient, admin_user: User):
# Admin tries to impersonate themselves
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": admin_user.id},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 403
assert response.json()["detail"] == "You cannot impersonate yourself."
@pytest.mark.anyio
async def test_impersonate_user_admin_target(
http_client: AsyncClient, admin_user: User, settings: Settings
):
other_admin_id = uuid4().hex
tiny_id = shortuuid.uuid()[:8]
account = Account(
id=other_admin_id, username=f"u_{tiny_id}", email=f"u_{tiny_id}@lnbits.com"
)
account.hash_password("secret1234")
await create_user_account(account)
settings.lnbits_admin_users.append(other_admin_id)
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": other_admin_id},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 403
assert response.json()["detail"] == "You cannot impersonate another admin user."
@pytest.mark.anyio
async def test_impersonate_user_nonexistent(
http_client: AsyncClient, admin_user: User, settings: Settings
):
http_client.cookies.clear()
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
fake_id = "deadbeef"
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": fake_id},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 401
assert response.json()["detail"] == "User ID does not exist."
@pytest.mark.anyio
async def test_impersonate_user_invalid_data(
http_client: AsyncClient, admin_user: User
):
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
response = await http_client.post(
"/api/v1/auth/impersonate",
json={"usr": None},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 400
assert "type_error.none.not_allowed" in str(response.json()["detail"])
@pytest.mark.anyio
async def test_impersonate_user_missing_usr_field(
http_client: AsyncClient, admin_user: User
):
response = await http_client.post(
"/api/v1/auth", json={"username": admin_user.username, "password": "secret1234"}
)
admin_token = response.json()["access_token"]
response = await http_client.post(
"/api/v1/auth/impersonate",
json={},
headers={"Authorization": f"Bearer {admin_token}"},
)
assert response.status_code == 400
assert "value_error.missing" in str(response.json()["detail"])
@pytest.mark.anyio
async def test_impersonate_user_by_non_admin(http_client: AsyncClient, user_alan: User):
response = await http_client.post(
"/api/v1/auth", json={"username": user_alan.username, "password": "secret1234"}
)
alan_token = response.json()["access_token"]
response = await http_client.post(
"/api/v1/auth/impersonate",
json={},
headers={"Authorization": f"Bearer {alan_token}"},
)
assert response.status_code == 403
assert response.json()["detail"] == "User not authorized. No admin privileges."
@pytest.mark.anyio
async def test_stop_impersonate_user_by_non_admin(
http_client: AsyncClient, user_alan: User
):
response = await http_client.post(
"/api/v1/auth", json={"username": user_alan.username, "password": "secret1234"}
)
response = await http_client.delete("/api/v1/auth/impersonate")
assert response.status_code == 401
assert (
response.json()["detail"]
== "No admin access token found to stop impersonation."
)
+5 -5
View File
@@ -14,9 +14,9 @@ from lnbits.core.crud import (
delete_account,
get_account_by_username,
get_payment,
get_user,
update_payment,
)
from lnbits.core.crud.users import get_user_from_account
from lnbits.core.models import Account, CreateInvoice, PaymentState, User
from lnbits.core.models.users import UpdateSuperuserPassword
from lnbits.core.services import create_user_account, update_wallet_balance
@@ -114,10 +114,10 @@ async def user_alan():
@pytest.fixture(scope="session")
async def admin_user():
username = "admin"
account = await get_account_by_username(username)
if account:
return await get_user_from_account(account)
username = "admin_" + uuid4().hex[:8]
user = await get_user(ADMIN_USER_ID)
if user:
return user
account = Account(
id=ADMIN_USER_ID,
+60 -20
View File
@@ -100,8 +100,8 @@ def insert_to_pg(query, data):
logger.error(exc)
logger.error(f"Failed to insert {d}")
else:
logger.error("query:", query)
logger.error("data:", d)
logger.error("query: " + query)
logger.error("data: " + str(d))
raise ValueError(f"Failed to insert {d}") from exc
connection.commit()
@@ -125,6 +125,7 @@ def migrate_ext(file: str):
migrate_db(file, schema)
logger.info(f"✅ Migrated ext: {schema}")
except Exception as exc:
logger.error(exc)
logger.error(f"🛑 Failed to migrate extension {schema}: {exc}")
@@ -134,8 +135,8 @@ def migrate_db(file: str, schema: str, exclude_tables: list[str] | None = None):
exclude_tables = []
assert os.path.isfile(file), f"{file} does not exist!"
cursor = get_sqlite_cursor(file)
tables = cursor.execute(
sqlite_cursor = get_sqlite_cursor(file)
tables = sqlite_cursor.execute(
"""
SELECT name FROM sqlite_master
WHERE type='table' AND name not like 'sqlite?_%' escape '?'
@@ -151,16 +152,18 @@ def migrate_db(file: str, schema: str, exclude_tables: list[str] | None = None):
if exclude_tables and table_name in exclude_tables:
continue
columns = cursor.execute(f"PRAGMA table_info({table_name})").fetchall()
columns = build_table_columns(file, schema, table_name)
q = build_insert_query(schema, table_name, columns)
data = cursor.execute(f"SELECT * FROM {table_name};").fetchall()
data = sqlite_cursor.execute(f"SELECT * FROM {table_name};").fetchall()
if len(data) == 0:
logger.warning(f"🛑 You sneaky dev! Table {table_name} is empty!")
logger.warning(f"⚠️ You sneaky dev! Table {table_name} is empty!")
continue
insert_to_pg(q, data)
cursor.close()
logger.info(f"✅ Migrated table '{schema}.{table_name}' successfully")
sqlite_cursor.close()
def build_insert_query(schema, table_name, columns):
@@ -174,6 +177,30 @@ def build_insert_query(schema, table_name, columns):
"""
def build_table_columns(file: str, schema: str, table_name: str):
sqlite_cursor = get_sqlite_cursor(file)
pg_cursor = get_postgres_cursor()
sqlite_columns = sqlite_cursor.execute(
f"PRAGMA table_info({table_name})"
).fetchall()
pg_cursor.execute(
f"""
SELECT table_name, column_name, udt_name FROM information_schema.columns
WHERE table_schema = '{schema}'AND table_name = '{table_name}';"""
)
pg_columns = pg_cursor.fetchall()
columns = []
for sqlite_col in sqlite_columns:
for pg_col in pg_columns:
if sqlite_col[1].lower() == pg_col[1].lower():
columns.append((sqlite_col[0], sqlite_col[1], pg_col[2]))
break
sqlite_cursor.close()
return columns
def build_on_conflict_query_statement(schema, table_name, columns):
unique_cols = table_unique_columns(schema, table_name)
if len(unique_cols) == 0:
@@ -191,23 +218,36 @@ def build_on_conflict_query_statement(schema, table_name, columns):
def table_unique_columns(schema, table_name):
cursor = get_postgres_cursor()
query = f"""
SELECT a.attname
FROM pg_index i
JOIN pg_attribute a ON a.attrelid = i.indrelid
AND a.attnum = ANY(i.indkey)
WHERE i.indrelid = '{schema}.{table_name}'::regclass
AND i.indisunique;
SELECT
array_agg(a.attname ORDER BY a.attnum) AS columns,
i.indisprimary as is_primary,
i.indexrelid::regclass AS index_name,
COUNT(*) AS column_count,
(COUNT(*) = 1) AS is_individual
FROM pg_index i
JOIN pg_attribute a
ON a.attrelid = i.indrelid
AND a.attnum = ANY (i.indkey)
WHERE i.indrelid = '{schema}.{table_name}'::regclass
AND i.indisunique
GROUP BY i.indexrelid;
"""
cursor.execute(query)
columns = [row[0] for row in cursor.fetchall()]
rows = cursor.fetchall()
columns = [row[0] for row in rows if not row[1]] # exclude primary keys
if len(columns) == 0:
# use primary keys if no unique keys found
columns = [row[0] for row in rows if row[1]]
cursor.close()
return columns
if len(columns) == 0:
return []
return columns[0]
def to_column_type(column_type):
if column_type == "TIMESTAMP":
def to_column_type(column_type: str):
if column_type.upper() == "TIMESTAMP":
return "to_timestamp(%s)"
if column_type in ["BOOLEAN", "BOOL"]:
if column_type.upper() in ["BOOLEAN", "BOOL"]:
return "%s::boolean"
return "%s"
@@ -255,7 +295,7 @@ parser.add_argument(
args = parser.parse_args()
logger.info("Selected path: ", args.sqlite_path)
logger.info("Selected path: " + args.sqlite_path)
if os.path.isdir(args.sqlite_path):
exclude_tables = ["dbversions"]
Generated
+422 -439
View File
File diff suppressed because it is too large Load Diff