acryl-datahub 0.15.0.6rc1__py3-none-any.whl → 0.15.0.6rc3__py3-none-any.whl

datahub/sdk/main_client.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from typing import Optional, overload
+
+from datahub.errors import SdkUsageError
+from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import DatahubClientConfig
+from datahub.sdk.entity_client import EntityClient
+from datahub.sdk.resolver_client import ResolverClient
+
+
+class DataHubClient:
+    @overload
+    def __init__(self, *, server: str, token: Optional[str] = None): ...
+    @overload
+    def __init__(self, *, config: DatahubClientConfig): ...
+    @overload
+    def __init__(self, *, graph: DataHubGraph): ...
+    def __init__(
+        self,
+        *,
+        server: Optional[str] = None,
+        token: Optional[str] = None,
+        graph: Optional[DataHubGraph] = None,
+        config: Optional[DatahubClientConfig] = None,
+    ):
+        if server is not None:
+            if config is not None:
+                raise SdkUsageError("Cannot specify both server and config")
+            if graph is not None:
+                raise SdkUsageError("Cannot specify both server and graph")
+            graph = DataHubGraph(config=DatahubClientConfig(server=server, token=token))
+        elif config is not None:
+            if graph is not None:
+                raise SdkUsageError("Cannot specify both config and graph")
+            graph = DataHubGraph(config=config)
+        elif graph is None:
+            raise SdkUsageError("Must specify either server, config, or graph")
+
+        self._graph = graph
+
+    @classmethod
+    def from_env(cls) -> "DataHubClient":
+        # Inspired by the DockerClient.from_env() method.
+        # TODO: This one also reads from ~/.datahubenv, so the "from_env" name might be a bit confusing.
+        # That file is part of the "environment", but is not a traditional "env variable".
+        graph = get_default_graph()
+        return cls(graph=graph)
+
+    @property
+    def entities(self) -> EntityClient:
+        return EntityClient(self)
+
+    @property
+    def resolve(self) -> ResolverClient:
+        return ResolverClient(self)

datahub/sdk/resolver_client.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional, overload
+
+from datahub.errors import ItemNotFoundError, MultipleItemsFoundError, SdkUsageError
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.metadata.urns import (
+    CorpUserUrn,
+    DomainUrn,
+    GlossaryTermUrn,
+)
+
+if TYPE_CHECKING:
+    from datahub.sdk.main_client import DataHubClient
+
+
+class ResolverClient:
+    def __init__(self, client: DataHubClient):
+        self._client = client
+
+    # TODO: add caching to this method
+
+    @property
+    def _graph(self) -> DataHubGraph:
+        return self._client._graph
+
+    def domain(self, *, name: str) -> DomainUrn:
+        urn_str = self._graph.get_domain_urn_by_name(name)
+        if urn_str is None:
+            raise ItemNotFoundError(f"Domain with name {name} not found")
+        return DomainUrn.from_string(urn_str)
+
+    @overload
+    def user(self, *, name: str) -> CorpUserUrn: ...
+    @overload
+    def user(self, *, email: str) -> CorpUserUrn: ...
+    def user(
+        self, *, name: Optional[str] = None, email: Optional[str] = None
+    ) -> CorpUserUrn:
+        filter_explanation: str
+        filters = []
+        if name is not None:
+            if email is not None:
+                raise SdkUsageError("Cannot specify both name and email for auto_user")
+            # TODO: do we filter on displayName or fullName?
+            filter_explanation = f"with name {name}"
+            filters.append(
+                {
+                    "field": "fullName",
+                    "values": [name],
+                    "condition": "EQUAL",
+                }
+            )
+        elif email is not None:
+            filter_explanation = f"with email {email}"
+            filters.append(
+                {
+                    "field": "email",
+                    "values": [email],
+                    "condition": "EQUAL",
+                }
+            )
+        else:
+            raise SdkUsageError("Must specify either name or email for auto_user")
+
+        users = list(
+            self._graph.get_urns_by_filter(
+                entity_types=[CorpUserUrn.ENTITY_TYPE],
+                extraFilters=filters,
+            )
+        )
+        if len(users) == 0:
+            # TODO: In auto methods, should we just create the user/domain/etc if it doesn't exist?
+            raise ItemNotFoundError(f"User {filter_explanation} not found")
+        elif len(users) > 1:
+            raise MultipleItemsFoundError(
+                f"Multiple users found {filter_explanation}: {users}"
+            )
+        else:
+            return CorpUserUrn.from_string(users[0])
+
+    def term(self, *, name: str) -> GlossaryTermUrn:
+        # TODO: Add some limits on the graph fetch
+        terms = list(
+            self._graph.get_urns_by_filter(
+                entity_types=[GlossaryTermUrn.ENTITY_TYPE],
+                extraFilters=[
+                    {
+                        "field": "id",
+                        "values": [name],
+                        "condition": "EQUAL",
+                    }
+                ],
+            )
+        )
+        if len(terms) == 0:
+            raise ItemNotFoundError(f"Term with name {name} not found")
+        elif len(terms) > 1:
+            raise SdkUsageError(f"Multiple terms found with name {name}: {terms}")
+        else:
+            return GlossaryTermUrn.from_string(terms[0])

