refactor: Refactor datasource module (#1309)

dbgpt/datasource/rdbms/dialect/__init__.py

@@ -0,0 +1 @@
+"""Module for RDBMS dialects."""

dbgpt/datasource/rdbms/dialect/starrocks/__init__.py

@@ -1,4 +1,4 @@
-#! /usr/bin/python3
+"""StarRocks dialect for SQLAlchemy."""
 # Copyright 2021-present StarRocks, Inc. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

dbgpt/datasource/rdbms/dialect/starrocks/sqlalchemy/__init__.py

@@ -1,4 +1,4 @@
-#! /usr/bin/python3
+"""SQLAlchemy dialect for StarRocks."""
 # Copyright 2021-present StarRocks, Inc. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

dbgpt/datasource/rdbms/dialect/starrocks/sqlalchemy/datatype.py

@@ -1,3 +1,5 @@
+"""SQLAlchemy data types for StarRocks."""
+
 import logging
 import re
 from typing import Any, Dict, List, Optional, Type
@@ -10,50 +12,71 @@ logger = logging.getLogger(__name__)
 
 
 class TINYINT(Integer):  # pylint: disable=no-init
+    """StarRocks TINYINT type."""
+
     __visit_name__ = "TINYINT"
 
 
 class LARGEINT(Integer):  # pylint: disable=no-init
+    """StarRocks LARGEINT type."""
+
     __visit_name__ = "LARGEINT"
 
 
 class DOUBLE(Float):  # pylint: disable=no-init
+    """StarRocks DOUBLE type."""
+
     __visit_name__ = "DOUBLE"
 
 
 class HLL(Numeric):  # pylint: disable=no-init
+    """StarRocks HLL type."""
+
     __visit_name__ = "HLL"
 
 
 class BITMAP(Numeric):  # pylint: disable=no-init
+    """StarRocks BITMAP type."""
+
     __visit_name__ = "BITMAP"
 
 
 class PERCENTILE(Numeric):  # pylint: disable=no-init
+    """StarRocks PERCENTILE type."""
+
     __visit_name__ = "PERCENTILE"
 
 
 class ARRAY(TypeEngine):  # pylint: disable=no-init
+    """StarRocks ARRAY type."""
+
     __visit_name__ = "ARRAY"
 
     @property
-    def python_type(self) -> Optional[Type[List[Any]]]:
+    def python_type(self) -> Optional[Type[List[Any]]]:  # type: ignore
+        """Return the Python type for this SQL type."""
         return list
 
 
 class MAP(TypeEngine):  # pylint: disable=no-init
+    """StarRocks MAP type."""
+
     __visit_name__ = "MAP"
 
     @property
-    def python_type(self) -> Optional[Type[Dict[Any, Any]]]:
+    def python_type(self) -> Optional[Type[Dict[Any, Any]]]:  # type: ignore
+        """Return the Python type for this SQL type."""
         return dict
 
 
 class STRUCT(TypeEngine):  # pylint: disable=no-init
+    """StarRocks STRUCT type."""
+
     __visit_name__ = "STRUCT"
 
     @property
-    def python_type(self) -> Optional[Type[Any]]:
+    def python_type(self) -> Optional[Type[Any]]:  # type: ignore
+        """Return the Python type for this SQL type."""
         return None
 
 
@@ -90,6 +113,7 @@ _type_map = {
 
 
 def parse_sqltype(type_str: str) -> TypeEngine:
+    """Parse a SQL type string into a SQLAlchemy type object."""
     type_str = type_str.strip().lower()
     match = re.match(r"^(?P<type>\w+)\s*(?:\((?P<options>.*)\))?", type_str)
     if not match:

dbgpt/datasource/rdbms/dialect/starrocks/sqlalchemy/dialect.py

@@ -1,4 +1,4 @@
-#! /usr/bin/python3
+"""StarRocks dialect for SQLAlchemy."""
 # Copyright 2021-present StarRocks, Inc. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, cast
 
 from sqlalchemy import exc, log, text
 from sqlalchemy.dialects.mysql.pymysql import MySQLDialect_pymysql
@@ -25,7 +25,9 @@ logger = logging.getLogger(__name__)
 
 
 @log.class_logger
-class StarRocksDialect(MySQLDialect_pymysql):
+class StarRocksDialect(MySQLDialect_pymysql):  # type: ignore
+    """StarRocks dialect for SQLAlchemy."""
+
     # Caching
     # Warnings are generated by SQLAlchmey if this flag is not explicitly set
     # and tests are needed before being enabled
@@ -34,9 +36,11 @@ class StarRocksDialect(MySQLDialect_pymysql):
     name = "starrocks"
 
     def __init__(self, *args, **kw):
+        """Create a new StarRocks dialect."""
         super(StarRocksDialect, self).__init__(*args, **kw)
 
-    def has_table(self, connection, table_name, schema=None, **kw):
+    def has_table(self, connection, table_name, schema: Optional[str] = None, **kw):
+        """Return True if the given table is present in the database."""
         self._ensure_has_table_connection(connection)
 
         if schema is None:
@@ -53,15 +57,13 @@ class StarRocksDialect(MySQLDialect_pymysql):
         return res.first() is not None
 
     def get_schema_names(self, connection, **kw):
+        """Return a list of schema names available in the database."""
         rp = connection.exec_driver_sql("SHOW schemas")
         return [r[0] for r in rp]
 
-    def get_table_names(self, connection, schema=None, **kw):
+    def get_table_names(self, connection, schema: Optional[str] = None, **kw):
         """Return a Unicode SHOW TABLES from a given schema."""
-        if schema is not None:
-            current_schema = schema
-        else:
-            current_schema = self.default_schema_name
+        current_schema: str = cast(str, schema or self.default_schema_name)
 
         charset = self._connection_charset
 
@@ -76,13 +78,15 @@ class StarRocksDialect(MySQLDialect_pymysql):
             if row[1] == "BASE TABLE"
         ]
 
-    def get_view_names(self, connection, schema=None, **kw):
+    def get_view_names(self, connection, schema: Optional[str] = None, **kw):
+        """Return a Unicode SHOW TABLES from a given schema."""
         if schema is None:
             schema = self.default_schema_name
+        current_schema = cast(str, schema)
         charset = self._connection_charset
         rp = connection.exec_driver_sql(
             "SHOW FULL TABLES FROM %s"
-            % self.identifier_preparer.quote_identifier(schema)
+            % self.identifier_preparer.quote_identifier(current_schema)
         )
         return [
             row[0]
@@ -90,9 +94,14 @@ class StarRocksDialect(MySQLDialect_pymysql):
             if row[1] in ("VIEW", "SYSTEM VIEW")
         ]
 
-    def get_columns(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
-    ) -> List[Dict[str, Any]]:
+    def get_columns(  # type: ignore
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
+    ) -> List[Dict[str, Any]]:  # type: ignore
+        """Return information about columns in `table_name`."""
         if not self.has_table(connection, table_name, schema):
             raise exc.NoSuchTableError(f"schema={schema}, table={table_name}")
         schema = schema or self._get_default_schema_name(connection)
@@ -114,60 +123,100 @@ class StarRocksDialect(MySQLDialect_pymysql):
             columns.append(column)
         return columns
 
-    def get_pk_constraint(self, connection, table_name, schema=None, **kw):
+    def get_pk_constraint(
+        self, connection, table_name, schema: Optional[str] = None, **kw
+    ):
+        """Return information about the primary key constraint."""
         return {  # type: ignore  # pep-655 not supported
             "name": None,
             "constrained_columns": [],
         }
 
-    def get_unique_constraints(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
+    def get_unique_constraints(  # type: ignore
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> List[Dict[str, Any]]:
+        """Return information about unique constraints."""
         return []
 
-    def get_check_constraints(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
+    def get_check_constraints(  # type: ignore
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> List[Dict[str, Any]]:
+        """Return information about check constraints."""
         return []
 
-    def get_foreign_keys(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
+    def get_foreign_keys(  # type: ignore
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> List[Dict[str, Any]]:
+        """Return information about foreign keys."""
         return []
 
     def get_primary_keys(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> List[str]:
+        """Return the primary key columns of the given table."""
         pk = self.get_pk_constraint(connection, table_name, schema)
         return pk.get("constrained_columns")  # type: ignore
 
-    def get_indexes(self, connection, table_name, schema=None, **kw):
+    def get_indexes(self, connection, table_name, schema: Optional[str] = None, **kw):
+        """Get table indexes about specified table."""
         return []
 
     def has_sequence(
-        self, connection: Connection, sequence_name: str, schema: str = None, **kw
+        self,
+        connection: Connection,
+        sequence_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> bool:
+        """Return True if the given sequence is present in the database."""
         return False
 
     def get_sequence_names(
-        self, connection: Connection, schema: str = None, **kw
+        self, connection: Connection, schema: Optional[str] = None, **kw
     ) -> List[str]:
+        """Return a list of sequence names."""
         return []
 
     def get_temp_view_names(
-        self, connection: Connection, schema: str = None, **kw
+        self, connection: Connection, schema: Optional[str] = None, **kw
     ) -> List[str]:
+        """Return a list of temporary view names."""
         return []
 
     def get_temp_table_names(
-        self, connection: Connection, schema: str = None, **kw
+        self, connection: Connection, schema: Optional[str] = None, **kw
     ) -> List[str]:
+        """Return a list of temporary table names."""
         return []
 
-    def get_table_options(self, connection, table_name, schema=None, **kw):
+    def get_table_options(
+        self, connection, table_name, schema: Optional[str] = None, **kw
+    ):
+        """Return a dictionary of options specified when the table was created."""
         return {}
 
-    def get_table_comment(
-        self, connection: Connection, table_name: str, schema: str = None, **kw
+    def get_table_comment(  # type: ignore
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: Optional[str] = None,
+        **kw,
     ) -> Dict[str, Any]:
+        """Return the comment for a table."""
         return dict(text=None)