From 2e506a1cd0c3bf86d9aa6eea053dde43df04f5ce Mon Sep 17 00:00:00 2001
From: alex <devkral@web.de>
Date: Sat, 14 Dec 2024 07:45:06 +0100
Subject: [PATCH] Changes:

- fix url template, should only use netloc and username
- allow bytes for identifier hasher
---
 edgy/cli/templates/url/env.py         | 10 ++++------
 edgy/cli/templates/url/script.py.mako | 20 ++++++++++++++------
 edgy/utils/hashing.py                 |  7 +++++--
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/edgy/cli/templates/url/env.py b/edgy/cli/templates/url/env.py
index adc62f45..9b4d866f 100644
--- a/edgy/cli/templates/url/env.py
+++ b/edgy/cli/templates/url/env.py
@@ -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
@@ -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,
@@ -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
@@ -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)
 
diff --git a/edgy/cli/templates/url/script.py.mako b/edgy/cli/templates/url/script.py.mako
index 5ccd897b..e08e2a78 100644
--- a/edgy/cli/templates/url/script.py.mako
+++ b/edgy/cli/templates/url/script.py.mako
@@ -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()
 
@@ -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
diff --git a/edgy/utils/hashing.py b/edgy/utils/hashing.py
index 59b50b58..848411a3 100644
--- a/edgy/utils/hashing.py
+++ b/edgy/utils/hashing.py
@@ -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.
@@ -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('=')}"