Annotate missing fields

Author: Sergey Vasilyev

data_diff/databases/_connect.py

@@ -95,6 +95,8 @@ def match_path(self, dsn):
 class Connect:
     """Provides methods for connecting to a supported database using a URL or connection dict."""
 
+    database_by_scheme: Dict[str, Database]
+    match_uri_path: Dict[str, MatchUriPath]
     conn_cache: MutableMapping[Hashable, Database]
 
     def __init__(self, database_by_scheme: Dict[str, Database] = DATABASE_BY_SCHEME):

data_diff/databases/base.py

@@ -5,7 +5,7 @@
 import math
 import sys
 import logging
-from typing import Any, Callable, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union, TypeVar
+from typing import Any, Callable, ClassVar, Dict, Generator, Tuple, Optional, Sequence, Type, List, Union, TypeVar
 from functools import partial, wraps
 from concurrent.futures import ThreadPoolExecutor
 import threading
@@ -179,6 +179,9 @@ class ThreadLocalInterpreter:
     Useful for cursor-sensitive operations, such as creating a temporary table.
     """
 
+    compiler: Compiler
+    gen: Generator
+
     def __init__(self, compiler: Compiler, gen: Generator):
         super().__init__()
         self.gen = gen
@@ -238,9 +241,9 @@ def optimizer_hints(self, hints: str) -> str:
 
 
 class BaseDialect(abc.ABC):
-    SUPPORTS_PRIMARY_KEY = False
-    SUPPORTS_INDEXES = False
-    TYPE_CLASSES: Dict[str, type] = {}
+    SUPPORTS_PRIMARY_KEY: ClassVar[bool] = False
+    SUPPORTS_INDEXES: ClassVar[bool] = False
+    TYPE_CLASSES: ClassVar[Dict[str, type]] = {}
 
     PLACEHOLDER_TABLE = None  # Used for Oracle
 
@@ -835,14 +838,13 @@ class Database(abc.ABC, _RuntypeHackToFixCicularRefrencedDatabase):
     Instanciated using :meth:`~data_diff.connect`
     """
 
-    default_schema: str = None
-    SUPPORTS_ALPHANUMS = True
-    SUPPORTS_UNIQUE_CONSTAINT = False
-
-    CONNECT_URI_KWPARAMS = []
+    SUPPORTS_ALPHANUMS: ClassVar[bool] = True
+    SUPPORTS_UNIQUE_CONSTAINT: ClassVar[bool] = False
+    CONNECT_URI_KWPARAMS: ClassVar[List[str]] = []
 
-    _interactive = False
-    is_closed = False
+    default_schema: Optional[str] = None
+    _interactive: bool = False
+    is_closed: bool = False
 
     @property
     def name(self):
