Use AI Assistant

Create data aggregators and extractors

If you need a custom data aggregator or extractor, you can ask AI Assistant to create one for your task, and then add the created script to the aggregators or extractors directory.

Providing the LLM with an example of a working script might shorten the number of prompts and corrections. In DataGrip, you can find bundled scripts in the Files tool window (Alt+2) under Scratches and Consoles | Extensions | Database Tools and SQL | data.

For more information about aggregators and extractors, refer to Aggregate view and Data extractors.

In this tutorial, we will create an aggregator and use the following AVG.groovy aggregator as an example:

/* * Available context bindings: * COLUMNS List<DataColumn> * ROWS Iterable<DataRow> * OUT { append() } * FORMATTER { format(row, col); formatValue(Object, col); getTypeName(Object, col); isStringLiteral(Object, col); } * TRANSPOSED Boolean * plus ALL_COLUMNS, TABLE, DIALECT * * where: * DataRow { rowNumber(); first(); last(); data(): List<Object>; value(column): Object } * DataColumn { columnNumber(), name() } */ import static java.math.MathContext.DECIMAL128 BigDecimal RES = 0 int i = 0 ROWS.each { row -> COLUMNS.each { column -> def value = row.value(column) if (value instanceof Number) { RES = RES.add(value, DECIMAL128) i++ } else if (value.toString().isBigDecimal()) { RES = RES.add(value.toString().toBigDecimal(), DECIMAL128) i++ } } } if (i > 0) { RES = RES.divide(i, DECIMAL128) OUT.append(RES.toString()) } else { OUT.append("Not enough values") }

In DataGrip, you can find this script in the Files tool window under Scratches and Consoles | Extensions | Database Tools and SQL | data | aggregators.

Step 1. Request a script

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the input field, type your prompt. Describe the required script and provide the LLM with an example of a working script.
In the upper-right corner of the field with the generated script, click Copy to Clipboard to copy the code to clipboard.

The generated script:

/* * Available context bindings: * COLUMNS List<DataColumn> * ROWS Iterable<DataRow> * OUT { append() } * FORMATTER { format(row, col); formatValue(Object, col); getTypeName(Object, col); isStringLiteral(Object, col); } * TRANSPOSED Boolean * plus ALL_COLUMNS, TABLE, DIALECT * * where: * DataRow { rowNumber(); first(); last(); data(): List<Object>; value(column): Object } * DataColumn { columnNumber(), name() } */ import java.util.* import java.util.stream.Collectors // Create a map to store string values and their frequencies Map<String, Integer> freqs = new HashMap<>() ROWS.each { row -> COLUMNS.each { column -> def value = row.value(column) if (value instanceof String) { freqs.put(value, freqs.containsKey(value) ? freqs.get(value) + 1 : 1) } } } // Find max frequency int maxFreq = Collections.max(freqs.values()) // Collect most frequent strings List<String> mostFrequent = freqs.entrySet().stream() .filter(e -> e.getValue() == maxFreq) .map(Map.Entry::getKey) .collect(Collectors.toList()) if (mostFrequent.size() > 0) { OUT.append(String.join(", ", mostFrequent)) } else { OUT.append("No string values found") }

Data aggregator and extractor script creation

Step 2. Save the script as a file

You can open the Files tool window by doing one of the following:
- In the main menu, go to View | Tool Windows | Files.
- On the tool window bar, click Files.
- Press Alt+2.
Navigate to Scratches and Consoles | Extensions | Database Tools and SQL | data | aggregators.
Right-click the aggregators node and select New | File.
In the New File popup, type the new file name. For example,
most_frequent_values.groovy.
Paste your new data aggregator script into the file.
(Optional) If the Convert Pasted Code dialog appears, press Cancel. In this tutorial, we save the generated script itself.

New data aggregator appears on the aggregator list in status bar for data editor and can be used with your data.

Request queries and information using a natural language

AI Assistant can give you insights about your schema, and it can also generate SQL queries for the schema based on a natural language request. Having access to your schema, it will analyze the structure and provide you with the requested query or information.

For illustration purposes, we will use the PostgreSQL Sakila schema imported to the ai_demo database. You can get the dump files by cloning the dumps repository.

