From c778031f0275cf611acca1fab4a37cda4efdf6f5 Mon Sep 17 00:00:00 2001
From: Ronen Lubin <63970571+ronenlu@users.noreply.github.com>
Date: Wed, 24 Jan 2024 11:52:12 +0200
Subject: [PATCH] support MariaDB dialect (#8)

---
 .../commands/atlas-provider-django.py         | 48 ++++++++++++-------
 1 file changed, 32 insertions(+), 16 deletions(-)

diff --git a/atlas_provider_django/management/commands/atlas-provider-django.py b/atlas_provider_django/management/commands/atlas-provider-django.py
index 7e48cf0..2a121ae 100644
--- a/atlas_provider_django/management/commands/atlas-provider-django.py
+++ b/atlas_provider_django/management/commands/atlas-provider-django.py
@@ -21,6 +21,7 @@
 
 class Dialect(str, Enum):
     mysql = "mysql"
+    mariadb = "mariadb"
     sqlite = "sqlite"
     postgresql = "postgresql"
 
@@ -74,26 +75,41 @@ def has_native_uuid_field(self):
         return False
 
 
+class MockMariaDBDatabaseFeatures(MySQLDatabaseFeatures):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def has_native_uuid_field(self):
+        return True
+
+
 # Returns the database connection wrapper for the given dialect.
 # Mocks some methods in order to get the sql statements without db connection.
 def get_connection_by_dialect(dialect):
     conn = None
-    if dialect == Dialect.sqlite:
-        conn = Sqlite3DatabaseWrapper({
-            "ENGINE": "django.db.backends.sqlite3",
-        }, "sqlite3")
-        conn.SchemaEditorClass = MockSqliteSchemaEditor
-    elif dialect == Dialect.postgresql:
-        conn = PGDatabaseWrapper({
-            "ENGINE": "django.db.backends.postgresql",
-        }, "postgresql")
-        conn.SchemaEditorClass = MockPGDatabaseSchemaEditor
-    elif dialect == Dialect.mysql:
-        conn = MySQLDatabaseWrapper({
-            "ENGINE": "django.db.backends.mysql",
-        }, "mysql")
-        conn.SchemaEditorClass = MockMySQLDatabaseSchemaEditor
-        conn.features = MockMySQLDatabaseFeatures
+    match dialect:
+        case Dialect.sqlite:
+            conn = Sqlite3DatabaseWrapper({
+                "ENGINE": "django.db.backends.sqlite3",
+            }, "sqlite3")
+            conn.SchemaEditorClass = MockSqliteSchemaEditor
+        case Dialect.postgresql:
+            conn = PGDatabaseWrapper({
+                "ENGINE": "django.db.backends.postgresql",
+            }, "postgresql")
+            conn.SchemaEditorClass = MockPGDatabaseSchemaEditor
+        case Dialect.mysql:
+            conn = MySQLDatabaseWrapper({
+                "ENGINE": "django.db.backends.mysql",
+            }, "mysql")
+            conn.SchemaEditorClass = MockMySQLDatabaseSchemaEditor
+            conn.features = MockMySQLDatabaseFeatures(conn)
+        case Dialect.mariadb:
+            conn = MySQLDatabaseWrapper({
+                "ENGINE": "django.db.backends.mysql",
+            }, "mysql")
+            conn.SchemaEditorClass = MockMySQLDatabaseSchemaEditor
+            conn.features = MockMariaDBDatabaseFeatures(conn)
     return conn