acryl-datahub 1.0.0rc4__py3-none-any.whl → 1.0.0rc5__py3-none-any.whl

datahub/sql_parsing/split_statements.py

@@ -8,11 +8,11 @@ END_KEYWORD = "END"
 
 CONTROL_FLOW_KEYWORDS = [
     "GO",
-    r"BEGIN\w+TRY",
-    r"BEGIN\w+CATCH",
+    r"BEGIN\s+TRY",
+    r"BEGIN\s+CATCH",
     "BEGIN",
-    r"END\w+TRY",
-    r"END\w+CATCH",
+    r"END\s+TRY",
+    r"END\s+CATCH",
     # This isn't strictly correct, but we assume that IF | (condition) | (block) should all be split up
     # This mainly ensures that IF statements don't get tacked onto the previous statement incorrectly
     "IF",
@@ -73,25 +73,31 @@ class _StatementSplitter:
         # what a given END is closing.
         self.current_case_statements = 0
 
-    def _is_keyword_at_position(self, pos: int, keyword: str) -> bool:
+    def _is_keyword_at_position(self, pos: int, keyword: str) -> Tuple[bool, str]:
         """
         Check if a keyword exists at the given position using regex word boundaries.
         """
         sql = self.sql
 
-        if pos + len(keyword) > len(sql):
-            return False
+        keyword_length = len(keyword.replace(r"\s+", " "))
+
+        if pos + keyword_length > len(sql):
+            return False, ""
 
         # If we're not at a word boundary, we can't generate a keyword.
         if pos > 0 and not (
             bool(re.match(r"\w\W", sql[pos - 1 : pos + 1]))
             or bool(re.match(r"\W\w", sql[pos - 1 : pos + 1]))
         ):
-            return False
+            return False, ""
 
-        pattern = rf"^{re.escape(keyword)}\b"
+        pattern = rf"^{keyword}\b"
         match = re.match(pattern, sql[pos:], re.IGNORECASE)
-        return bool(match)
+        is_match = bool(match)
+        actual_match = (
+            sql[pos:][match.start() : match.end()] if match is not None else ""
+        )
+        return is_match, actual_match
 
     def _look_ahead_for_keywords(self, keywords: List[str]) -> Tuple[bool, str, int]:
         """
@@ -99,7 +105,8 @@ class _StatementSplitter:
         """
 
         for keyword in keywords:
-            if self._is_keyword_at_position(self.i, keyword):
+            is_match, keyword = self._is_keyword_at_position(self.i, keyword)
+            if is_match:
                 return True, keyword, len(keyword)
         return False, "", 0
 
@@ -118,7 +125,7 @@ class _StatementSplitter:
 
     def process(self) -> Iterator[str]:
         if not self.sql or not self.sql.strip():
-            return
+            yield from ()
 
         prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
         while self.i < len(self.sql):
@@ -181,7 +188,7 @@ class _StatementSplitter:
     def _process_normal(self, most_recent_real_char: str) -> Iterator[str]:
         c = self.sql[self.i]
 
-        if self._is_keyword_at_position(self.i, CASE_KEYWORD):
+        if self._is_keyword_at_position(self.i, CASE_KEYWORD)[0]:
             self.current_case_statements += 1
 
         is_control_keyword, keyword, keyword_len = self._look_ahead_for_keywords(

datahub/utilities/file_backed_collections.py

@@ -10,13 +10,11 @@ import tempfile
 import threading
 from dataclasses import dataclass, field
 from datetime import datetime
-from enum import Enum
 from types import TracebackType
 from typing import (
     Any,
     Callable,
     Dict,
-    Final,
     Generic,
     Iterator,
     List,
@@ -31,6 +29,7 @@ from typing import (
 )
 
 from datahub.ingestion.api.closeable import Closeable
+from datahub.utilities.sentinels import Unset, unset
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -59,16 +58,6 @@ SqliteValue = Union[int, float, str, bytes, datetime, None]
 _VT = TypeVar("_VT")
 
 
-class Unset(Enum):
-    token = 0
-
-
-# It's pretty annoying to create a true sentinel that works with typing.
-# https://peps.python.org/pep-0484/#support-for-singleton-types-in-unions
-# Can't wait for https://peps.python.org/pep-0661/
-_unset: Final = Unset.token
-
-
 class ConnectionWrapper:
     """
     Wraps a SQlite connection, allowing connection reuse across multiple FileBacked* objects.
@@ -372,7 +361,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         self,
         /,
         key: str,
-        default: Union[_VT, Unset] = _unset,
+        default: Union[_VT, Unset] = unset,
     ) -> _VT:
         # If key is in the dictionary, this is similar to __getitem__ + mark_dirty.
         # If key is not in the dictionary, this is similar to __setitem__.
@@ -383,7 +372,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
             self.mark_dirty(key)
             return value
         except KeyError:
-            if default is _unset:
+            if default is unset:
                 raise
 
             self[key] = default

datahub/utilities/sentinels.py

@@ -0,0 +1,22 @@
+from enum import Enum
+from typing import Final
+
+# It's pretty annoying to create a true sentinel that works with typing.
+# This approach using enums is inspired by:
+# https://peps.python.org/pep-0484/#support-for-singleton-types-in-unions
+#
+# Can't wait for https://peps.python.org/pep-0661/
+
+
+class Unset(Enum):
+    token = 0
+
+
+unset: Final = Unset.token
+
+
+class Auto(Enum):
+    token = 0
+
+
+auto: Final = Auto.token