git clone https://github.com/DataGrip/dumps.git

For more information about cloning repositories, refer to Set up a Git repository. For more information about running dump files in DataGrip, refer to Import SQL dump files.

Step 1. Attach your schema

To grant the LLM access to your database schema, attach the schema to chat in the AI Assistant tool window.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the AI Assistant tool window, click Attach Schema above input field and select the schema that you want to attach. In our case, it is ai_demo.public.
If the Attach Schema dialog appears, click Attach to attach the schema.

For more information about attaching schemas to the chat, refer to Attach schema using schema selector.

Attach the ai_demo.public schema to AI Assistant chat

Step 2. Ask in the chat

Ask AI Assistant to generate a certain query or to provide you with certain information about your schema. Consider the following examples:
1. Request to generate a query.
2. Request for the insights about your schema.
3. Request to perform a search.

Optimize your schema and queries using EXPLAIN PLAN

AI Assistant can help you with optimizing your schema and queries using the EXPLAIN PLAN command. Having access to your schema, it can analyze the plan and suggest optimizations.

For illustration purposes, we will use the PostgreSQL Sakila schema imported to the ai_demo database. You can get the dump files by cloning the dumps repository.

git clone https://github.com/DataGrip/dumps.git

For more information about cloning repositories, refer to Set up a Git repository. For more information about running dump files in DataGrip, refer to Import SQL dump files.

In this tutorial, we will use the following query:

            SELECT
                customer.first_name,
                customer.last_name,
                category.name
            FROM
                customer
                    INNER JOIN
                rental ON customer.customer_id = rental.customer_id
                    INNER JOIN
                inventory ON rental.inventory_id = inventory.inventory_id
                    INNER JOIN
                film_category ON inventory.film_id = film_category.film_id
                    INNER JOIN
                category ON film_category.category_id = category.category_id
            WHERE
                category.name = 'Comedy';
        

Step 1. Run the EXPLAIN PLAN command

Obtain and copy the EXPLAIN PLAN that you want AI Assistant to analyze.

In the query console, right-click your query and navigate to Explain Plan | Explain Plan (Raw).
In the result tab, press Ctrl+A and Ctrl+C to copy the result to the clipboard.

For more information about EXPLAIN PLAN commands, refer to Query execution plan.

EXPLAIN PLAN for the original query:

Hash Join (cost=141.42..159.52 rows=86 width=81) Hash Cond: (customer.customer_id = rental.customer_id) -> Seq Scan on customer (cost=0.00..14.99 rows=599 width=17) -> Hash (cost=140.35..140.35 rows=86 width=70) -> Nested Loop (cost=38.41..140.35 rows=86 width=70) -> Hash Join (cost=38.12..126.34 rows=24 width=72) Hash Cond: (inventory.film_id = film_category.film_id) -> Seq Scan on inventory (cost=0.00..70.81 rows=4581 width=6) -> Hash (cost=38.06..38.06 rows=5 width=70) -> Hash Join (cost=19.43..38.06 rows=5 width=70) Hash Cond: (film_category.category_id = category.category_id) -> Seq Scan on film_category (cost=0.00..16.00 rows=1000 width=4) -> Hash (cost=19.38..19.38 rows=4 width=72) -> Seq Scan on category (cost=0.00..19.38 rows=4 width=72) Filter: ((name)::text = 'Comedy'::text) -> Index Scan using idx_fk_inventory_id on rental (cost=0.29..0.54 rows=4 width=6) Index Cond: (inventory_id = inventory.inventory_id)

Step 2. Attach your schema

To grant the LLM access to your database schema, attach the schema to chat in the AI Assistant tool window.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the AI Assistant tool window, click Attach Schema above input field and select the schema that you want to attach. In our case, it is ai_demo.public.
If the Attach Schema dialog appears, click Attach to attach the schema.

For more information about attaching schemas to the chat, refer to Attach schema using schema selector.

Step 3. Request optimization suggestions and commands

Ask AI Assistant to analyze the EXPLAIN PLAN of your query and to suggest the possible optimizations.

In the input field, type your prompt, press Ctrl+V to paste the EXPLAIN PLAN, and press Enter.
Type another prompt requesting the necessary commands for the suggested optimizations and press Enter.

