Method QueryConn

backend/plugin/db/cockroachdb/cockroachdb.go:611–700

QueryConn queries a SQL statement in a given connection.

(ctx context.Context, conn *sql.Conn, statement string, queryContext db.QueryContext)

// QueryConn queries a SQL statement in a given connection.
func (*Driver) QueryConn(ctx context.Context, conn *sql.Conn, statement string, queryContext db.QueryContext) ([]*v1pb.QueryResult, error) {
	singleSQLs, err := crdbparser.SplitSQLStatement(statement)
	if err != nil {
		return nil, err
	}
	if len(singleSQLs) == 0 {
		return nil, nil
	}

	var results []*v1pb.QueryResult
	for _, singleSQL := range singleSQLs {
		statement := singleSQL
		if queryContext.Explain {
			statement = fmt.Sprintf("EXPLAIN %s", statement)
		} else if queryContext.Limit > 0 {
			statement = getStatementWithResultLimit(statement, queryContext.Limit)
		}

		_, allQuery, err := base.ValidateSQLForEditor(storepb.Engine_POSTGRES, statement)
		if err != nil {
			return nil, err
		}

		// Sanitize the schema name by escaping any quotes.
		safeSchemeName := strings.ReplaceAll(queryContext.Schema, "\"", "\"\"")

		// If the queryContext.Schema is not empty, set the search path for the database connection to the specified schema.
		if queryContext.Schema != "" {
			if err := crdb.Execute(func() error {
				_, err := conn.ExecContext(ctx, fmt.Sprintf(`SET search_path TO "%s";`, safeSchemeName)) // NOSONAR(go:S2077) safeSchemeName is sanitized by escaping double quotes above
				return err
			}); err != nil {
				return nil, err
			}
		}

		startTime := time.Now()
		queryResult, err := func() (*v1pb.QueryResult, error) {
			if allQuery {
				var r *v1pb.QueryResult
				if err := crdb.Execute(func() error {
					rows, err := conn.QueryContext(ctx, statement)
					if err != nil {
						return err
					}
					defer rows.Close()
					r, err = util.RowsToQueryResult(rows, makeValueByTypeName, convertValue, queryContext.MaximumSQLResultSize)
					if err != nil {
						return err
					}
					err = rows.Err()
					return err
				}); err != nil {
					return nil, err
				}
				return r, nil
			}