support multithreading for databricks

The databricks connector is not thread-safe so we should inherit ThreadedDatabase class

Author: pik94

data_diff/databases/databricks.py

@@ -13,7 +13,7 @@
  ColType,
  UnknownColType,
 )
-from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, BaseDialect, Database, import_helper, parse_table_name
+from .base import MD5_HEXDIGITS, CHECKSUM_HEXDIGITS, BaseDialect, ThreadedDatabase, import_helper, parse_table_name
 
 
 @import_helper(text="You can install it using 'pip install databricks-sql-connector'")
@@ -68,43 +68,45 @@ def _convert_db_precision_to_digits(self, p: int) -> int:
  return max(super()._convert_db_precision_to_digits(p) - 1, 0)
 
 
-class Databricks(Database):
+class Databricks(ThreadedDatabase):
  dialect = Dialect()
 
- def __init__(
- self,
- http_path: str,
- access_token: str,
- server_hostname: str,
- catalog: str = "hive_metastore",
- schema: str = "default",
- **kwargs,
- ):
- databricks = import_databricks()
-
- self._conn = databricks.sql.connect(
- server_hostname=server_hostname, http_path=http_path, access_token=access_token, catalog=catalog
- )
-
+ def __init__(self, *, thread_count, **kw):
  logging.getLogger("databricks.sql").setLevel(logging.WARNING)
 
- self.catalog = catalog
- self.default_schema = schema
- self.kwargs = kwargs
+ self._args = kw
+ self.default_schema = kw.get('schema', 'hive_metastore')
+ super().__init__(thread_count=thread_count)
 
- def _query(self, sql_code: str) -> list:
- "Uses the standard SQL cursor interface"
- return self._query_conn(self._conn, sql_code)
+ def create_connection(self):
+ databricks = import_databricks()
+
+ try:
+ return databricks.sql.connect(
+ server_hostname=self._args['server_hostname'],
+ http_path=self._args['http_path'],
+ access_token=self._args['access_token'],
+ catalog=self._args['catalog'],
+ )
+ except databricks.sql.exc.Error as e:
+ raise ConnectionError(*e.args) from e
 
  def query_table_schema(self, path: DbPath) -> Dict[str, tuple]:
  # Databricks has INFORMATION_SCHEMA only for Databricks Runtime, not for Databricks SQL.
  # https://docs.databricks.com/spark/latest/spark-sql/language-manual/information-schema/columns.html
  # So, to obtain information about schema, we should use another approach.
 
+ conn = self.create_connection()
+
  schema, table = self._normalize_table_path(path)
- with self._conn.cursor() as cursor:
- cursor.columns(catalog_name=self.catalog, schema_name=schema, table_name=table)
- rows = cursor.fetchall()
+ with conn.cursor() as cursor:
+ cursor.columns(catalog_name=self._args['catalog'], schema_name=schema, table_name=table)
+ try:
+ rows = cursor.fetchall()
+ except:
+ rows = None
+ finally:
+ conn.close()
  if not rows:
  raise RuntimeError(f"{self.name}: Table '{'.'.join(path)}' does not exist, or has no columns")
 
@@ -121,7 +123,7 @@ def _process_table_schema(
  resulted_rows = []
  for row in rows:
  row_type = "DECIMAL" if row[1].startswith("DECIMAL") else row[1]
- type_cls = self.TYPE_CLASSES.get(row_type, UnknownColType)
+ type_cls = self.dialect.TYPE_CLASSES.get(row_type, UnknownColType)
 
  if issubclass(type_cls, Integer):
  row = (row[0], row_type, None, None, 0)
@@ -152,9 +154,6 @@ def parse_table_name(self, name: str) -> DbPath:
  path = parse_table_name(name)
  return self._normalize_table_path(path)
 
- def close(self):
- self._conn.close()
-
  @property
  def is_autocommit(self) -> bool:
  return True