Skip to content

Commit 1e338a1

Browse files
ESnarkESnark
committed
fix: enhance databricks permission check logic
1 parent 8c1300e commit 1e338a1

File tree

3 files changed

+146
-21
lines changed

3 files changed

+146
-21
lines changed

‎.changeset/brave-aliens-fold.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"ansible-database-mcp": patch
3+
---
4+
5+
enhance databricks permission check logic

‎src/services/databricks-permission-checker.ts‎

Lines changed: 137 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -29,27 +29,100 @@ export async function isDatabricksReadOnlySession(adapter: DatabricksAdapter): P
2929
const catalogResult = await adapter.raw('SELECT current_catalog() as catalog');
3030
const currentCatalog = catalogResult.rows[0]?.catalog;
3131

32-
// Check if user can create tables (write permission indicator)
32+
// Define write permissions that we want to check for
33+
// Note: EXECUTE is excluded as it only allows running functions/procedures, not direct data modification
34+
const writePermissions = ['CREATE', 'MODIFY', 'DELETE', 'DROP', 'ALTER', 'UPDATE', 'INSERT', 'WRITE_FILES', 'REFRESH'];
35+
36+
// Collect all grants to check
37+
const allGrants: any[] = [];
38+
3339
try {
34-
// For Unity Catalog, we need to specify the catalog name
40+
// Check catalog-level permissions
3541
if (currentCatalog) {
36-
await adapter.raw(`SHOW GRANTS ON CATALOG ${currentCatalog}`);
37-
} else {
38-
// If no catalog is set, try to check schema permissions
39-
const schemaResult = await adapter.raw('SELECT current_schema() as schema');
40-
const currentSchema = schemaResult.rows[0]?.schema;
41-
if (currentSchema) {
42-
await adapter.raw(`SHOW GRANTS ON SCHEMA ${currentSchema}`);
42+
const catalogGrants = await adapter.raw(`SHOW GRANTS ON CATALOG ${currentCatalog}`);
43+
allGrants.push(...(catalogGrants.rows || []));
44+
}
45+
46+
// Check schema-level permissions
47+
const schemaResult = await adapter.raw('SELECT current_schema() as schema');
48+
const currentSchema = schemaResult.rows[0]?.schema;
49+
if (currentSchema && currentCatalog) {
50+
try {
51+
const schemaGrants = await adapter.raw(`SHOW GRANTS ON SCHEMA ${currentCatalog}.${currentSchema}`);
52+
allGrants.push(...(schemaGrants.rows || []));
53+
} catch (error) {
54+
// Schema grants might fail if no specific schema permissions
55+
console.debug('Could not fetch schema grants:', error);
4356
}
4457
}
4558

46-
// If we can show grants, parse them to check for write permissions
47-
// In Databricks, write permissions include: CREATE, MODIFY, DELETE, etc.
48-
// For now, we'll do a simple check
49-
return false; // Assume not read-only if we can see grants
50-
} catch (error) {
51-
// If we can't even show grants, assume read-only
59+
// Note: Databricks doesn't support "SHOW GRANTS TO user" syntax
60+
// We can only show grants ON specific objects, not grants TO a user
61+
// The catalog and schema level grants above should capture user permissions
62+
63+
// Parse grants to check for write permissions
64+
let hasExecutePermission = false;
65+
for (const grant of allGrants) {
66+
// The grant row might have different column names depending on Databricks version
67+
// Common formats: action_type, ActionType, privilege, Privilege
68+
const action = grant.action_type || grant.ActionType || grant.privilege || grant.Privilege || '';
69+
const upperAction = action.toUpperCase();
70+
71+
// Check for EXECUTE permission
72+
if (upperAction.includes('EXECUTE')) {
73+
hasExecutePermission = true;
74+
}
75+
76+
// Check if any write permission exists
77+
if (writePermissions.some(perm => upperAction.includes(perm))) {
78+
console.log(`Found write permission: ${action} for user ${currentUser}`);
79+
return false; // Has write permissions, not read-only
80+
}
81+
}
82+
83+
// Warn if EXECUTE permission is found
84+
if (hasExecutePermission) {
85+
console.warn(`⚠️ Warning: User ${currentUser} has EXECUTE permission. While this connection is allowed, ` +
86+
'be aware that EXECUTE permission could potentially run functions that modify data. ' +
87+
'Ensure that only read-only functions are available in this environment.');
88+
}
89+
90+
// If we only found SELECT or no permissions, it's read-only
5291
return true;
92+
93+
} catch (error) {
94+
// If we can't fetch grants, we need to test actual operations
95+
console.debug('Could not fetch grants, testing with actual operations:', error);
96+
97+
// Try to perform a write operation that should fail if read-only
98+
try {
99+
// Try to create a temporary view (least invasive write operation)
100+
const testViewName = `_mcp_permission_test_${Date.now()}`;
101+
await adapter.raw(`CREATE OR REPLACE TEMPORARY VIEW ${testViewName} AS SELECT 1 as test`);
102+
103+
// If we got here, we have write permissions - clean up
104+
try {
105+
await adapter.raw(`DROP VIEW IF EXISTS ${testViewName}`);
106+
} catch (cleanupError) {
107+
// Ignore cleanup errors
108+
}
109+
110+
return false; // Has write permissions
111+
} catch (writeError: any) {
112+
// Check if error is due to lack of permissions
113+
const errorMessage = writeError.message || writeError.toString();
114+
if (errorMessage.includes('PERMISSION_DENIED') ||
115+
errorMessage.includes('does not have privilege') ||
116+
errorMessage.includes('insufficient privileges') ||
117+
errorMessage.includes('CREATE') ||
118+
errorMessage.includes('denied')) {
119+
return true; // Confirmed read-only
120+
}
121+
122+
// Unknown error, assume read-only for safety
123+
console.warn('Unknown error during permission test:', writeError);
124+
return true;
125+
}
53126
}
54127
} catch (error) {
55128
console.error('Error checking Databricks permissions:', error);
@@ -72,19 +145,63 @@ export async function checkDatabricksPermissions(adapter: DatabricksAdapter): Pr
72145
const currentCatalog = catalogResult.rows[0]?.catalog || 'unknown';
73146
const currentSchema = schemaResult.rows[0]?.schema || 'unknown';
74147

75-
// Try to gather permissions
76-
const permissions: string[] = [];
148+
// Collect all unique permissions
149+
const permissionsSet = new Set<string>();
77150

78151
// Check basic SELECT permission
79152
try {
80153
await adapter.raw('SELECT 1');
81-
permissions.push('SELECT');
154+
permissionsSet.add('SELECT');
82155
} catch (error) {
83156
// No SELECT permission
84157
}
85158

86-
// For now, we'll assume read-only if we only have SELECT
87-
const isStrictlyReadOnly = permissions.length === 1 && permissions[0] === 'SELECT';
159+
// Collect grants from various levels
160+
try {
161+
// Catalog grants
162+
if (currentCatalog && currentCatalog !== 'unknown') {
163+
const catalogGrants = await adapter.raw(`SHOW GRANTS ON CATALOG ${currentCatalog}`);
164+
for (const grant of catalogGrants.rows || []) {
165+
const action = grant.action_type || grant.ActionType || grant.privilege || grant.Privilege || '';
166+
if (action) {
167+
permissionsSet.add(action.toUpperCase());
168+
}
169+
}
170+
}
171+
172+
// Schema grants
173+
if (currentSchema && currentSchema !== 'unknown' && currentCatalog && currentCatalog !== 'unknown') {
174+
try {
175+
const schemaGrants = await adapter.raw(`SHOW GRANTS ON SCHEMA ${currentCatalog}.${currentSchema}`);
176+
for (const grant of schemaGrants.rows || []) {
177+
const action = grant.action_type || grant.ActionType || grant.privilege || grant.Privilege || '';
178+
if (action) {
179+
permissionsSet.add(action.toUpperCase());
180+
}
181+
}
182+
} catch (error) {
183+
// Schema grants might not be available
184+
}
185+
}
186+
187+
// Note: Databricks doesn't support "SHOW GRANTS TO user" syntax
188+
// The catalog and schema level grants above should capture user permissions
189+
} catch (error) {
190+
console.debug('Could not fetch all grants:', error);
191+
}
192+
193+
const permissions = Array.from(permissionsSet).sort();
194+
195+
// Define write permissions
196+
// Note: EXECUTE is excluded as it only allows running functions/procedures, not direct data modification
197+
const writePermissions = ['CREATE', 'MODIFY', 'DELETE', 'DROP', 'ALTER', 'UPDATE', 'INSERT', 'WRITE_FILES', 'REFRESH'];
198+
199+
// Check if strictly read-only (only SELECT or no permissions)
200+
const hasWritePermissions = permissions.some(perm =>
201+
writePermissions.some(writePerm => perm.includes(writePerm))
202+
);
203+
204+
const isStrictlyReadOnly = !hasWritePermissions;
88205

89206
return {
90207
isStrictlyReadOnly,

‎src/tools/query.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,10 @@ import environment from "../config/environment.js";
66

77
const definition: ToolDefinition = {
88
name: 'execute-sql-query',
9-
description: 'Execute SQL query and return results. Validates write permissions on connection to ensure read-only access.',
9+
description: 'Execute SQL query and return results. Validates write permissions on connection to ensure read-only access. ' +
10+
'IMPORTANT: For Databricks connections with EXECUTE permissions - if your query calls any functions or stored procedures, ' +
11+
'please review them carefully before execution as they could potentially modify data. Ask the user to confirm before ' +
12+
'executing any CALL statements or function invocations on Databricks.',
1013
inputSchema: {
1114
query: z.string().describe('SQL query to execute'),
1215
database: z.string().describe('Database name to connect to'),

0 commit comments

Comments
 (0)