InfoTracker 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

infotracker/parser.py

@@ -3,6 +3,7 @@ SQL parsing and lineage extraction using SQLGlot.
 """
 from __future__ import annotations
 
+import logging
 import re
 from typing import List, Optional, Set, Dict, Any
 
@@ -11,9 +12,11 @@ from sqlglot import expressions as exp
 
 from .models import (
     ColumnReference, ColumnSchema, TableSchema, ColumnLineage, 
-    TransformationType, ObjectInfo, SchemaRegistry
+    TransformationType, ObjectInfo, SchemaRegistry, ColumnNode
 )
 
+logger = logging.getLogger(__name__)
+
 
 class SqlParser:
     """Parser for SQL statements using SQLGlot."""
@@ -21,12 +24,153 @@ class SqlParser:
     def __init__(self, dialect: str = "tsql"):
         self.dialect = dialect
         self.schema_registry = SchemaRegistry()
+        self.default_database: Optional[str] = None  # Will be set from config
+    
+    def set_default_database(self, default_database: Optional[str]):
+        """Set the default database for qualification."""
+        self.default_database = default_database
+    
+    def _preprocess_sql(self, sql: str) -> str:
+        """
+        Preprocess SQL to remove control lines and join INSERT INTO #temp EXEC patterns.
+        """
+        import re
+        
+        lines = sql.split('\n')
+        processed_lines = []
+        
+        for line in lines:
+            stripped_line = line.strip()
+            
+            # Skip lines starting with DECLARE, SET, PRINT (case-insensitive)
+            if re.match(r'(?i)^(DECLARE|SET|PRINT)\b', stripped_line):
+                continue
+            
+            # Skip IF OBJECT_ID('tempdb..#...') patterns and DROP TABLE #temp patterns
+            if (re.match(r"(?i)^IF\s+OBJECT_ID\('tempdb\.\.#", stripped_line) or
+                re.match(r'(?i)^DROP\s+TABLE\s+#\w+', stripped_line)):
+                continue
+            
+            processed_lines.append(line)
+        
+        # Join the lines back together
+        processed_sql = '\n'.join(processed_lines)
+        
+        # Join two-line INSERT INTO #temp + EXEC patterns
+        processed_sql = re.sub(
+            r'(?i)(INSERT\s+INTO\s+#\w+)\s*\n\s*(EXEC\b)',
+            r'\1 \2',
+            processed_sql
+        )
+        
+        return processed_sql
+    
+    def _try_insert_exec_fallback(self, sql_content: str, object_hint: Optional[str] = None) -> Optional[ObjectInfo]:
+        """
+        Fallback parser for INSERT INTO #temp EXEC pattern when SQLGlot fails.
+        """
+        import re
+        
+        # Look for INSERT INTO #temp EXEC pattern
+        pattern = r'(?is)INSERT\s+INTO\s+(#\w+)\s+EXEC\s+([^\s(]+)'
+        match = re.search(pattern, sql_content)
+        
+        if not match:
+            return None
+        
+        temp_table = match.group(1)  # e.g., "#customer_metrics"
+        proc_name = match.group(2)   # e.g., "dbo.usp_customer_metrics_dataset"
+        
+        # Qualify procedure name if needed
+        if '.' not in proc_name and self.default_database:
+            qualified_proc_name = f"{self.default_database}.dbo.{proc_name}"
+        else:
+            qualified_proc_name = proc_name
+        
+        # Create placeholder columns for the temp table
+        placeholder_columns = [
+            ColumnSchema(
+                name="output_col_1",
+                data_type="unknown",
+                nullable=True,
+                ordinal=0
+            ),
+            ColumnSchema(
+                name="output_col_2", 
+                data_type="unknown",
+                nullable=True,
+                ordinal=1
+            )
+        ]
+        
+        # Create schema for temp table
+        schema = TableSchema(
+            namespace="tempdb",
+            name=temp_table,
+            columns=placeholder_columns
+        )
+        
+        # Create lineage for each placeholder column
+        lineage = []
+        for col in placeholder_columns:
+            lineage.append(ColumnLineage(
+                output_column=col.name,
+                input_fields=[
+                    ColumnReference(
+                        namespace="mssql://localhost/InfoTrackerDW",
+                        table_name=qualified_proc_name,
+                        column_name="*"
+                    )
+                ],
+                transformation_type=TransformationType.EXEC,
+                transformation_description=f"INSERT INTO {temp_table} EXEC {proc_name}"
+            ))
+        
+        # Set dependencies to the procedure
+        dependencies = {qualified_proc_name}
+        
+        # Register schema in registry
+        self.schema_registry.register(schema)
+        
+        # Create and return ObjectInfo
+        return ObjectInfo(
+            name=temp_table,
+            object_type="temp_table",
+            schema=schema,
+            lineage=lineage,
+            dependencies=dependencies
+        )
+    
+    def _find_last_select_string(self, sql_content: str, dialect: str = "tsql") -> str | None:
+        """Find the last SELECT statement in SQL content using SQLGlot AST."""
+        import sqlglot
+        from sqlglot import exp
+        try:
+            parsed = sqlglot.parse(sql_content, read=dialect)
+            selects = []
+            for stmt in parsed:
+                selects.extend(list(stmt.find_all(exp.Select)))
+            if not selects:
+                return None
+            return str(selects[-1])
+        except Exception:
+            return None
     
     def parse_sql_file(self, sql_content: str, object_hint: Optional[str] = None) -> ObjectInfo:
         """Parse a SQL file and extract object information."""
         try:
-            # Parse the SQL statement
-            statements = sqlglot.parse(sql_content, read=self.dialect)
+            # First check if this is a function or procedure using string matching
+            sql_upper = sql_content.upper()
+            if "CREATE FUNCTION" in sql_upper or "CREATE OR ALTER FUNCTION" in sql_upper:
+                return self._parse_function_string(sql_content, object_hint)
+            elif "CREATE PROCEDURE" in sql_upper or "CREATE OR ALTER PROCEDURE" in sql_upper:
+                return self._parse_procedure_string(sql_content, object_hint)
+            
+            # Preprocess the SQL content to handle demo script patterns
+            preprocessed_sql = self._preprocess_sql(sql_content)
+            
+            # Parse the SQL statement with SQLGlot
+            statements = sqlglot.parse(preprocessed_sql, read=self.dialect)
             if not statements:
                 raise ValueError("No valid SQL statements found")
             
@@ -37,10 +181,18 @@ class SqlParser:
                 return self._parse_create_statement(statement, object_hint)
             elif isinstance(statement, exp.Select) and self._is_select_into(statement):
                 return self._parse_select_into(statement, object_hint)
+            elif isinstance(statement, exp.Insert) and self._is_insert_exec(statement):
+                return self._parse_insert_exec(statement, object_hint)
             else:
                 raise ValueError(f"Unsupported statement type: {type(statement)}")
                 
         except Exception as e:
+            # Try fallback for INSERT INTO #temp EXEC pattern
+            fallback_result = self._try_insert_exec_fallback(sql_content, object_hint)
+            if fallback_result:
+                return fallback_result
+            
+            logger.warning("parse failed: %s", e)
             # Return an object with error information
             return ObjectInfo(
                 name=object_hint or "unknown",
@@ -58,6 +210,17 @@ class SqlParser:
         """Check if this is a SELECT INTO statement."""
         return statement.args.get('into') is not None
     
+    def _is_insert_exec(self, statement: exp.Insert) -> bool:
+        """Check if this is an INSERT INTO ... EXEC statement."""
+        # Check if the expression is a command (EXEC)
+        expression = statement.expression
+        return (
+            hasattr(expression, 'expressions') and 
+            expression.expressions and 
+            isinstance(expression.expressions[0], exp.Command) and
+            str(expression.expressions[0]).upper().startswith('EXEC')
+        )
+    
     def _parse_select_into(self, statement: exp.Select, object_hint: Optional[str] = None) -> ObjectInfo:
         """Parse SELECT INTO statement."""
         # Get target table name from INTO clause
@@ -95,12 +258,97 @@ class SqlParser:
             dependencies=dependencies
         )
     
+    def _parse_insert_exec(self, statement: exp.Insert, object_hint: Optional[str] = None) -> ObjectInfo:
+        """Parse INSERT INTO ... EXEC statement."""
+        # Get target table name from INSERT INTO clause
+        table_name = self._get_table_name(statement.this, object_hint)
+        namespace = "mssql://localhost/InfoTrackerDW"
+        
+        # Normalize temp table names
+        if table_name.startswith('#'):
+            namespace = "tempdb"
+        
+        # Extract the EXEC command
+        expression = statement.expression
+        if hasattr(expression, 'expressions') and expression.expressions:
+            exec_command = expression.expressions[0]
+            
+            # Extract procedure name and dependencies
+            dependencies = set()
+            procedure_name = None
+            
+            # Parse the EXEC command text
+            exec_text = str(exec_command)
+            if exec_text.upper().startswith('EXEC'):
+                # Extract procedure name (first identifier after EXEC)
+                parts = exec_text.split()
+                if len(parts) > 1:
+                    procedure_name = parts[1].strip('()').split('(')[0]
+                    dependencies.add(procedure_name)
+            
+            # For EXEC temp tables, we create placeholder columns since we can't determine 
+            # the actual structure without executing the procedure
+            # Create at least 2 output columns as per the requirement
+            output_columns = [
+                ColumnSchema(
+                    name="output_col_1",
+                    data_type="unknown",
+                    ordinal=0,
+                    nullable=True
+                ),
+                ColumnSchema(
+                    name="output_col_2",
+                    data_type="unknown",
+                    ordinal=1,
+                    nullable=True
+                )
+            ]
+            
+            # Create placeholder lineage pointing to the procedure
+            lineage = []
+            if procedure_name:
+                for i, col in enumerate(output_columns):
+                    lineage.append(ColumnLineage(
+                        output_column=col.name,
+                        input_fields=[ColumnReference(
+                            namespace="mssql://localhost/InfoTrackerDW",
+                            table_name=procedure_name,
+                            column_name="*"  # Wildcard since we don't know the procedure output
+                        )],
+                        transformation_type=TransformationType.EXEC,
+                        transformation_description=f"INSERT INTO {table_name} EXEC {procedure_name}"
+                    ))
+            
+            schema = TableSchema(
+                namespace=namespace,
+                name=table_name,
+                columns=output_columns
+            )
+            
+            # Register schema for future reference
+            self.schema_registry.register(schema)
+            
+            return ObjectInfo(
+                name=table_name,
+                object_type="temp_table" if table_name.startswith('#') else "table",
+                schema=schema,
+                lineage=lineage,
+                dependencies=dependencies
+            )
+        
+        # Fallback if we can't parse the EXEC command
+        raise ValueError("Could not parse INSERT INTO ... EXEC statement")
+    
     def _parse_create_statement(self, statement: exp.Create, object_hint: Optional[str] = None) -> ObjectInfo:
-        """Parse CREATE TABLE or CREATE VIEW statement."""
+        """Parse CREATE TABLE, CREATE VIEW, CREATE FUNCTION, or CREATE PROCEDURE statement."""
         if statement.kind == "TABLE":
             return self._parse_create_table(statement, object_hint)
         elif statement.kind == "VIEW":
             return self._parse_create_view(statement, object_hint)
+        elif statement.kind == "FUNCTION":
+            return self._parse_create_function(statement, object_hint)
+        elif statement.kind == "PROCEDURE":
+            return self._parse_create_procedure(statement, object_hint)
         else:
             raise ValueError(f"Unsupported CREATE statement: {statement.kind}")
     
@@ -189,31 +437,117 @@ class SqlParser:
             dependencies=dependencies
         )
     
+    def _parse_create_function(self, statement: exp.Create, object_hint: Optional[str] = None) -> ObjectInfo:
+        """Parse CREATE FUNCTION statement (table-valued functions only)."""
+        function_name = self._get_table_name(statement.this, object_hint)
+        namespace = "mssql://localhost/InfoTrackerDW"
+        
+        # Check if this is a table-valued function
+        if not self._is_table_valued_function(statement):
+            # For scalar functions, create a simple object without lineage
+            return ObjectInfo(
+                name=function_name,
+                object_type="function",
+                schema=TableSchema(
+                    namespace=namespace,
+                    name=function_name,
+                    columns=[]
+                ),
+                lineage=[],
+                dependencies=set()
+            )
+        
+        # Handle table-valued functions
+        lineage, output_columns, dependencies = self._extract_tvf_lineage(statement, function_name)
+        
+        schema = TableSchema(
+            namespace=namespace,
+            name=function_name,
+            columns=output_columns
+        )
+        
+        # Register schema for future reference
+        self.schema_registry.register(schema)
+        
+        return ObjectInfo(
+            name=function_name,
+            object_type="function",
+            schema=schema,
+            lineage=lineage,
+            dependencies=dependencies
+        )
+    
+    def _parse_create_procedure(self, statement: exp.Create, object_hint: Optional[str] = None) -> ObjectInfo:
+        """Parse CREATE PROCEDURE statement."""
+        procedure_name = self._get_table_name(statement.this, object_hint)
+        namespace = "mssql://localhost/InfoTrackerDW"
+        
+        # Extract the procedure body and find the last SELECT statement
+        lineage, output_columns, dependencies = self._extract_procedure_lineage(statement, procedure_name)
+        
+        schema = TableSchema(
+            namespace=namespace,
+            name=procedure_name,
+            columns=output_columns
+        )
+        
+        # Register schema for future reference
+        self.schema_registry.register(schema)
+        
+        return ObjectInfo(
+            name=procedure_name,
+            object_type="procedure",
+            schema=schema,
+            lineage=lineage,
+            dependencies=dependencies
+        )
+    
     def _get_table_name(self, table_expr: exp.Expression, hint: Optional[str] = None) -> str:
-        """Extract table name from expression."""
+        """Extract table name from expression and qualify with default database if needed."""
+        from .openlineage_utils import qualify_identifier
+        
         if isinstance(table_expr, exp.Table):
-            # Handle qualified names like dbo.table_name
-            if table_expr.db:
-                return f"{table_expr.db}.{table_expr.name}"
-            return str(table_expr.name)
+            # Handle three-part names: database.schema.table
+            if table_expr.catalog and table_expr.db:
+                return f"{table_expr.catalog}.{table_expr.db}.{table_expr.name}"
+            # Handle two-part names like dbo.table_name (legacy format)
+            elif table_expr.db:
+                table_name = f"{table_expr.db}.{table_expr.name}"
+                return qualify_identifier(table_name, self.default_database)
+            else:
+                table_name = str(table_expr.name)
+                return qualify_identifier(table_name, self.default_database)
         elif isinstance(table_expr, exp.Identifier):
-            return str(table_expr.this)
+            table_name = str(table_expr.this)
+            return qualify_identifier(table_name, self.default_database)
         return hint or "unknown"
     
     def _extract_column_type(self, column_def: exp.ColumnDef) -> str:
         """Extract column type from column definition."""
         if column_def.kind:
             data_type = str(column_def.kind)
-            # Convert to match expected format (lowercase for simple types)
-            if data_type.upper().startswith('VARCHAR'):
-                data_type = data_type.replace('VARCHAR', 'nvarchar')
-            elif data_type.upper() == 'INT':
-                data_type = 'int'
-            elif data_type.upper() == 'DATE':
-                data_type = 'date'
-            elif 'DECIMAL' in data_type.upper():
+            
+            # Type normalization mappings - adjust these as needed for your environment
+            # Note: This aggressive normalization can be modified by updating the mappings below
+            TYPE_MAPPINGS = {
+                'VARCHAR': 'nvarchar',  # SQL Server: VARCHAR -> NVARCHAR
+                'INT': 'int',
+                'DATE': 'date',
+            }
+            
+            data_type_upper = data_type.upper()
+            for old_type, new_type in TYPE_MAPPINGS.items():
+                if data_type_upper.startswith(old_type):
+                    data_type = data_type.replace(old_type, new_type)
+                    break
+                elif data_type_upper == old_type:
+                    data_type = new_type
+                    break
+            
+            if 'DECIMAL' in data_type_upper:
                 # Normalize decimal formatting: "DECIMAL(10, 2)" -> "decimal(10,2)"
                 data_type = data_type.replace(' ', '').lower()
+            
             return data_type.lower()
         return "unknown"
     
@@ -670,76 +1004,97 @@ class SqlParser:
         lineage = []
         output_columns = []
         
-        # Get source tables and their aliases
-        source_tables = []
-        table_aliases = {}
+        # Process all SELECT expressions, including both stars and explicit columns
+        ordinal = 0
         
-        # Check for explicit aliased star (o.*, c.*)
         for select_expr in select_stmt.expressions:
-            if isinstance(select_expr, exp.Star) and select_expr.table:
-                # This is an aliased star like o.* or c.*
-                alias = str(select_expr.table)
-                table_name = self._resolve_table_from_alias(alias, select_stmt)
-                if table_name != "unknown":
-                    columns = self._infer_table_columns(table_name)
-                    ordinal = len(output_columns)
+            if isinstance(select_expr, exp.Star):
+                if hasattr(select_expr, 'table') and select_expr.table:
+                    # This is an aliased star like o.* or c.*
+                    alias = str(select_expr.table)
+                    table_name = self._resolve_table_from_alias(alias, select_stmt)
+                    if table_name != "unknown":
+                        columns = self._infer_table_columns(table_name)
+                        
+                        for column_name in columns:
+                            output_columns.append(ColumnSchema(
+                                name=column_name,
+                                data_type="unknown",
+                                nullable=True,
+                                ordinal=ordinal
+                            ))
+                            ordinal += 1
+                            
+                            lineage.append(ColumnLineage(
+                                output_column=column_name,
+                                input_fields=[ColumnReference(
+                                    namespace="mssql://localhost/InfoTrackerDW",
+                                    table_name=table_name,
+                                    column_name=column_name
+                                )],
+                                transformation_type=TransformationType.IDENTITY,
+                                transformation_description=f"SELECT {alias}.{column_name}"
+                            ))
+                else:
+                    # Handle unqualified * - expand all tables
+                    source_tables = []
+                    for table in select_stmt.find_all(exp.Table):
+                        table_name = self._get_table_name(table)
+                        if table_name != "unknown":
+                            source_tables.append(table_name)
                     
-                    for column_name in columns:
-                        output_columns.append(ColumnSchema(
-                            name=column_name,
-                            data_type="unknown",
-                            nullable=True,
-                            ordinal=ordinal
-                        ))
-                        ordinal += 1
+                    for table_name in source_tables:
+                        columns = self._infer_table_columns(table_name)
                         
-                        lineage.append(ColumnLineage(
-                            output_column=column_name,
-                            input_fields=[ColumnReference(
-                                namespace="mssql://localhost/InfoTrackerDW",
-                                table_name=table_name,
-                                column_name=column_name
-                            )],
-                            transformation_type=TransformationType.IDENTITY,
-                            transformation_description=f"SELECT {alias}.{column_name}"
-                        ))
-                return lineage, output_columns
-        
-        # Handle unqualified * - expand all tables
-        for table in select_stmt.find_all(exp.Table):
-            table_name = self._get_table_name(table)
-            if table_name != "unknown":
-                source_tables.append(table_name)
-        
-        if not source_tables:
-            return lineage, output_columns
-        
-        # For unqualified *, expand columns from all tables
-        ordinal = 0
-        for table_name in source_tables:
-            columns = self._infer_table_columns(table_name)
-            
-            for column_name in columns:
+                        for column_name in columns:
+                            output_columns.append(ColumnSchema(
+                                name=column_name,
+                                data_type="unknown",
+                                nullable=True,
+                                ordinal=ordinal
+                            ))
+                            ordinal += 1
+                            
+                            lineage.append(ColumnLineage(
+                                output_column=column_name,
+                                input_fields=[ColumnReference(
+                                    namespace="mssql://localhost/InfoTrackerDW",
+                                    table_name=table_name,
+                                    column_name=column_name
+                                )],
+                                transformation_type=TransformationType.IDENTITY,
+                                transformation_description=f"SELECT * (from {table_name})"
+                            ))
+            else:
+                # Handle explicit column expressions (like "1 as extra_col")
+                col_name = self._extract_column_alias(select_expr) or f"col_{ordinal}"
                 output_columns.append(ColumnSchema(
-                    name=column_name,
+                    name=col_name,
                     data_type="unknown",
                     nullable=True,
                     ordinal=ordinal
                 ))
                 ordinal += 1
                 
-                lineage.append(ColumnLineage(
-                    output_column=column_name,
-                    input_fields=[ColumnReference(
+                # Try to extract lineage for this column
+                input_refs = self._extract_column_references(select_expr, select_stmt)
+                if not input_refs:
+                    # If no specific references found, treat as expression
+                    input_refs = [ColumnReference(
                         namespace="mssql://localhost/InfoTrackerDW",
-                        table_name=table_name,
-                        column_name=column_name
-                    )],
-                    transformation_type=TransformationType.IDENTITY,
-                    transformation_description=f"SELECT * (from {table_name})"
+                        table_name="LITERAL",
+                        column_name=str(select_expr)
+                    )]
+                
+                lineage.append(ColumnLineage(
+                    output_column=col_name,
+                    input_fields=input_refs,
+                    transformation_type=TransformationType.EXPRESSION,
+                    transformation_description=f"SELECT {str(select_expr)}"
                 ))
         
         return lineage, output_columns
+
     
     def _handle_union_lineage(self, stmt: exp.Expression, view_name: str) -> tuple[List[ColumnLineage], List[ColumnSchema]]:
         """Handle UNION operations."""
@@ -790,8 +1145,21 @@ class SqlParser:
         return lineage, output_columns
     
     def _infer_table_columns(self, table_name: str) -> List[str]:
-        """Infer table columns based on known schemas or naming patterns."""
-        # This is a simplified approach - you'd typically query the database
+        """Infer table columns from schema registry or fallback to patterns."""
+        # First try to get from schema registry
+        # Try different namespace combinations
+        namespaces_to_try = [
+            "mssql://localhost/InfoTrackerDW",
+            "dbo", 
+            "",
+        ]
+        
+        for namespace in namespaces_to_try:
+            schema = self.schema_registry.get(namespace, table_name)
+            if schema:
+                return [col.name for col in schema.columns]
+        
+        # Fallback to patterns if not in registry
         table_simple = table_name.split('.')[-1].lower()
         
         if 'orders' in table_simple:
@@ -805,3 +1173,407 @@ class SqlParser:
         else:
             # Generic fallback
             return ['Column1', 'Column2', 'Column3']
+
+    def _parse_function_string(self, sql_content: str, object_hint: Optional[str] = None) -> ObjectInfo:
+        """Parse CREATE FUNCTION using string-based approach."""
+        function_name = self._extract_function_name(sql_content) or object_hint or "unknown_function"
+        namespace = "mssql://localhost/InfoTrackerDW"
+        
+        # Check if this is a table-valued function
+        if not self._is_table_valued_function_string(sql_content):
+            # For scalar functions, create a simple object without lineage
+            return ObjectInfo(
+                name=function_name,
+                object_type="function",
+                schema=TableSchema(
+                    namespace=namespace,
+                    name=function_name,
+                    columns=[]
+                ),
+                lineage=[],
+                dependencies=set()
+            )
+        
+        # Handle table-valued functions
+        lineage, output_columns, dependencies = self._extract_tvf_lineage_string(sql_content, function_name)
+        
+        schema = TableSchema(
+            namespace=namespace,
+            name=function_name,
+            columns=output_columns
+        )
+        
+        # Register schema for future reference
+        self.schema_registry.register(schema)
+        
+        return ObjectInfo(
+            name=function_name,
+            object_type="function",
+            schema=schema,
+            lineage=lineage,
+            dependencies=dependencies
+        )
+    
+    def _parse_procedure_string(self, sql_content: str, object_hint: Optional[str] = None) -> ObjectInfo:
+        """Parse CREATE PROCEDURE using string-based approach."""
+        procedure_name = self._extract_procedure_name(sql_content) or object_hint or "unknown_procedure"
+        namespace = "mssql://localhost/InfoTrackerDW"
+        
+        # Extract the procedure body and find the last SELECT statement
+        lineage, output_columns, dependencies = self._extract_procedure_lineage_string(sql_content, procedure_name)
+        
+        schema = TableSchema(
+            namespace=namespace,
+            name=procedure_name,
+            columns=output_columns
+        )
+        
+        # Register schema for future reference
+        self.schema_registry.register(schema)
+        
+        return ObjectInfo(
+            name=procedure_name,
+            object_type="procedure",
+            schema=schema,
+            lineage=lineage,
+            dependencies=dependencies
+        )
+
+    def _extract_function_name(self, sql_content: str) -> Optional[str]:
+        """Extract function name from CREATE FUNCTION statement."""
+        match = re.search(r'CREATE\s+(?:OR\s+ALTER\s+)?FUNCTION\s+([^\s\(]+)', sql_content, re.IGNORECASE)
+        return match.group(1).strip() if match else None
+    
+    def _extract_procedure_name(self, sql_content: str) -> Optional[str]:
+        """Extract procedure name from CREATE PROCEDURE statement."""
+        match = re.search(r'CREATE\s+(?:OR\s+ALTER\s+)?PROCEDURE\s+([^\s\(]+)', sql_content, re.IGNORECASE)
+        return match.group(1).strip() if match else None
+
+    def _is_table_valued_function_string(self, sql_content: str) -> bool:
+        """Check if this is a table-valued function (returns TABLE)."""
+        sql_upper = sql_content.upper()
+        return "RETURNS TABLE" in sql_upper or "RETURNS @" in sql_upper
+    
+    def _extract_tvf_lineage_string(self, sql_content: str, function_name: str) -> tuple[List[ColumnLineage], List[ColumnSchema], Set[str]]:
+        """Extract lineage from a table-valued function using string parsing."""
+        lineage = []
+        output_columns = []
+        dependencies = set()
+        
+        sql_upper = sql_content.upper()
+        
+        # Handle inline TVF (RETURN AS SELECT or RETURN (SELECT))
+        if "RETURN" in sql_upper and ("AS" in sql_upper or "(" in sql_upper):
+            select_sql = self._extract_select_from_return_string(sql_content)
+            if select_sql:
+                try:
+                    parsed = sqlglot.parse(select_sql, read=self.dialect)
+                    if parsed and isinstance(parsed[0], exp.Select):
+                        lineage, output_columns = self._extract_column_lineage(parsed[0], function_name)
+                        dependencies = self._extract_dependencies(parsed[0])
+                except Exception:
+                    # Fallback to basic analysis
+                    output_columns = self._extract_basic_select_columns(select_sql)
+                    dependencies = self._extract_basic_dependencies(select_sql)
+        
+        # Handle multi-statement TVF (RETURNS @table TABLE)
+        elif "RETURNS @" in sql_upper:
+            output_columns = self._extract_table_variable_schema_string(sql_content)
+            dependencies = self._extract_basic_dependencies(sql_content)
+        
+        return lineage, output_columns, dependencies
+    
+    def _extract_procedure_lineage_string(self, sql_content: str, procedure_name: str) -> tuple[List[ColumnLineage], List[ColumnSchema], Set[str]]:
+        """Extract lineage from a procedure using string parsing."""
+        lineage = []
+        output_columns = []
+        dependencies = set()
+        
+        # Find the last SELECT statement in the procedure body
+        last_select_sql = self._find_last_select_string(sql_content)
+        if last_select_sql:
+            try:
+                parsed = sqlglot.parse(last_select_sql, read=self.dialect)
+                if parsed and isinstance(parsed[0], exp.Select):
+                    lineage, output_columns = self._extract_column_lineage(parsed[0], procedure_name)
+                    dependencies = self._extract_dependencies(parsed[0])
+            except Exception:
+                # Fallback to basic analysis
+                output_columns = self._extract_basic_select_columns(last_select_sql)
+                dependencies = self._extract_basic_dependencies(last_select_sql)
+        
+        return lineage, output_columns, dependencies
+    
+    def _extract_select_from_return_string(self, sql_content: str) -> Optional[str]:
+        """Extract SELECT statement from RETURN clause using regex."""
+        # Handle RETURN (SELECT ...)
+        match = re.search(r'RETURN\s*\(\s*(SELECT.*?)\s*\)(?:\s*;)?$', sql_content, re.IGNORECASE | re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        
+        # Handle RETURN AS (SELECT ...)
+        match = re.search(r'RETURN\s+AS\s*\(\s*(SELECT.*?)\s*\)', sql_content, re.IGNORECASE | re.DOTALL)
+        if match:
+            return match.group(1).strip()
+        
+        return None
+    
+    def _extract_table_variable_schema_string(self, sql_content: str) -> List[ColumnSchema]:
+        """Extract column schema from @table TABLE definition using regex."""
+        output_columns = []
+        
+        # Look for @Variable TABLE (column definitions)
+        match = re.search(r'@\w+\s+TABLE\s*\((.*?)\)', sql_content, re.IGNORECASE | re.DOTALL)
+        if match:
+            columns_def = match.group(1)
+            # Simple parsing of column definitions
+            for i, col_def in enumerate(columns_def.split(',')):
+                col_def = col_def.strip()
+                if col_def:
+                    parts = col_def.split()
+                    if len(parts) >= 2:
+                        col_name = parts[0].strip()
+                        col_type = parts[1].strip()
+                        output_columns.append(ColumnSchema(
+                            name=col_name,
+                            data_type=col_type,
+                            nullable=True,
+                            ordinal=i
+                        ))
+        
+        return output_columns
+        
+
+    
+    def _extract_basic_select_columns(self, select_sql: str) -> List[ColumnSchema]:
+        """Basic extraction of column names from SELECT statement."""
+        output_columns = []
+        
+        # Extract the SELECT list (between SELECT and FROM)
+        match = re.search(r'SELECT\s+(.*?)\s+FROM', select_sql, re.IGNORECASE | re.DOTALL)
+        if match:
+            select_list = match.group(1)
+            columns = [col.strip() for col in select_list.split(',')]
+            
+            for i, col in enumerate(columns):
+                # Handle aliases (column AS alias or column alias)
+                if ' AS ' in col.upper():
+                    col_name = col.split(' AS ')[-1].strip()
+                elif ' ' in col and not any(func in col.upper() for func in ['SUM', 'COUNT', 'MAX', 'MIN', 'AVG', 'CAST', 'CASE']):
+                    parts = col.strip().split()
+                    col_name = parts[-1]  # Last part is usually the alias
+                else:
+                    # Extract the base column name
+                    col_name = col.split('.')[-1] if '.' in col else col
+                    col_name = re.sub(r'[^\w]', '', col_name)  # Remove non-alphanumeric
+                
+                if col_name:
+                    output_columns.append(ColumnSchema(
+                        name=col_name,
+                        data_type="varchar",  # Default type
+                        nullable=True,
+                        ordinal=i
+                    ))
+        
+        return output_columns
+    
+    def _extract_basic_dependencies(self, sql_content: str) -> Set[str]:
+        """Basic extraction of table dependencies from SQL."""
+        dependencies = set()
+        
+        # Find FROM and JOIN clauses
+        from_matches = re.findall(r'FROM\s+([^\s\(]+)', sql_content, re.IGNORECASE)
+        join_matches = re.findall(r'JOIN\s+([^\s\(]+)', sql_content, re.IGNORECASE)
+        
+        for match in from_matches + join_matches:
+            table_name = match.strip()
+            # Clean up table name (remove aliases, schema qualifiers for dependency tracking)
+            if ' ' in table_name:
+                table_name = table_name.split()[0]
+            dependencies.add(table_name.lower())
+        
+        return dependencies
+
+    def _is_table_valued_function(self, statement: exp.Create) -> bool:
+        """Check if this is a table-valued function (returns TABLE)."""
+        # Simple heuristic: check if the function has RETURNS TABLE
+        sql_text = str(statement).upper()
+        return "RETURNS TABLE" in sql_text or "RETURNS @" in sql_text
+    
+    def _extract_tvf_lineage(self, statement: exp.Create, function_name: str) -> tuple[List[ColumnLineage], List[ColumnSchema], Set[str]]:
+        """Extract lineage from a table-valued function."""
+        lineage = []
+        output_columns = []
+        dependencies = set()
+        
+        sql_text = str(statement)
+        
+        # Handle inline TVF (RETURN AS SELECT)
+        if "RETURN AS" in sql_text.upper() or "RETURN(" in sql_text.upper():
+            # Find the SELECT statement in the RETURN clause
+            select_stmt = self._extract_select_from_return(statement)
+            if select_stmt:
+                lineage, output_columns = self._extract_column_lineage(select_stmt, function_name)
+                dependencies = self._extract_dependencies(select_stmt)
+        
+        # Handle multi-statement TVF (RETURN @table TABLE)
+        elif "RETURNS @" in sql_text.upper():
+            # Extract the table variable definition and find INSERT statements
+            output_columns = self._extract_table_variable_schema(statement)
+            lineage, dependencies = self._extract_mstvf_lineage(statement, function_name, output_columns)
+        
+        return lineage, output_columns, dependencies
+    
+    def _extract_procedure_lineage(self, statement: exp.Create, procedure_name: str) -> tuple[List[ColumnLineage], List[ColumnSchema], Set[str]]:
+        """Extract lineage from a procedure that returns a dataset."""
+        lineage = []
+        output_columns = []
+        dependencies = set()
+        
+        # Find the last SELECT statement in the procedure body
+        last_select = self._find_last_select_in_procedure(statement)
+        if last_select:
+            lineage, output_columns = self._extract_column_lineage(last_select, procedure_name)
+            dependencies = self._extract_dependencies(last_select)
+        
+        return lineage, output_columns, dependencies
+    
+    def _extract_select_from_return(self, statement: exp.Create) -> Optional[exp.Select]:
+        """Extract SELECT statement from RETURN AS clause."""
+        # This is a simplified implementation - in practice would need more robust parsing
+        try:
+            sql_text = str(statement)
+            return_as_match = re.search(r'RETURN\s*\(\s*(SELECT.*?)\s*\)', sql_text, re.IGNORECASE | re.DOTALL)
+            if return_as_match:
+                select_sql = return_as_match.group(1)
+                parsed = sqlglot.parse(select_sql, read=self.dialect)
+                if parsed and isinstance(parsed[0], exp.Select):
+                    return parsed[0]
+        except Exception:
+            pass
+        return None
+    
+    def _extract_table_variable_schema(self, statement: exp.Create) -> List[ColumnSchema]:
+        """Extract column schema from @table TABLE definition."""
+        # Simplified implementation - would need more robust parsing for production
+        output_columns = []
+        sql_text = str(statement)
+        
+        # Look for @Result TABLE (col1 type1, col2 type2, ...)
+        table_def_match = re.search(r'@\w+\s+TABLE\s*\((.*?)\)', sql_text, re.IGNORECASE | re.DOTALL)
+        if table_def_match:
+            columns_def = table_def_match.group(1)
+            # Parse column definitions
+            for i, col_def in enumerate(columns_def.split(',')):
+                col_parts = col_def.strip().split()
+                if len(col_parts) >= 2:
+                    col_name = col_parts[0].strip()
+                    col_type = col_parts[1].strip()
+                    output_columns.append(ColumnSchema(
+                        name=col_name,
+                        data_type=col_type,
+                        nullable=True,
+                        ordinal=i
+                    ))
+        
+        return output_columns
+    
+    def _extract_mstvf_lineage(self, statement: exp.Create, function_name: str, output_columns: List[ColumnSchema]) -> tuple[List[ColumnLineage], Set[str]]:
+        """Extract lineage from multi-statement table-valued function."""
+        lineage = []
+        dependencies = set()
+        
+        # Find INSERT statements into the @table variable
+        sql_text = str(statement)
+        insert_matches = re.finditer(r'INSERT\s+INTO\s+@\w+.*?SELECT(.*?)(?:FROM|$)', sql_text, re.IGNORECASE | re.DOTALL)
+        
+        for match in insert_matches:
+            try:
+                select_part = "SELECT" + match.group(1)
+                parsed = sqlglot.parse(select_part, read=self.dialect)
+                if parsed and isinstance(parsed[0], exp.Select):
+                    select_stmt = parsed[0]
+                    stmt_lineage, _ = self._extract_column_lineage(select_stmt, function_name)
+                    lineage.extend(stmt_lineage)
+                    dependencies.update(self._extract_dependencies(select_stmt))
+            except Exception:
+                continue
+        
+        return lineage, dependencies
+    
+    def _find_last_select_in_procedure(self, statement: exp.Create) -> Optional[exp.Select]:
+        """Find the last SELECT statement in a procedure body."""
+        sql_text = str(statement)
+        
+        # Find all SELECT statements that are not part of INSERT/UPDATE/DELETE
+        select_matches = list(re.finditer(r'(?<!INSERT\s)(?<!UPDATE\s)(?<!DELETE\s)SELECT\s+.*?(?=(?:FROM|$))', sql_text, re.IGNORECASE | re.DOTALL))
+        
+        if select_matches:
+            # Get the last SELECT statement
+            last_match = select_matches[-1]
+            try:
+                select_sql = last_match.group(0)
+                # Find the FROM clause and complete SELECT
+                from_match = re.search(r'FROM.*?(?=(?:WHERE|GROUP|ORDER|HAVING|;|$))', sql_text[last_match.end():], re.IGNORECASE | re.DOTALL)
+                if from_match:
+                    select_sql += from_match.group(0)
+                
+                parsed = sqlglot.parse(select_sql, read=self.dialect)
+                if parsed and isinstance(parsed[0], exp.Select):
+                    return parsed[0]
+            except Exception:
+                pass
+        
+        return None
+    
+    def _extract_column_alias(self, select_expr: exp.Expression) -> Optional[str]:
+        """Extract column alias from a SELECT expression."""
+        if hasattr(select_expr, 'alias') and select_expr.alias:
+            return str(select_expr.alias)
+        elif isinstance(select_expr, exp.Alias):
+            return str(select_expr.alias)
+        elif isinstance(select_expr, exp.Column):
+            return str(select_expr.this)
+        else:
+            # Try to extract from the expression itself
+            expr_str = str(select_expr)
+            if ' AS ' in expr_str.upper():
+                parts = expr_str.split()
+                as_idx = -1
+                for i, part in enumerate(parts):
+                    if part.upper() == 'AS':
+                        as_idx = i
+                        break
+                if as_idx >= 0 and as_idx + 1 < len(parts):
+                    return parts[as_idx + 1].strip("'\"")
+        return None
+    
+    def _extract_column_references(self, select_expr: exp.Expression, select_stmt: exp.Select) -> List[ColumnReference]:
+        """Extract column references from a SELECT expression."""
+        refs = []
+        
+        # Find all column references in the expression
+        for column_expr in select_expr.find_all(exp.Column):
+            table_name = "unknown"
+            column_name = str(column_expr.this)
+            
+            # Try to resolve table name from table reference or alias
+            if hasattr(column_expr, 'table') and column_expr.table:
+                table_alias = str(column_expr.table)
+                table_name = self._resolve_table_from_alias(table_alias, select_stmt)
+            else:
+                # If no table specified, try to infer from FROM clause
+                tables = []
+                for table in select_stmt.find_all(exp.Table):
+                    tables.append(self._get_table_name(table))
+                if len(tables) == 1:
+                    table_name = tables[0]
+            
+            if table_name != "unknown":
+                refs.append(ColumnReference(
+                    namespace="mssql://localhost/InfoTrackerDW",
+                    table_name=table_name,
+                    column_name=column_name
+                ))
+        
+        return refs