datahub-agent-context 1.3.1.10rc1__py3-none-any.whl → 1.4.0rc2__py3-none-any.whl

datahub_agent_context/snowflake/generators/network_rules.py

@@ -0,0 +1,53 @@
+"""Generate Snowflake network rules SQL."""
+
+
+def generate_network_rules_sql(datahub_domain: str) -> str:
+    """Generate network rules SQL that uses configuration variables.
+
+    Args:
+        datahub_domain: The domain extracted from the DataHub URL (e.g., 'fieldeng.acryl.io')
+    """
+    return f"""-- ============================================================================
+-- Step 1: Network Rules and External Access Integration for DataHub API
+-- ============================================================================
+-- This script creates the network rule and external access integration
+-- required for UDFs to make outbound calls to DataHub
+--
+-- Prerequisites:
+-- - Run 00_configuration.sql first to set variables
+-- - You must have ACCOUNTADMIN role or equivalent privileges
+-- ============================================================================
+
+-- Use the configured database and schema
+USE DATABASE IDENTIFIER($SF_DATABASE);
+USE SCHEMA IDENTIFIER($SF_SCHEMA);
+USE WAREHOUSE IDENTIFIER($SF_WAREHOUSE);
+
+-- Step 1: Create network rule for DataHub API
+-- For EGRESS mode with external access, use HOST_PORT type with domain name
+-- Note: Snowflake resolves the domain to IPs automatically
+CREATE OR REPLACE NETWORK RULE datahub_api_rule
+  MODE = EGRESS
+  TYPE = HOST_PORT
+  VALUE_LIST = ('{datahub_domain}')
+  COMMENT = 'Network rule for {datahub_domain} DataHub API access';
+
+-- Step 2: Create external access integration
+CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION datahub_access
+  ALLOWED_NETWORK_RULES = (datahub_api_rule)
+  ALLOWED_AUTHENTICATION_SECRETS = (datahub_url, datahub_token)
+  ENABLED = TRUE
+  COMMENT = 'External access integration for DataHub API calls';
+
+-- Step 3: Grant usage on integration to the specified role
+GRANT USAGE ON INTEGRATION datahub_access TO ROLE IDENTIFIER($SF_ROLE);
+
+-- Verify the setup
+SHOW NETWORK RULES LIKE 'datahub_api_rule';
+SHOW INTEGRATIONS LIKE 'datahub_access';
+
+SELECT
+    'Network rule and external access integration created successfully!' AS status,
+    '{datahub_domain}' AS datahub_domain,
+    $SF_DATABASE || '.' || $SF_SCHEMA AS location;
+"""

datahub_agent_context/snowflake/generators/stored_procedure.py

@@ -0,0 +1,87 @@
+"""Generate Snowflake stored procedure SQL."""
+
+
+def generate_stored_procedure_sql() -> str:
+    """Generate stored procedure SQL that uses configuration variables."""
+    return """-- ============================================================================
+-- Step 3: Create JavaScript Stored Procedure for SQL Execution
+-- ============================================================================
+-- This creates a stored procedure that can execute SELECT queries dynamically
+-- and return results as JSON
+--
+-- Prerequisites:
+-- - Run 00_configuration.sql first to set variables
+-- ============================================================================
+
+USE DATABASE IDENTIFIER($SF_DATABASE);
+USE SCHEMA IDENTIFIER($SF_SCHEMA);
+USE WAREHOUSE IDENTIFIER($SF_WAREHOUSE);
+
+-- Create the JavaScript stored procedure (using CREATE OR REPLACE, so no need to drop first)
+CREATE OR REPLACE PROCEDURE EXECUTE_DYNAMIC_SQL(SQL_TEXT STRING)
+RETURNS VARIANT
+LANGUAGE JAVASCRIPT
+EXECUTE AS CALLER
+AS
+$$
+  // Validate that it's a SELECT query only
+  var queryUpper = SQL_TEXT.trim().toUpperCase();
+  if (!queryUpper.startsWith('SELECT')) {
+    return {
+      "success": false,
+      "error": "Only SELECT queries are allowed. Query must start with SELECT.",
+      "status": 400
+    };
+  }
+
+  try {
+    var stmt = snowflake.createStatement({ sqlText: SQL_TEXT });
+    var rs = stmt.execute();
+
+    var results = [];
+    var columns = [];
+
+    // Get column names from the first row metadata
+    var columnCount = rs.getColumnCount();
+    for (var i = 1; i <= columnCount; i++) {
+      columns.push(rs.getColumnName(i));
+    }
+
+    // Process all rows
+    while (rs.next()) {
+      var row = {};
+      for (var i = 1; i <= columnCount; i++) {
+        var colName = rs.getColumnName(i);
+        var colValue = rs.getColumnValue(i);
+        // Handle null values and convert to JSON-serializable format
+        row[colName] = colValue;
+      }
+      results.push(row);
+    }
+
+    return {
+      "success": true,
+      "columns": columns,
+      "rows": results,
+      "row_count": results.length
+    };
+
+  } catch (err) {
+    return {
+      "success": false,
+      "error": err.message,
+      "status": 500
+    };
+  }
+$$;
+
+-- Grant usage on the procedure
+GRANT USAGE ON PROCEDURE EXECUTE_DYNAMIC_SQL(STRING) TO ROLE IDENTIFIER($SF_ROLE);
+
+-- Verify it was created
+DESCRIBE PROCEDURE EXECUTE_DYNAMIC_SQL(STRING);
+
+SELECT
+    'Stored procedure EXECUTE_DYNAMIC_SQL created successfully!' AS status,
+    $SF_DATABASE || '.' || $SF_SCHEMA || '.EXECUTE_DYNAMIC_SQL' AS procedure_name;
+"""