Repository: reata/sqllineage

Issue:
missing database/schema in lineage result

Comments:

Patch:
diff --git a/sqllineage/core.py b/sqllineage/core.py
index 2b1b77b..3ebe4dc 100644
--- a/sqllineage/core.py
+++ b/sqllineage/core.py
@@ -120,7 +120,7 @@ Target Tables:
                         # referring https://github.com/andialbrecht/sqlparse/issues/218 for further information
                         pass
                     else:
-                        self._source_tables.add(Table(sub_token.get_real_name()))
+                        self._source_tables.add(Table.create(sub_token))
                     source_table_token_flag = False
             elif target_table_token_flag:
                 if self.__token_negligible_before_tablename(sub_token):
@@ -131,7 +131,7 @@ Target Tables:
                     if not isinstance(sub_token.token_first(skip_cm=True), Identifier):
                         raise SQLLineageException("An Identifier is expected")
                     self._target_tables.add(
-                        Table(sub_token.token_first(skip_cm=True).get_real_name())
+                        Table.create(sub_token.token_first(skip_cm=True))
                     )
                 elif isinstance(sub_token, Comparison):
                     # create table tab1 like tab2, tab1 like tab2 will be parsed as Comparison
@@ -141,12 +141,12 @@ Target Tables:
                         and isinstance(sub_token.right, Identifier)
                     ):
                         raise SQLLineageException("An Identifier is expected")
-                    self._target_tables.add(Table(sub_token.left.get_real_name()))
-                    self._source_tables.add(Table(sub_token.right.get_real_name()))
+                    self._target_tables.add(Table.create(sub_token.left))
+                    self._source_tables.add(Table.create(sub_token.right))
                 else:
                     if not isinstance(sub_token, Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._target_tables.add(Table(sub_token.get_real_name()))
+                    self._target_tables.add(Table.create(sub_token))
                 target_table_token_flag = False
             elif temp_table_token_flag:
                 if self.__token_negligible_before_tablename(sub_token):
@@ -154,23 +154,19 @@ Target Tables:
                 else:
                     if not isinstance(sub_token, Identifier):
                         raise SQLLineageException("An Identifier is expected")
-                    self._source_tables.add(Table(sub_token.get_real_name()))
-                    self._target_tables.add(Table(sub_token.get_real_name()))
+                    self._source_tables.add(Table.create(sub_token))
+                    self._target_tables.add(Table.create(sub_token))
                     self._extract_from_DML(sub_token)
                     temp_table_token_flag = False
 
     def _extract_from_DDL_DROP(self, stmt: Statement) -> None:
         for st_tables in (self._source_tables, self._target_tables):
             st_tables -= {
-                Table(t.get_real_name())
-                for t in stmt.tokens
-                if isinstance(t, Identifier)
+                Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)
             }
 
     def _extract_from_DDL_ALTER(self, stmt: Statement) -> None:
-        tables = [
-            Table(t.get_real_name()) for t in stmt.tokens if isinstance(t, Identifier)
-        ]
+        tables = [Table.create(t) for t in stmt.tokens if isinstance(t, Identifier)]
         keywords = [t for t in stmt.tokens if t.ttype is Keyword]
         if any(k.normalized == "RENAME" for k in keywords) and len(tables) == 2:
             for st_tables in (self._source_tables, self._target_tables):
diff --git a/sqllineage/models.py b/sqllineage/models.py
index 6ec1349..25bbdef 100644
--- a/sqllineage/models.py
+++ b/sqllineage/models.py
@@ -1,8 +1,14 @@
-from typing import Optional
+import warnings
 
+from sqlparse.sql import Identifier
 
-class Database:
-    def __init__(self, name: Optional[str] = "<unknown>"):
+from sqllineage.exceptions import SQLLineageException
+
+
+class Schema:
+    unknown = "<unknown>"
+
+    def __init__(self, name: str = unknown):
         self.raw_name = name
 
     def __str__(self):
@@ -17,14 +23,26 @@ class Database:
     def __hash__(self):
         return hash(str(self))
 
+    def __bool__(self):
+        return str(self) != self.unknown
+
 
 class Table:
-    def __init__(self, name: str, database: Optional[Database] = Database()):
-        self.database = database
-        self.raw_name = name
+    def __init__(self, name: str, schema: Schema = Schema()):
+        if len(name.split(".")) == 2:
+            schema_name, table_name = name.split(".")
+            self.schema = Schema(schema_name)
+            self.raw_name = table_name
+            if schema:
+                warnings.warn("Name is in schema.table format, schema param is ignored")
+        elif "." not in name:
+            self.schema = schema
+            self.raw_name = name
+        else:
+            raise SQLLineageException("Invalid format for table name: %s", name)
 
     def __str__(self):
-        return "{}.{}".format(self.database, self.raw_name.lower())
+        return "{}.{}".format(self.schema, self.raw_name.lower())
 
     def __repr__(self):
         return "Table: " + str(self)
@@ -35,6 +53,15 @@ class Table:
     def __hash__(self):
         return hash(str(self))
 
+    @staticmethod
+    def create(identifier: Identifier):
+        schema = (
+            Schema(identifier.get_parent_name())
+            if identifier.get_parent_name() is not None
+            else Schema()
+        )
+        return Table(identifier.get_real_name(), schema)
+
 
 class Partition:
     pass


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
index 147bfe6..47f09bc 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1,13 +1,13 @@
-from sqllineage.models import Database, Table
+from sqllineage.models import Schema, Table
 
 
 def test_repr_dummy():
-    assert repr(Database())
+    assert repr(Schema())
     assert repr(Table(""))
 
 
 def test_hash_eq():
-    assert Database("a") == Database("a")
-    assert len({Database("a"), Database("a")}) == 1
+    assert Schema("a") == Schema("a")
+    assert len({Schema("a"), Schema("a")}) == 1
     assert Table("a") == Table("a")
     assert len({Table("a"), Table("a")}) == 1
diff --git a/tests/test_select.py b/tests/test_select.py
index 4f53d07..30d8ee3 100644
--- a/tests/test_select.py
+++ b/tests/test_select.py
@@ -5,6 +5,10 @@ def test_select():
     helper("SELECT col1 FROM tab1", {"tab1"})
 
 
+def test_select_with_schema():
+    helper("SELECT col1 FROM schema1.tab1", {"schema1.tab1"})
+
+
 def test_select_multi_line():
     helper(
         """SELECT col1 FROM