Step 4. Run the suggested CREATE INDEX commands

Select the commands and run them by clicking Execute on the toolbar, or by pressing Ctrl+Enter.

The suggested commands:

CREATE INDEX idx_customer_customer_id ON customer(customer_id); CREATE INDEX idx_rental_customer_id ON rental(customer_id); CREATE INDEX idx_rental_inventory_id ON rental(inventory_id); CREATE INDEX idx_inventory_inventory_id ON inventory(inventory_id); CREATE INDEX idx_inventory_film_id ON inventory(film_id); CREATE INDEX idx_film_category_film_id ON film_category(film_id); CREATE INDEX idx_film_category_category_id ON film_category(category_id); CREATE INDEX idx_category_category_id ON category(category_id); CREATE INDEX idx_category_name ON category(name);

IDE will create the defined database object, you can view the output in the Output tab of Services tool window.

Execution of the commands for the suggested optimizations

Step 5. Run the EXPLAIN PLAN for optimized schema and query

In the upper-right corner of the field with the generated CREATE INDEX commands, click Insert Snippet at Caret to insert the code to query console.
The suggested optimized query:
SELECT customer.first_name, customer.last_name, comedic_category.name FROM customer INNER JOIN rental ON customer.customer_id = rental.customer_id INNER JOIN inventory ON rental.inventory_id = inventory.inventory_id INNER JOIN film_category ON inventory.film_id = film_category.film_id INNER JOIN ( SELECT * FROM category WHERE category.name = 'Comedy' ) AS comedic_category ON film_category.category_id = comedic_category.category_id;
EXPLAIN PLAN for the optimized query:
Hash Join (cost=27.80..236.36 rows=1003 width=81) Hash Cond: (rental.customer_id = customer.customer_id) -> Nested Loop (cost=5.32..211.23 rows=1003 width=70) -> Nested Loop (cost=5.04..49.31 rows=286 width=72) -> Nested Loop (cost=4.76..13.35 rows=62 width=70) -> Seq Scan on category (cost=0.00..1.20 rows=1 width=72) Filter: ((name)::text = 'Comedy'::text) -> Bitmap Heap Scan on film_category (cost=4.76..11.53 rows=62 width=4) Recheck Cond: (category_id = category.category_id) -> Bitmap Index Scan on idx_film_category_category_id (cost=0.00..4.74 rows=62 width=0) Index Cond: (category_id = category.category_id) -> Index Scan using idx_inventory_film_id on inventory (cost=0.28..0.53 rows=5 width=6) Index Cond: (film_id = film_category.film_id) -> Index Scan using idx_rental_inventory_id on rental (cost=0.29..0.53 rows=4 width=6) Index Cond: (inventory_id = inventory.inventory_id) -> Hash (cost=14.99..14.99 rows=599 width=17) -> Seq Scan on customer (cost=0.00..14.99 rows=599 width=17)

Compare the DDL of two database objects

Once provided with access to your database schema, AI Assistant can analyze it and compare DDL of the database objects.

Let us compare the actor and actor_test tables of the testing.public schema.

The actor and actor_test tables in Database Explorer

Step 1. Attach your schema

To grant the LLM access to your database schema, attach the schema to chat in the AI Assistant tool window.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the AI Assistant tool window, click Attach Schema above input field and select the schema that you want to attach. In our case, it is ai_demo.public.
If the Attach Schema dialog appears, click Attach to attach the schema.

For more information about attaching schemas to the chat, refer to Attach schema using schema selector.

Attach the testing.public schema to AI Assistant chat

Step 2. Compare the DDL of your tables

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the input field, type your prompt and press Enter.
AI Assistant will compare the DDL and explain the differences.

Comparison of DLL of two tables by AI Assistant

Fix code mistakes

If some code is not working, you can ask AI Assistant to check if there are any mistakes in it and to correct them.

Let us use the following code with a mistake in it as an example:

            CREATE FUNCTION film_not_in_stock(p_film_id integer, p_store_id integer, OUT p_film_count integer) RETURNS SETOF integer
            AS $_$
            SELECT inventory_id
            FROM inventory
            WHERE film_id = $1
                AND store_id = $2
                AND NOT inventory_in_stock(inventory_id);
            $_$
                LANGUAGE sql;
        