@@ -1109,6 +1111,10 @@ class ThreadedDatabase(Database):
     Used for database connectors that do not support sharing their connection between different threads.
     """
 
+    _init_error: Optional[Exception]
+    _queue: ThreadPoolExecutor
+    thread_local: threading.local
+
     def __init__(self, thread_count=1):
         super().__init__()
         self._init_error = None

data_diff/databases/bigquery.py

@@ -223,6 +223,10 @@ class BigQuery(Database):
     CONNECT_URI_PARAMS = ["dataset"]
     dialect = Dialect()
 
+    project: str
+    dataset: str
+    _client: Any
+
     def __init__(self, project, *, dataset, bigquery_credentials=None, **kw):
         super().__init__()
         credentials = bigquery_credentials

data_diff/databases/clickhouse.py

@@ -1,4 +1,4 @@
-from typing import Optional, Type
+from typing import Any, Dict, Optional, Type
 
 from data_diff.databases.base import (
     MD5_HEXDIGITS,
@@ -167,6 +167,8 @@ class Clickhouse(ThreadedDatabase):
     CONNECT_URI_HELP = "clickhouse://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
 
+    _args: Dict[str, Any]
+
     def __init__(self, *, thread_count: int, **kw):
         super().__init__(thread_count=thread_count)
 

data_diff/databases/databricks.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, Sequence
+from typing import Any, Dict, Sequence
 import logging
 
 from data_diff.abcs.database_types import (
@@ -103,13 +103,16 @@ class Databricks(ThreadedDatabase):
     CONNECT_URI_HELP = "databricks://:<access_token>@<server_hostname>/<http_path>"
     CONNECT_URI_PARAMS = ["catalog", "schema"]
 
+    catalog: str
+    _args: Dict[str, Any]
+
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         logging.getLogger("databricks.sql").setLevel(logging.WARNING)
 
         self._args = kw
         self.default_schema = kw.get("schema", "default")
-        self.catalog = self._args.get("catalog", "hive_metastore")
+        self.catalog = kw.get("catalog", "hive_metastore")
 
     def create_connection(self):
         databricks = import_databricks()

data_diff/databases/duckdb.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Any, Dict, Union
 
 from data_diff.utils import match_regexps
 from data_diff.abcs.database_types import (
@@ -134,14 +134,17 @@ def current_timestamp(self) -> str:
 class DuckDB(Database):
     dialect = Dialect()
     SUPPORTS_UNIQUE_CONSTAINT = False  # Temporary, until we implement it
-    default_schema = "main"
     CONNECT_URI_HELP = "duckdb://<dbname>@<filepath>"
     CONNECT_URI_PARAMS = ["database", "dbpath"]
 
+    _args: Dict[str, Any]
+    _conn: Any
+
     def __init__(self, **kw):
         super().__init__()
         self._args = kw
         self._conn = self.create_connection()
+        self.default_schema = "main"
 
     @property
     def is_autocommit(self) -> bool:

data_diff/databases/mssql.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Any, Dict, Optional
 from data_diff.abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue
 from data_diff.databases.base import (
     CHECKSUM_HEXDIGITS,
@@ -160,10 +160,12 @@ def constant_values(self, rows) -> str:
 
 class MsSQL(ThreadedDatabase):
     dialect = Dialect()
-    #
     CONNECT_URI_HELP = "mssql://<user>:<password>@<host>/<database>/<schema>"
     CONNECT_URI_PARAMS = ["database", "schema"]
 
+    default_database: str
+    _args: Dict[str, Any]
+
     def __init__(self, host, port, user, password, *, database, thread_count, **kw):
         args = dict(server=host, port=port, database=database, user=user, password=password, **kw)
         self._args = {k: v for k, v in args.items() if v is not None}

data_diff/databases/mysql.py

@@ -1,3 +1,5 @@
+from typing import Any, Dict
+
 from data_diff.abcs.database_types import (
     Datetime,
     Timestamp,
@@ -137,6 +139,8 @@ class MySQL(ThreadedDatabase):
     CONNECT_URI_HELP = "mysql://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
 
+    _args: Dict[str, Any]
+
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         self._args = kw

data_diff/databases/oracle.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 
 from data_diff.utils import match_regexps
 from data_diff.abcs.database_types import (
@@ -181,6 +181,8 @@ class Oracle(ThreadedDatabase):
     CONNECT_URI_HELP = "oracle://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
 
+    kwargs: Dict[str, Any]
+
     def __init__(self, *, host, database, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         self.kwargs = dict(dsn=f"{host}/{database}" if database else host, **kw)

data_diff/databases/postgresql.py

@@ -1,5 +1,6 @@
-from typing import List
+from typing import Any, ClassVar, Dict, List, Type
 from data_diff.abcs.database_types import (
+    ColType,
     DbPath,
     JSON,
     Timestamp,
@@ -68,7 +69,7 @@ class PostgresqlDialect(
     SUPPORTS_PRIMARY_KEY = True
     SUPPORTS_INDEXES = True
 
-    TYPE_CLASSES = {
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
         # Timestamps
         "timestamp with time zone": TimestampTZ,
         "timestamp without time zone": Timestamp,
@@ -125,11 +126,12 @@ class PostgreSQL(ThreadedDatabase):
     CONNECT_URI_HELP = "postgresql://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
 
-    default_schema = "public"
+    _args: Dict[str, Any]
 
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         self._args = kw
+        self.default_schema = "public"
 
     def create_connection(self):
         if not self._args:

data_diff/databases/presto.py

@@ -1,5 +1,6 @@
 from functools import partial
 import re
+from typing import Any
 
 from data_diff.utils import match_regexps
 
@@ -158,10 +159,11 @@ class Presto(Database):
     CONNECT_URI_HELP = "presto://<user>@<host>/<catalog>/<schema>"
     CONNECT_URI_PARAMS = ["catalog", "schema"]
 
-    default_schema = "public"
+    _conn: Any
 
     def __init__(self, **kw):
         super().__init__()
+        self.default_schema = "public"
         prestodb = import_presto()
 
         if kw.get("schema"):

data_diff/databases/redshift.py

@@ -1,5 +1,6 @@
-from typing import List, Dict
+from typing import ClassVar, List, Dict, Type
 from data_diff.abcs.database_types import (
+    ColType,
     Float,
     JSON,
     TemporalType,
@@ -53,7 +54,7 @@ def normalize_json(self, value: str, _coltype: JSON) -> str:
 
 class Dialect(PostgresqlDialect, Mixin_MD5, Mixin_NormalizeValue, AbstractMixin_MD5, AbstractMixin_NormalizeValue):
     name = "Redshift"
-    TYPE_CLASSES = {
+    TYPE_CLASSES: ClassVar[Dict[str, Type[ColType]]] = {
         **PostgresqlDialect.TYPE_CLASSES,
         "double": Float,
         "real": Float,

data_diff/databases/snowflake.py

@@ -1,4 +1,4 @@
-from typing import Union, List
+from typing import Any, Union, List
 import logging
 
 from data_diff.abcs.database_types import (
@@ -154,6 +154,8 @@ class Snowflake(Database):
     CONNECT_URI_PARAMS = ["database", "schema"]
     CONNECT_URI_KWPARAMS = ["warehouse"]
 
+    _conn: Any
+
     def __init__(self, *, schema: str, **kw):
         super().__init__()
         snowflake, serialization, default_backend = import_snowflake()

data_diff/databases/trino.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from data_diff.abcs.mixins import AbstractMixin_MD5, AbstractMixin_NormalizeValue
 from data_diff.abcs.database_types import TemporalType, ColType_UUID
 from data_diff.databases import presto
@@ -39,6 +41,8 @@ class Trino(presto.Presto):
     CONNECT_URI_HELP = "trino://<user>@<host>/<catalog>/<schema>"
     CONNECT_URI_PARAMS = ["catalog", "schema"]
 
+    _conn: Any
+
     def __init__(self, **kw):
         super().__init__()
         trino = import_trino()

data_diff/databases/vertica.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Any, Dict, List
 
 from data_diff.utils import match_regexps
 from data_diff.databases.base import (
@@ -156,12 +156,13 @@ class Vertica(ThreadedDatabase):
     CONNECT_URI_HELP = "vertica://<user>:<password>@<host>/<database>"
     CONNECT_URI_PARAMS = ["database?"]
 
-    default_schema = "public"
+    _args: Dict[str, Any]
 
     def __init__(self, *, thread_count, **kw):
         super().__init__(thread_count=thread_count)
         self._args = kw
         self._args["AUTOCOMMIT"] = False
+        self.default_schema = "public"
 
     def create_connection(self):
         vertica = import_vertica()

data_diff/thread_utils.py

@@ -45,6 +45,11 @@ class ThreadedYielder(Iterable):
     Priority for the iterator can be provided via the keyword argument 'priority'. (higher runs first)
     """
 
+    _pool: ThreadPoolExecutor
+    _futures: deque
+    _yield: deque
+    _exception: Optional[None]
+
     def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._pool = PriorityThreadPoolExecutor(max_workers)

data_diff/utils.py

@@ -4,7 +4,7 @@
 import re
 import string
 from abc import abstractmethod
-from typing import Any, Dict, Iterable, Iterator, List, MutableMapping, Sequence, TypeVar, Union
+from typing import Any, Dict, Iterable, Iterator, List, MutableMapping, Optional, Sequence, TypeVar, Union
 from urllib.parse import urlparse
 import operator
 import threading
@@ -175,6 +175,9 @@ def alphanums_to_numbers(s1: str, s2: str):
 
 
 class ArithAlphanumeric(ArithString):
+    _str: str
+    _max_len: Optional[int]
+
     def __init__(self, s: str, max_len=None):
         super().__init__()