datahub/sql_parsing/split_statements.py

@@ -1,6 +1,10 @@
 import re
 from enum import Enum
-from typing import Generator, List, Tuple
+from typing import Iterator, List, Tuple
+
+SELECT_KEYWORD = "SELECT"
+CASE_KEYWORD = "CASE"
+END_KEYWORD = "END"
 
 CONTROL_FLOW_KEYWORDS = [
     "GO",
@@ -9,18 +13,36 @@ CONTROL_FLOW_KEYWORDS = [
     "BEGIN",
     r"END\w+TRY",
     r"END\w+CATCH",
-    "END",
+    # 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",
+    # For things like CASE, END does not mean the end of a statement.
+    # We have special handling for this.
+    END_KEYWORD,
+    # "ELSE",  # else is also valid in CASE, so we we can't use it here.
 ]
 
 # There's an exception to this rule, which is when the statement
-# is preceeded by a CTE.
-FORCE_NEW_STATEMENT_KEYWORDS = [
+# is preceded by a CTE. For those, we have to check if the character
+# before this is a ")".
+NEW_STATEMENT_KEYWORDS = [
     # SELECT is used inside queries as well, so we can't include it here.
+    "CREATE",
     "INSERT",
     "UPDATE",
     "DELETE",
     "MERGE",
 ]
+STRICT_NEW_STATEMENT_KEYWORDS = [
+    # For these keywords, a SELECT following it does indicate a new statement.
+    "DROP",
+    "TRUNCATE",
+]
+
+
+class _AlreadyIncremented(Exception):
+    # Using exceptions for control flow isn't great - but the code is clearer so it's fine.
+    pass
 
 
 class ParserState(Enum):
@@ -30,134 +52,199 @@ class ParserState(Enum):
     MULTILINE_COMMENT = 4
 
 
-def _is_keyword_at_position(sql: str, pos: int, keyword: str) -> bool:
-    """
-    Check if a keyword exists at the given position using regex word boundaries.
-    """
-    if pos + len(keyword) > len(sql):
-        return False
+class _StatementSplitter:
+    def __init__(self, sql: str):
+        self.sql = sql
 
-    # 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
+        # Main parser state.
+        self.i = 0
+        self.state = ParserState.NORMAL
+        self.current_statement: List[str] = []
 
-    pattern = rf"^{re.escape(keyword)}\b"
-    match = re.match(pattern, sql[pos:], re.IGNORECASE)
-    return bool(match)
+        # Additional parser state.
 
+        # If we see a SELECT, should we start a new statement?
+        # If we previously saw a drop/truncate/etc, a SELECT does mean a new statement.
+        # But if we're in a select/create/etc, a select could just be a subquery.
+        self.does_select_mean_new_statement = False
 
-def _look_ahead_for_keywords(
-    sql: str, pos: int, keywords: List[str]
-) -> Tuple[bool, str, int]:
-    """
-    Look ahead for SQL keywords at the current position.
-    """
+        # The END keyword terminates CASE and BEGIN blocks.
+        # We need to match the CASE statements with END blocks to determine
+        # what a given END is closing.
+        self.current_case_statements = 0
 
-    for keyword in keywords:
-        if _is_keyword_at_position(sql, pos, keyword):
-            return True, keyword, len(keyword)
-    return False, "", 0
+    def _is_keyword_at_position(self, pos: int, keyword: str) -> bool:
+        """
+        Check if a keyword exists at the given position using regex word boundaries.
+        """
+        sql = self.sql
 
+        if pos + len(keyword) > len(sql):
+            return False
 
-def split_statements(sql: str) -> Generator[str, None, None]:
-    """
-    Split T-SQL code into individual statements, handling various SQL constructs.
-    """
-    if not sql or not sql.strip():
-        return
+        # 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
+
+        pattern = rf"^{re.escape(keyword)}\b"
+        match = re.match(pattern, sql[pos:], re.IGNORECASE)
+        return bool(match)
 
-    current_statement: List[str] = []
-    state = ParserState.NORMAL
-    i = 0
+    def _look_ahead_for_keywords(self, keywords: List[str]) -> Tuple[bool, str, int]:
+        """
+        Look ahead for SQL keywords at the current position.
+        """
 
-    def yield_if_complete() -> Generator[str, None, None]:
-        statement = "".join(current_statement).strip()
+        for keyword in keywords:
+            if self._is_keyword_at_position(self.i, keyword):
+                return True, keyword, len(keyword)
+        return False, "", 0
+
+    def _yield_if_complete(self) -> Iterator[str]:
+        statement = "".join(self.current_statement).strip()
         if statement:
+            # Subtle - to avoid losing full whitespace, they get merged into the next statement.
             yield statement
-            current_statement.clear()
-
-    prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
-    while i < len(sql):
-        c = sql[i]
-        next_char = sql[i + 1] if i < len(sql) - 1 else "\0"
-
-        if state == ParserState.NORMAL:
-            if c == "'":
-                state = ParserState.STRING
-                current_statement.append(c)
-                prev_real_char = c
-            elif c == "-" and next_char == "-":
-                state = ParserState.COMMENT
-                current_statement.append(c)
-                current_statement.append(next_char)
-                i += 1
-            elif c == "/" and next_char == "*":
-                state = ParserState.MULTILINE_COMMENT
-                current_statement.append(c)
-                current_statement.append(next_char)
-                i += 1
-            else:
-                most_recent_real_char = prev_real_char
-                if not c.isspace():
+            self.current_statement.clear()
+
+        # Reset current_statement-specific state.
+        self.does_select_mean_new_statement = False
+        if self.current_case_statements != 0:
+            breakpoint()
+        self.current_case_statements = 0
+
+    def process(self) -> Iterator[str]:
+        if not self.sql or not self.sql.strip():
+            return
+
+        prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
+        while self.i < len(self.sql):
+            c = self.sql[self.i]
+            next_char = self.sql[self.i + 1] if self.i < len(self.sql) - 1 else "\0"
+
+            if self.state == ParserState.NORMAL:
+                if c == "'":
+                    self.state = ParserState.STRING
+                    self.current_statement.append(c)
                     prev_real_char = c
-
-                is_control_keyword, keyword, keyword_len = _look_ahead_for_keywords(
-                    sql, i, keywords=CONTROL_FLOW_KEYWORDS
-                )
-                if is_control_keyword:
-                    # Yield current statement if any
-                    yield from yield_if_complete()
-                    # Yield keyword as its own statement
-                    yield keyword
-                    i += keyword_len
-                    continue
-
-                (
-                    is_force_new_statement_keyword,
-                    keyword,
-                    keyword_len,
-                ) = _look_ahead_for_keywords(
-                    sql, i, keywords=FORCE_NEW_STATEMENT_KEYWORDS
-                )
-                if (
-                    is_force_new_statement_keyword and most_recent_real_char != ")"
-                ):  # usually we'd have a close paren that closes a CTE
-                    # Force termination of current statement
-                    yield from yield_if_complete()
-
-                    current_statement.append(keyword)
-                    i += keyword_len
-                    continue
-
-                elif c == ";":
-                    yield from yield_if_complete()
+                elif c == "-" and next_char == "-":
+                    self.state = ParserState.COMMENT
+                    self.current_statement.append(c)
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                elif c == "/" and next_char == "*":
+                    self.state = ParserState.MULTILINE_COMMENT
+                    self.current_statement.append(c)
+                    self.current_statement.append(next_char)
+                    self.i += 1
                 else:
-                    current_statement.append(c)
-
-        elif state == ParserState.STRING:
-            current_statement.append(c)
-            if c == "'" and next_char == "'":
-                current_statement.append(next_char)
-                i += 1
-            elif c == "'":
-                state = ParserState.NORMAL
-
-        elif state == ParserState.COMMENT:
-            current_statement.append(c)
-            if c == "\n":
-                state = ParserState.NORMAL
-
-        elif state == ParserState.MULTILINE_COMMENT:
-            current_statement.append(c)
-            if c == "*" and next_char == "/":
-                current_statement.append(next_char)
-                i += 1
-                state = ParserState.NORMAL
-
-        i += 1
-
-    # Handle the last statement
-    yield from yield_if_complete()
+                    most_recent_real_char = prev_real_char
+                    if not c.isspace():
+                        prev_real_char = c
+
+                    try:
+                        yield from self._process_normal(
+                            most_recent_real_char=most_recent_real_char
+                        )
+                    except _AlreadyIncremented:
+                        # Skip the normal i += 1 step.
+                        continue
+
+            elif self.state == ParserState.STRING:
+                self.current_statement.append(c)
+                if c == "'" and next_char == "'":
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                elif c == "'":
+                    self.state = ParserState.NORMAL
+
+            elif self.state == ParserState.COMMENT:
+                self.current_statement.append(c)
+                if c == "\n":
+                    self.state = ParserState.NORMAL
+
+            elif self.state == ParserState.MULTILINE_COMMENT:
+                self.current_statement.append(c)
+                if c == "*" and next_char == "/":
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                    self.state = ParserState.NORMAL
+
+            self.i += 1
+
+        # Handle the last statement
+        yield from self._yield_if_complete()
+
+    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):
+            self.current_case_statements += 1
+
+        is_control_keyword, keyword, keyword_len = self._look_ahead_for_keywords(
+            keywords=CONTROL_FLOW_KEYWORDS
+        )
+        if (
+            is_control_keyword
+            and keyword == END_KEYWORD
+            and self.current_case_statements > 0
+        ):
+            # If we're closing a CASE statement with END, we can just decrement the counter and continue.
+            self.current_case_statements -= 1
+        elif is_control_keyword:
+            # Yield current statement if any
+            yield from self._yield_if_complete()
+            # Yield keyword as its own statement
+            yield keyword
+            self.i += keyword_len
+            self.does_select_mean_new_statement = True
+            raise _AlreadyIncremented()
+
+        (
+            is_strict_new_statement_keyword,
+            keyword,
+            keyword_len,
+        ) = self._look_ahead_for_keywords(keywords=STRICT_NEW_STATEMENT_KEYWORDS)
+        if is_strict_new_statement_keyword:
+            yield from self._yield_if_complete()
+            self.current_statement.append(keyword)
+            self.i += keyword_len
+            self.does_select_mean_new_statement = True
+            raise _AlreadyIncremented()
+
+        (
+            is_force_new_statement_keyword,
+            keyword,
+            keyword_len,
+        ) = self._look_ahead_for_keywords(
+            keywords=(
+                NEW_STATEMENT_KEYWORDS
+                + ([SELECT_KEYWORD] if self.does_select_mean_new_statement else [])
+            ),
+        )
+        if (
+            is_force_new_statement_keyword and most_recent_real_char != ")"
+        ):  # usually we'd have a close paren that closes a CTE
+            # Force termination of current statement
+            yield from self._yield_if_complete()
+
+            self.current_statement.append(keyword)
+            self.i += keyword_len
+            raise _AlreadyIncremented()
+
+        if c == ";":
+            yield from self._yield_if_complete()
+        else:
+            self.current_statement.append(c)
+
+
+def split_statements(sql: str) -> Iterator[str]:
+    """
+    Split T-SQL code into individual statements, handling various SQL constructs.
+    """
+
+    splitter = _StatementSplitter(sql)
+    yield from splitter.process()