The mistake here is that the declared function has an output parameter p_film_count, but the SQL body does not calculate or return it. Additionally, the function is declared as returning SETOF integer, but only returns result from a SELECT statement.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the input field, enter your prompt, paste the code, request code correction, and press Enter.
If required, make corrections by adding them to your next prompts.

The corrected code:

CREATE FUNCTION film_not_in_stock(p_film_id integer, p_store_id integer) RETURNS integer AS $$ DECLARE p_film_count integer; BEGIN SELECT COUNT(inventory_id) INTO p_film_count FROM inventory WHERE film_id = p_film_id AND store_id = p_store_id AND NOT inventory_in_stock(inventory_id); RETURN p_film_count; END; $$ LANGUAGE plpgsql;

SQL code mistakes correction request and fixed code

Migrate your code from one DBMS to another

Migrating database objects from one DBMS to another is a complex task that involves steps such as data migration, schema conversion, and so on. You can ask AI Assistant to generate migration scripts for the source code of your database objects.

Let us use the following inventory_held_by_customer procedure of a PostgreSQL Sakila schema as an example:

            CREATE FUNCTION inventory_held_by_customer(p_inventory_id integer) RETURNS integer
                LANGUAGE plpgsql
            AS
            $$
            DECLARE
                v_customer_id INTEGER;
            BEGIN

                SELECT customer_id INTO v_customer_id
                FROM rental
                WHERE return_date IS NULL
                AND inventory_id = p_inventory_id;

                RETURN v_customer_id;
            END $$;

            ALTER FUNCTION inventory_held_by_customer(integer) OWNER TO guest;
        

For AI Assistant to migrate your code, provide the migration details in your prompt.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the input field, type your prompt. Define the original DBMS and the target DBMS, paste the DDL of the database object that you want to migrate, and press Enter.
If required, make corrections by adding them to your next prompts.

The generated migration script:

CREATE FUNCTION inventory_held_by_customer(p_inventory_id integer) RETURNS INTEGER DETERMINISTIC BEGIN DECLARE v_customer_id INTEGER; SELECT customer_id INTO v_customer_id FROM rental WHERE return_date IS NULL AND inventory_id = p_inventory_id; RETURN v_customer_id; END;

Format your code

It can be difficult to read the code that has no indentation, line breaks, consistent capitalization, and so on. To make such code easier to read, use AI Assistant to format it in accordance with your preferred code style.

We will use the following poorly-formatted SQL code as an example:

            SELECT e.firstName,e.lastName,d.name,p.projectName FROM employees
            e join departments d on e.departmentId=d.id INNER JOIN employee_projects ep on e.id=ep.employeeId
            INNER JOIN projects p on ep.projectId=p.id where e.salary>50000 AND d.name='Sales' AND p.status='active'
        

For AI Assistant to format your code, describe your preferred code formatting rules in the prompt.

On the toolbar, click More tool windows and select AI Assistant to open the AI Assistant tool window.
In the input field, type your prompt with the formatting rules, paste the code that you want to format, and press Enter.
If required, make corrections by adding them to your next prompts.

The formatted code:

SELECT e.first_name, e.last_name, d.name, p.project_name FROM employees AS e JOIN departments AS d ON e.department_id = d.id INNER JOIN employee_projects AS ep ON e.id = ep.employee_id INNER JOIN projects AS p ON ep.project_id = p.id WHERE e.salary > 50000 AND d.name = 'Sales' AND p.status = 'active';

DataGrip 2024.3 Help

Install the AI Assistant plugin

Create data aggregators and extractors

Step 1. Request a script

Step 2. Save the script as a file

Request queries and information using a natural language

Step 1. Attach your schema

Step 2. Ask in the chat

Optimize your schema and queries using EXPLAIN PLAN

Step 1. Run the EXPLAIN PLAN command

Step 2. Attach your schema

Step 3. Request optimization suggestions and commands

Step 4. Run the suggested CREATE INDEX commands

Step 5. Run the EXPLAIN PLAN for optimized schema and query

Compare the DDL of two database objects

Step 1. Attach your schema

Step 2. Compare the DDL of your tables

Fix code mistakes

Migrate your code from one DBMS to another

Format your code