From a7f9392836e95cd8fa6a8b781eac204b5307d297 Mon Sep 17 00:00:00 2001
From: Olivier Tassinari <olivier.tassinari@gmail.com>
Date: Sat, 6 Feb 2021 18:24:47 +0100
Subject: [PATCH] [DataGrid] Resolve the api ref at the same time as any other
 ref

---
 .../_modules_/grid/hooks/features/useApiRef.ts  | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/packages/grid/_modules_/grid/hooks/features/useApiRef.ts b/packages/grid/_modules_/grid/hooks/features/useApiRef.ts
index 5765d3e5e063d..25584cf1b8805 100644
--- a/packages/grid/_modules_/grid/hooks/features/useApiRef.ts
+++ b/packages/grid/_modules_/grid/hooks/features/useApiRef.ts
@@ -10,14 +10,17 @@ function createGridApi(): GridApi {
   return new EventEmitter() as GridApi;
 }
 
-export function useApiRef(apiRefProp?: ApiRef): ApiRef {
-  const apiRef = React.useRef<GridApi>(createGridApi());
+// Public developers facing overload
+export function useApiRef(): ApiRef;
 
-  React.useEffect(() => {
-    if (apiRefProp?.current) {
-      apiRefProp.current = apiRef.current;
-    }
-  });
+// Internal grid facing overload
+export function useApiRef(apiRefProp: ApiRef | undefined): ApiRef;
+
+export function useApiRef(...args): any {
+  const apiRefProp = args[0];
+  const apiRef = React.useRef<GridApi>(args.length === 0 ? null : createGridApi());
+
+  React.useImperativeHandle(apiRefProp, () => apiRef.current, [apiRef]);
 
   return apiRef;
 }