Skip to content

Commit

Permalink
Changes:
Browse files Browse the repository at this point in the history
- fix url template, should only use netloc and username
- allow bytes for identifier hasher
  • Loading branch information
devkral committed Dec 14, 2024
1 parent 2750b19 commit 2e506a1
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 14 deletions.
10 changes: 4 additions & 6 deletions edgy/cli/templates/url/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
from rich.console import Console

import edgy
from edgy.core.connection import Database, Registry
from edgy.core.connection import Database, DatabaseURL, Registry

if TYPE_CHECKING:
import sqlalchemy
Expand Down Expand Up @@ -81,7 +81,7 @@ def run_migrations_offline() -> Any:
registry = edgy.get_migration_prepared_registry()
for name, db, metadata in iter_databases(registry):
# db is maybe overwritten, so use the original url
orig_url = str(registry.database.url) if name is None else str(registry.extra[name].url)
orig_url = registry.database.url if name is None else registry.extra[name].url
context.configure(
url=str(db.url),
target_metadata=metadata,
Expand All @@ -93,7 +93,7 @@ def run_migrations_offline() -> Any:


def do_run_migrations(
connection: Any, url: str, orig_url: str, name: str, metadata: "sqlalchemy.Metadata"
connection: Any, url: str, orig_url: DatabaseURL, name: str, metadata: "sqlalchemy.Metadata"
) -> Any:
# this callback is used to prevent an auto-migration from being generated
# when there are no changes to the schema
Expand Down Expand Up @@ -137,9 +137,7 @@ async def run_migrations_online() -> Any:
async with registry:
for name, db, metadata in iter_databases(registry):
# db is maybe overwritten, so use the original url
orig_url = (
str(registry.database.url) if name is None else str(registry.extra[name].url)
)
orig_url = registry.database.url if name is None else registry.extra[name].url
async with db as database:
await database.run_sync(do_run_migrations, str(db.url), orig_url, name, metadata)

Expand Down
20 changes: 14 additions & 6 deletions edgy/cli/templates/url/script.py.mako
Original file line number Diff line number Diff line change
Expand Up @@ -5,25 +5,32 @@ Revises: ${down_revision | comma,n}
Create Date: ${create_date}

"""
from typing import TYPE_CHECKING, Optional

from alembic import op
import sqlalchemy as sa
from edgy.utils.hashing import hash_to_identifier
${imports if imports else ""}

if TYPE_CHECKING:
from edgy.core.connection import DatabaseURL

# revision identifiers, used by Alembic.
revision = ${repr(up_revision)}
down_revision = ${repr(down_revision)}
branch_labels = ${repr(branch_labels)}
depends_on = ${repr(depends_on)}

def upgrade(url: str = "") -> None:
fn = globals().get(f"upgrade_{hash_to_identifier(url)}")
def upgrade(url: Optional["DatabaseURL"] = None) -> None:
urlstring = "" if url is None else f"{url.username}:{url.netloc}"
fn = globals().get(f"upgrade_{hash_to_identifier(urlstring)}")
if fn is not None:
fn()


def downgrade(url: str = "") -> None:
fn = globals().get(f"downgrade_{hash_to_identifier(url)}")
def downgrade(url: Optional["DatabaseURL"] = None) -> None:
urlstring = "" if url is None else f"{url.username}:{url.netloc}"
fn = globals().get(f"downgrade_{hash_to_identifier(urlstring)}")
if fn is not None:
fn()

Expand All @@ -35,9 +42,10 @@ def downgrade(url: str = "") -> None:
def url_for_name(name):
if name:
return str(monkay.instance.registry.extra[name].url)
url = monkay.instance.registry.extra[name].url
else:
return str(monkay.instance.registry.database.url)
url = monkay.instance.registry.database.url
return f"{url.username}:{url.netloc}"
%>

## generate an "upgrade_<xyz>() / downgrade_<xyz>()" function
Expand Down
7 changes: 5 additions & 2 deletions edgy/utils/hashing.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from base64 import b32encode
from hashlib import blake2b
from typing import Union


def hash_to_identifier(key: str) -> str:
def hash_to_identifier(key: Union[str, bytes]) -> str:
"""
A generic hasher for keys, which output stays a valid name for python
and other languages.
Expand All @@ -12,5 +13,7 @@ def hash_to_identifier(key: str) -> str:
See edgy/cli/templates/default/script.py or edgy/core/db/querysets/base.py for the usage.
"""
if isinstance(key, str):
key = key.encode()
# prefix with _ for preventing a name starting with a number
return f"_{b32encode(blake2b(key.encode(), digest_size=16).digest()).decode().rstrip('=')}"
return f"_{b32encode(blake2b(key, digest_size=16).digest()).decode().rstrip('=')}"

0 comments on commit 2e506a1

Please sign in to comment.