fix #53 (#55)

* fix #53

* fix get_database_version

Author: jayvynl

CHANGELOG.md

@@ -1,6 +1,7 @@
 ### 1.1.3
 - Fix #50 partition by single expression raise TypeError.
 - Fix #51 .
+- Fix #53 .
 
 ### 1.1.2
 - Use [flake8](https://flake8.pycqa.org/) to lint code.

clickhouse_backend/__init__.py

@@ -1,5 +1,5 @@
 from clickhouse_backend.utils.version import get_version
 
-VERSION = (1, 1, 3, "alpha", 0)
+VERSION = (1, 1, 3, "final", 0)
 
 __version__ = get_version(VERSION)

clickhouse_backend/backend/base.py

@@ -1,3 +1,5 @@
+import threading
+
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.db.backends.base.base import BaseDatabaseWrapper
@@ -16,6 +18,9 @@
 
 
 class DatabaseWrapper(BaseDatabaseWrapper):
+    _connection_sharing_lock = threading.Lock()
+    _clickhouse_connections = {}
+
     vendor = "clickhouse"
     display_name = "ClickHouse"
     # This dictionary maps Field objects to their associated ClickHouse column
@@ -200,8 +205,17 @@ def get_connection_params(self):
 
     @async_unsafe
     def get_new_connection(self, conn_params):
-        connection = Database.connect(**conn_params)
-        return connection
+        # Fix https://github.com/jayvynl/django-clickhouse-backend/issues/53
+        with self._connection_sharing_lock:
+            if self.alias in self._clickhouse_connections:
+                params, conn = self._clickhouse_connections[self.alias]
+                if conn_params == params:
+                    return conn
+
+            conn = Database.connect(**conn_params)
+            self._clickhouse_connections[self.alias] = (conn_params, conn)
+
+        return conn
 
     def init_connection_state(self):
         pass
@@ -222,6 +236,13 @@ def _savepoint_commit(self, sid):
     def _set_autocommit(self, autocommit):
         pass
 
+    def _close(self):
+        """Close database connection.
+
+        This is a noop, because inner connection is shared between threads.
+        """
+        pass
+
     def is_usable(self):
         try:
             # Use a clickhouse_driver cursor directly, bypassing Django's utilities.

tests/backends/tests.py

@@ -259,11 +259,7 @@ def create_squares(self, args, paramstyle, multiple):
         if paramstyle == "format":
             query = "INSERT INTO %s (%s, %s) VALUES" % (tbl, f1, f2)
         elif paramstyle == "pyformat":
-            query = "INSERT INTO %s (%s, %s) VALUES" % (
-                tbl,
-                f1,
-                f2,
-            )
+            query = "INSERT INTO %s (%s, %s) VALUES" % (tbl, f1, f2)
         else:
             raise ValueError("unsupported paramstyle in test")
         with connection.cursor() as cursor:
@@ -369,7 +365,7 @@ def test_unicode_password(self):
             self.fail("Unexpected error raised with Unicode password: %s" % e)
         finally:
             connection.settings_dict["PASSWORD"] = old_password
-            connection.connection.close()
+            connection.close()
             connection.connection = None
 
     def test_database_operations_helper_class(self):
@@ -431,11 +427,10 @@ def test_is_usable_after_database_disconnects(self):
         # Open a connection to the database.
         with connection.cursor():
             pass
-        # Emulate a connection close by the database.
+        # _close is a noop, because underlying connection is shared between threads.
         connection._close()
-        # Even then is_usable() should not raise an exception.
         try:
-            self.assertFalse(connection.is_usable())
+            self.assertTrue(connection.is_usable())
         finally:
             # Clean up the mess created by connection._close(). Since the
             # connection is already closed, this crashes on some backends.
@@ -705,6 +700,7 @@ def test_check_constraints_sql_keywords(self):
 
 class ThreadTests(TransactionTestCase):
     available_apps = ["backends"]
+    databases = {"default", "s1r2", "s2r1"}
 
     def test_default_connection_thread_local(self):
         """
@@ -736,15 +732,17 @@ def runner():
                 t = threading.Thread(target=runner)
                 t.start()
                 t.join()
-            # Each created connection got different inner connection.
+            # The inner connection is shared between threads.
             self.assertEqual(
-                len({conn.connection for conn in connections_dict.values()}), 3
+                len({conn.connection for conn in connections_dict.values()}), 1
             )
         finally:
             # Finish by closing the connections opened by the other threads
             # (the connection opened in the main thread will automatically be
             # closed on teardown).
             for conn in connections_dict.values():
+                # Regression test for https://github.com/jayvynl/django-clickhouse-backend/issues/53
+                self.assertIs(conn.connection, connection.connection)
                 if conn is not connection and conn.allow_thread_sharing:
                     conn.close()
                     conn.dec_thread_sharing()
@@ -783,6 +781,10 @@ def runner():
             # (the connection opened in the main thread will automatically be
             # closed on teardown).
             for conn in connections_dict.values():
+                # Regression test for https://github.com/jayvynl/django-clickhouse-backend/issues/53
+                if conn.vendor == "clickhouse":
+                    conn.ensure_connection()
+                    self.assertIs(conn.connection, connections[conn.alias].connection)
                 if conn is not connection and conn.allow_thread_sharing:
                     conn.close()
                     conn.dec_thread_sharing()