diff --git a/querybook/server/datasources/search.py b/querybook/server/datasources/search.py index 08251c86d..4cec0b8c6 100644 --- a/querybook/server/datasources/search.py +++ b/querybook/server/datasources/search.py @@ -130,16 +130,24 @@ def vector_search_tables( @register("/suggest//tables/", methods=["GET"]) -def suggest_tables(metastore_id, prefix, limit=10): +def suggest_tables(metastore_id, prefix, limit=10, active_schema='default'): api_assert(limit is None or limit <= 100, "Requesting too many tables") verify_metastore_permission(metastore_id) + possible_schema = prefix.split(".")[0] + omit_schema = not active_schema.startswith(possible_schema) + query = construct_suggest_table_query(prefix, limit, metastore_id) options = get_matching_suggestions(query, ES_CONFIG["tables"]["index_name"]) texts = [ - "{}.{}".format( - option.get("_source", {}).get("schema", ""), - option.get("_source", {}).get("name", ""), + ( + option.get("_source", {}).get("name", "") + if omit_schema and active_schema == option.get("_source", {}).get("schema", "") + else + "{}.{}".format( + option.get("_source", {}).get("schema", ""), + option.get("_source", {}).get("name", ""), + ) ) for option in options ] diff --git a/querybook/server/models/admin.py b/querybook/server/models/admin.py index 17f90a7b1..8de96b461 100644 --- a/querybook/server/models/admin.py +++ b/querybook/server/models/admin.py @@ -105,6 +105,19 @@ class QueryEngine(CRUDMixin, Base): ), ) + def get_default_schema(self): + default_schema = "default" + connection_string = self.executor_params.get("connection_string") + try: + connection_url = sql.engine.make_url(connection_string) + except Exception: + connection_url = None + + if connection_url: + default_schema = connection_url.database or default_schema + + return default_schema + def to_dict(self): # IMPORTANT: do not expose executor params unless it is for admin return { @@ -115,6 +128,7 @@ def to_dict(self): "metastore_id": self.metastore_id, "feature_params": self.get_feature_params(), "executor": self.executor, + "default_schema": self.get_default_schema(), } def to_dict_admin(self): diff --git a/querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx b/querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx index 708a80ee1..65e00880f 100644 --- a/querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx +++ b/querybook/webapp/components/QueryEditor/BoundQueryEditor.tsx @@ -29,6 +29,7 @@ export const BoundQueryEditor = React.forwardRef< | 'functionDocumentationByNameByLanguage' | 'metastoreId' | 'language' + | 'defaultSchema' > & { engine?: IQueryEngine; cellId?: number; @@ -109,6 +110,7 @@ export const BoundQueryEditor = React.forwardRef< } metastoreId={engine?.metastore_id} language={engine?.language} + defaultSchema={engine?.default_schema} /> ); diff --git a/querybook/webapp/components/QueryEditor/QueryEditor.tsx b/querybook/webapp/components/QueryEditor/QueryEditor.tsx index bc62326e5..7370b7683 100644 --- a/querybook/webapp/components/QueryEditor/QueryEditor.tsx +++ b/querybook/webapp/components/QueryEditor/QueryEditor.tsx @@ -50,6 +50,7 @@ export interface IQueryEditorProps extends IStyledQueryEditorProps { lineWrapping?: boolean; readOnly?: boolean; language?: string; + defaultSchema?: string; theme?: string; functionDocumentationByNameByLanguage?: FunctionDocumentationCollection; metastoreId?: number; @@ -107,6 +108,7 @@ export const QueryEditor: React.FC< lineWrapping = false, readOnly, language = 'hive', + defaultSchema = 'default', theme = 'default', functionDocumentationByNameByLanguage = {}, metastoreId, @@ -136,6 +138,7 @@ export const QueryEditor: React.FC< const { codeAnalysis, codeAnalysisRef } = useCodeAnalysis({ language, query: value, + defaultSchema }); const autoCompleterRef = useAutoComplete( metastoreId, diff --git a/querybook/webapp/const/queryEngine.ts b/querybook/webapp/const/queryEngine.ts index df0341c66..eda6e02ae 100644 --- a/querybook/webapp/const/queryEngine.ts +++ b/querybook/webapp/const/queryEngine.ts @@ -6,6 +6,7 @@ export interface IQueryEngine { name: string; language: string; description: string; + default_schema: string; metastore_id: number; executor: string; diff --git a/querybook/webapp/hooks/queryEditor/useCodeAnalysis.ts b/querybook/webapp/hooks/queryEditor/useCodeAnalysis.ts index f69206bf7..abac169c4 100644 --- a/querybook/webapp/hooks/queryEditor/useCodeAnalysis.ts +++ b/querybook/webapp/hooks/queryEditor/useCodeAnalysis.ts @@ -7,9 +7,10 @@ import { analyzeCode } from 'lib/web-worker'; interface IUseCodeAnalysisParams { language: string; query: string; + defaultSchema: string; } -export function useCodeAnalysis({ language, query }: IUseCodeAnalysisParams) { +export function useCodeAnalysis({ language, query, defaultSchema }: IUseCodeAnalysisParams) { /** * the ref version is used to pass into functions in codemirror * this is to prevent unnecessary codemirror refreshes @@ -19,13 +20,13 @@ export function useCodeAnalysis({ language, query }: IUseCodeAnalysisParams) { const debouncedQuery = useDebounce(query, 500); useEffect(() => { - analyzeCode(debouncedQuery, 'autocomplete', language).then( + analyzeCode(debouncedQuery, 'autocomplete', language, defaultSchema).then( (codeAnalysis) => { codeAnalysisRef.current = codeAnalysis; setCodeAnalysis(codeAnalysis); } ); - }, [debouncedQuery, language]); + }, [debouncedQuery, language, defaultSchema]); - return { codeAnalysisRef, codeAnalysis }; + return { codeAnalysis, codeAnalysisRef }; } diff --git a/querybook/webapp/lib/sql-helper/sql-autocompleter.ts b/querybook/webapp/lib/sql-helper/sql-autocompleter.ts index 2e294c208..b59a845c3 100644 --- a/querybook/webapp/lib/sql-helper/sql-autocompleter.ts +++ b/querybook/webapp/lib/sql-helper/sql-autocompleter.ts @@ -240,9 +240,12 @@ export class SqlAutoCompleter { return []; } + let omit_schema = this.codeAnalysis.lineage.activeSchema || 'default'; + const { data: names } = await SearchTableResource.suggest( metastoreId, - prefix + prefix, + omit_schema ); return names; } diff --git a/querybook/webapp/lib/sql-helper/sql-lexer.ts b/querybook/webapp/lib/sql-helper/sql-lexer.ts index d83ce80f5..865d6e3bb 100644 --- a/querybook/webapp/lib/sql-helper/sql-lexer.ts +++ b/querybook/webapp/lib/sql-helper/sql-lexer.ts @@ -65,6 +65,7 @@ const contextSensitiveKeyWord = { table: 'table', update: 'table', join: 'table', + on: 'column', set: 'column', desc: 'table', @@ -183,6 +184,7 @@ export interface ILinterWarning extends IRange { export interface ILineage { references: Record; aliases: Record>; + activeSchema: string; } export interface ICodeAnalysis { @@ -646,10 +648,13 @@ export function findWithStatementPlaceholder(statement: IToken[]) { return placeholders; } -export function findTableReferenceAndAlias(statements: IToken[][]) { - let defaultSchema = 'default'; +export function findTableReferenceAndAlias( + statements: IToken[][], + defaultSchema: string = 'default' +) { const references: Record = {}; const aliases: Record> = {}; + let activeSchema = defaultSchema; statements.forEach((statement, statementNum) => { if (statement.length === 0) { @@ -679,6 +684,7 @@ export function findTableReferenceAndAlias(statements: IToken[][]) { const secondToken = statement[tokenCounter++]; if (secondToken && secondToken.type === 'VARIABLE') { defaultSchema = secondToken.text; + activeSchema = defaultSchema; } } else { const placeholders: Set = new Set( @@ -807,6 +813,7 @@ export function findTableReferenceAndAlias(statements: IToken[][]) { return { references, aliases, + activeSchema, }; } diff --git a/querybook/webapp/lib/web-worker/index.ts b/querybook/webapp/lib/web-worker/index.ts index 1a389f4fd..f6a73d17e 100644 --- a/querybook/webapp/lib/web-worker/index.ts +++ b/querybook/webapp/lib/web-worker/index.ts @@ -11,7 +11,8 @@ let sqlEditorWorker = null; export function analyzeCode( code: string, mode = 'autocomplete', - language = 'hive' + language = 'hive', + defaultSchema = 'default' ): Promise { if (!sqlEditorWorker) { sqlEditorWorker = new Worker( @@ -40,6 +41,7 @@ export function analyzeCode( mode, id, language, + defaultSchema }); }); } diff --git a/querybook/webapp/lib/web-worker/sql-editor.worker.ts b/querybook/webapp/lib/web-worker/sql-editor.worker.ts index f06de632b..68fc4097d 100644 --- a/querybook/webapp/lib/web-worker/sql-editor.worker.ts +++ b/querybook/webapp/lib/web-worker/sql-editor.worker.ts @@ -11,12 +11,12 @@ const context: Worker = self as any; context.addEventListener( 'message', (e) => { - const { id, mode, code, language } = e.data; + const { id, mode, code, language, defaultSchema } = e.data; const tokens = tokenize(code, { language }); const statements = simpleParse(tokens); const codeAnalysis: ICodeAnalysis = { - lineage: findTableReferenceAndAlias(statements), + lineage: findTableReferenceAndAlias(statements, defaultSchema), }; if (mode === 'autocomplete') { diff --git a/querybook/webapp/resource/search.ts b/querybook/webapp/resource/search.ts index e22f8a927..52fa3380c 100644 --- a/querybook/webapp/resource/search.ts +++ b/querybook/webapp/resource/search.ts @@ -33,9 +33,9 @@ export const SearchTableResource = { count: number; }>('/search/tables/vector/', { ...params }), - suggest: (metastoreId: number, prefix: string) => + suggest: (metastoreId: number, prefix: string, active_schema: string = 'default') => ds.fetch(`/suggest/${metastoreId}/tables/`, { - prefix, + prefix, active_schema: active_schema }), };