From aac223ed0159534a924ef705eff1c4736406023b Mon Sep 17 00:00:00 2001
From: Dmitry Maslennikov <mrdaimor@gmail.com>
Date: Thu, 18 Apr 2024 22:46:10 +1000
Subject: [PATCH] some irisasync support

---
 setup.cfg                    |  2 ++
 setup.py                     |  1 +
 sqlalchemy_iris/__init__.py  |  1 +
 sqlalchemy_iris/base.py      | 28 +++++++++++++++++++---------
 sqlalchemy_iris/irisasync.py | 17 +++++++++++++++++
 tests/conftest.py            |  3 +++
 6 files changed, 43 insertions(+), 9 deletions(-)
 create mode 100644 sqlalchemy_iris/irisasync.py

diff --git a/setup.cfg b/setup.cfg
index 3905741..1d89336 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -33,6 +33,8 @@ addopts= --tb native -v -r fxX --maxfail=25 -p no:warnings
 
 [db]
 default=iris://_SYSTEM:SYS@localhost:1972/USER
+iris=iris://_SYSTEM:SYS@localhost:1972/USER
+irisasync=iris+irisasync://_SYSTEM:SYS@localhost:1972/USER
 irisemb=iris+emb:///
 sqlite=sqlite:///:memory:
 
diff --git a/setup.py b/setup.py
index 49b7aa0..e43fc6d 100644
--- a/setup.py
+++ b/setup.py
@@ -8,6 +8,7 @@
         "sqlalchemy.dialects": [
             "iris = sqlalchemy_iris.iris:IRISDialect_iris",
             "iris.emb = sqlalchemy_iris.embedded:IRISDialect_emb",
+            "iris.irisasync = sqlalchemy_iris.irisasync:IRISDialect_irisasync",
         ]
     },
 )
diff --git a/sqlalchemy_iris/__init__.py b/sqlalchemy_iris/__init__.py
index 64882ab..04915b5 100644
--- a/sqlalchemy_iris/__init__.py
+++ b/sqlalchemy_iris/__init__.py
@@ -31,6 +31,7 @@
 
 _registry.register("iris.iris", "sqlalchemy_iris.iris", "IRISDialect_iris")
 _registry.register("iris.emb", "sqlalchemy_iris.embedded", "IRISDialect_emb")
+_registry.register("iris.irisasync", "sqlalchemy_iris.irisasync", "IRISDialect_irisasync")
 
 __all__ = [
     "BIGINT",
diff --git a/sqlalchemy_iris/base.py b/sqlalchemy_iris/base.py
index 44820e7..203f93d 100644
--- a/sqlalchemy_iris/base.py
+++ b/sqlalchemy_iris/base.py
@@ -884,8 +884,12 @@ def __init__(self, **kwargs):
 
     def _get_server_version_info(self, connection):
         server_version = connection.connection._connection_info._server_version
-        server_version = server_version[server_version.find("Version") + 8:].split(" ")[0].split(".")
-        return tuple([int(''.join(filter(str.isdigit, v))) for v in server_version])
+        server_version = (
+            server_version[server_version.find("Version") + 8 :]
+            .split(" ")[0]
+            .split(".")
+        )
+        return tuple([int("".join(filter(str.isdigit, v))) for v in server_version])
 
     _isolation_lookup = set(
         [
@@ -906,17 +910,23 @@ def on_connect(conn):
                 super_(conn)
 
             if self.embedded:
-                self.supports_vectors = conn.iris.cls("%SYSTEM.License").GetFeature(28) == 1
+                self.supports_vectors = (
+                    conn.iris.cls("%SYSTEM.License").GetFeature(28) == 1
+                )
             else:
                 try:
                     iris = IRISNative.createIRIS(conn)
-                    self.supports_vectors = iris.classMethodBoolean("%SYSTEM.License", "GetFeature", 28)
-                except: # noqa
+                    self.supports_vectors = iris.classMethodBoolean(
+                        "%SYSTEM.License", "GetFeature", 28
+                    )
+                except:  # noqa
                     self.supports_vectors = False
             if self.supports_vectors:
                 with conn.cursor() as cursor:
                     # Distance or similarity
-                    cursor.execute("select vector_cosine(to_vector('1'), to_vector('1'))")
+                    cursor.execute(
+                        "select vector_cosine(to_vector('1'), to_vector('1'))"
+                    )
                     self.vector_cosine_similarity = cursor.fetchone()[0] == 0
 
             self._dictionary_access = False
@@ -1511,9 +1521,9 @@ def get_multi_foreign_keys(
                 table_fkey[rfknm] = fkey = {
                     "name": rfknm,
                     "constrained_columns": [],
-                    "referred_schema": rschema
-                    if rschema != self.default_schema_name
-                    else None,
+                    "referred_schema": (
+                        rschema if rschema != self.default_schema_name else None
+                    ),
                     "referred_table": rtbl,
                     "referred_columns": [],
                     "options": {},
diff --git a/sqlalchemy_iris/irisasync.py b/sqlalchemy_iris/irisasync.py
new file mode 100644
index 0000000..0ded9e4
--- /dev/null
+++ b/sqlalchemy_iris/irisasync.py
@@ -0,0 +1,17 @@
+from .base import IRISDialect
+
+
+class IRISDialect_irisasync(IRISDialect):
+    driver = "irisasync"
+
+    is_async = True
+    supports_statement_cache = True
+
+    @classmethod
+    def import_dbapi(cls):
+        import intersystems_iris.dbapi._DBAPI as dbapi
+
+        return dbapi
+
+
+dialect = IRISDialect_irisasync
diff --git a/tests/conftest.py b/tests/conftest.py
index 5922b7d..f594811 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -8,6 +8,9 @@
 
 registry.register("iris.iris", "sqlalchemy_iris.iris", "IRISDialect_iris")
 registry.register("iris.emb", "sqlalchemy_iris.embedded", "IRISDialect_emb")
+registry.register(
+    "iris.irisasync", "sqlalchemy_iris.irisasync", "IRISDialect_irisasync"
+)
 
 pytest.register_assert_rewrite("sqlalchemy.testing.assertions")